using Revise
    using DiffEqOperators
    using LinearAlgebra, Plots, SparseArrays, Parameters, Setfield
    using DifferentialEquations, DiffEqOperators, Sundials

plotsol(x; kwargs...) = plot(heatmap(reshape(real.(x[1:div(length(x),2)]), Nx, Ny)',color=:viridis; kwargs...), heatmap(reshape(real.(x[div(length(x),2)+1:end]), Nx, Ny)',color=:viridis; kwargs...))
norminf = x -> norm(x, Inf)

function Laplacian2D(Nx, Ny, lx, ly, bc = :Dirichlet)
    hx = 2lx/Nx
    hy = 2ly/Ny
    D2x = CenteredDifference(2, 2, hx, Nx)
    D2y = CenteredDifference(2, 2, hy, Ny)
    if bc == :Neumann
        Qx = Neumann0BC(hx)
        Qy = Neumann0BC(hy)
    else
        Qx = Dirichlet0BC(typeof(hx))
        Qy = Dirichlet0BC(typeof(hy))
    end
    A = kron(sparse(I, Ny, Ny), sparse(D2x * Qx)[1]) + kron(sparse(D2y * Qy)[1], sparse(I, Nx, Nx))
    return A, D2x
end

function Fcgl(u, p)
    @unpack r, μ, ν, c3, c5, Δ = p
    n = div(length(u), 2)
    u1 = @view u[1:n]
    u2 = @view u[n+1:2n]

    ua = u1.^2 .+ u2.^2

    f = similar(u)
    f1 = @view f[1:n]
    f2 = @view f[n+1:2n]

    mul!(f, Δ, u)

    @. f1 .+= r * u1 - ν * u2 - ua * (c3 * u1 - μ * u2) - c5 * ua^2 * u1
    @. f2 .+= r * u2 + ν * u1 - ua * (c3 * u2 + μ * u1) - c5 * ua^2 * u2

    return f
end

# remark: I checked this against finite differences
function Jcgl(u, p)
    @unpack r, μ, ν, c3, c5, Δ = p

    n = div(length(u), 2)
    u1 = @view u[1:n]
    u2 = @view u[n+1:2n]

    ua = u1.^2 .+ u2.^2

    f1u = zero(u1)
    f2u = zero(u1)
    f1v = zero(u1)
    f2v = zero(u1)

    @. f1u =  r - 2 * u1 * (c3 * u1 - μ * u2) - c3 * ua - 4 * c5 * ua * u1^2 - c5 * ua^2
    @. f1v = -ν - 2 * u2 * (c3 * u1 - μ * u2)  + μ * ua - 4 * c5 * ua * u1 * u2
    @. f2u =  ν - 2 * u1 * (c3 * u2 + μ * u1)  - μ * ua - 4 * c5 * ua * u1 * u2
    @. f2v =  r - 2 * u2 * (c3 * u2 + μ * u1) - c3 * ua - 4 * c5 * ua * u2 ^2 - c5 * ua^2

    jacdiag = vcat(f1u, f2v)

    Δ + spdiagm(0 => jacdiag, n => f1v, -n => f2u)
end

function FOde(f, x, p, t)
    f .= Fcgl(x, p)
end

function JacOde(J, x, p, t)
    Jcgl = Jcgl(x, p)
    @show typeof(J) typeof(Jcgl)
    copyto!(J, Jcgl)
end

Nx = 41*1
    Ny = 21*1
    lx = pi
    ly = pi/2

    Δ = Laplacian2D(Nx, Ny, lx, ly)[1]
    par_cgl = (r = 0.5, μ = 0.1, ν = 1.0, c3 = -1.0, c5 = 1.0, Δ = kron(spdiagm(0 => ones(2)), Δ))
    sol0 = zeros(2Nx, Ny)

function FOde(f, x, p, t)
    f .= Fcgl(x, p)
end

function JacOde(J, x, p, t)
    Jsp = Jcgl(x, p)
    @show typeof(J) typeof(Jsp)
    copyto!(J, Jsp)
end

u0 = rand(2Nx*Ny)
ff = ODEFunction(FOde, jac = JacOde, jac_prototype = par_cgl.Δ)#jac_prototype = JacVecOperator{Float64}(FOde, u0, p))
prob = ODEProblem(ff, u0, (0., 2.), (@set par_cgl.r = 1.55 + 0.))
sol = @time solve(prob, Tsit5())#,save_everystep=false)
sol = @time solve(prob, CVODE_BDF(linear_solver=:GMRES))#,save_everystep=false)
sol = @time solve(prob, TRBDF2(linsolve=LinSolveGMRES()))#,save_everystep=false)
sol = @time solve(prob, TRBDF2())#,save_everystep=false)

heatmap(sol[:,:])

using ForwardDiff
function test_f(u0)
    println("#"^10)
    @show typeof(u0) eltype(u0)
    _prob = remake(prob; u0=u0)#, jac_prototype = convert.(ForwardDiff.Dual{ForwardDiff.Tag{typeof(test_f),Float64},Float64,12}, par_cgl.Δ))
    @show typeof(_prob.f.jac_prototype)

    solve(_prob,TRBDF2(),save_everystep=false)[end]
end

fd_res = ForwardDiff.jacobian(test_f, u0)