using Revise
    using LinearAlgebra, Plots, SparseArrays, Setfield, Parameters

f1(u, v) = u^2 * v
norminf = x -> norm(x, Inf)

function plotsol(x; kwargs...)
    N = div(length(x), 2)
    plot!(x[1:N], label="u"; kwargs...)
    plot!(x[N+1:2N], label="v"; kwargs...)
end

function Fbru!(f, x, p)
    @unpack α, β, D1, D2, l = p
    n = div(length(x), 2)
    h = 1.0 / (n+0); h2 = h*h
    c1 = D1 / l^2 / h2
    c2 = D2 / l^2 / h2

    u = @view x[1:n]
    v = @view x[n+1:2n]

    # Dirichlet boundary conditions
    f[1]   = c1 * (α     - 2u[1] + u[2] ) + α - (β + 1) * u[1] + f1(u[1], v[1])
    f[end] = c2 * (v[n-1] - 2v[n] + β / α)             + β * u[n] - f1(u[n], v[n])

    f[n]   = c1 * (u[n-1] - 2u[n] +  α   ) + α - (β + 1) * u[n] + f1(u[n], v[n])
    f[n+1] = c2 * (β / α  - 2v[1] + v[2])          + β * u[1] - f1(u[1], v[1])

    for i=2:n-1
          f[i] = c1 * (u[i-1] - 2u[i] + u[i+1]) + α - (β + 1) * u[i] + f1(u[i], v[i])
        f[n+i] = c2 * (v[i-1] - 2v[i] + v[i+1])           + β * u[i] - f1(u[i], v[i])
    end
    return f
end

function Fbru(x, p)
    f = similar(x)
    Fbru!(f, x, p)
end

function Jbru_sp(x, p)
    @unpack α, β, D1, D2, l = p
    # compute the Jacobian using a sparse representation
    n = div(length(x), 2)
    h = 1.0 / (n+0); h2 = h*h

    c1 = D1 / p.l^2 / h2
    c2 = D2 / p.l^2 / h2

    u = @view x[1:n]
    v = @view x[n+1:2n]

    diag   = zeros(eltype(x), 2n)
    diagp1 = zeros(eltype(x), 2n-1)
    diagm1 = zeros(eltype(x), 2n-1)

    diagpn = zeros(eltype(x), n)
    diagmn = zeros(eltype(x), n)

    @. diagmn = β - 2 * u * v
    @. diagm1[1:n-1] = c1
    @. diagm1[n+1:end] = c2

    @. diag[1:n]    = -2c1 - (β + 1) + 2 * u * v
    @. diag[n+1:2n] = -2c2 - u * u

    @. diagp1[1:n-1]   = c1
    @. diagp1[n+1:end] = c2

    @. diagpn = u * u
    J = spdiagm(0 => diag, 1 => diagp1, -1 => diagm1, n => diagpn, -n => diagmn)
    return J
end

function finalise_solution(z, tau, step, contResult)
    n = div(length(z.u), 2)
    printstyled(color=:red, "--> Solution constant = ", norm(diff(z.u[1:n])), " - ", norm(diff(z.u[n+1:2n])), "\n")
    return true
end

n = 100
####################################################################################################
# test for the Jacobian expression
# using ForwardDiff
# sol0 = rand(2n)
# J0 = ForwardDiff.jacobian(x-> Fbru(x, par_bru), sol0) |> sparse
# J1 = Jbru_sp(sol0, par_bru)
# J0 - J1
####################################################################################################
# different parameters to define the Brusselator model and guess for the stationary solution
par_bru = (α = 2., β = 5.45, D1 = 0.008, D2 = 0.004, l = 0.3)
    sol0 = vcat(par_bru.α * ones(n), par_bru.β/par_bru.α * ones(n))



using DifferentialEquations, DiffEqOperators, Sundials, ForwardDiff

function plotPeriodic(x)
    n = div(size(x,1),2)
    plot(heatmap(x[1:n,:]', ylabel="Time", color=:viridis),
            heatmap(x[n+1:end,:]', color=:viridis))
end

FOde(f, x, p, t) = Fbru!(f, x, p)
Jbru(x, dx, p) = ForwardDiff.derivative(t -> Fbru(x .+ t .* dx, p), 0.)
JacOde(J, x, p, t) = copyto!(J, Jbru_sp(x, p))

u0 = sol0 .+ 0.01 .* rand(2n)
par_hopf = (@set par_bru.l = 0.51792 + 0.01)

jac_prototype = Jbru_sp(ones(sol0|> length), @set par_bru.β = 0)
    jac_prototype.nzval .= ones(length(jac_prototype.nzval))

ff = ODEFunction(FOde, jac_prototype = JacVecOperator{Float64}(FOde, u0, par_hopf))
prob = ODEProblem(ff, u0, (0., 5400.), par_hopf)
probsundials = ODEProblem(FOde, u0, (0., 5200.), par_hopf)

M = 3
    sol = @time solve(probsundials, CVODE_BDF(linear_solver=:GMRES), saveat = LinRange(820-Thopf, 820, M+1))


function sectionP(x, po::AbstractMatrix, p)
    res = 1.0
    N, M = size(po)

    for ii = 1:size(x,2)
        res *= dot(x .- po[:, ii], Fbru(po[:, ii], p))
    end
    @show p res
    res
end

M = 1
affect!(integrator) = terminate!(integrator)
pSection(u, t, integrator) = (sectionP(u, Array(sol[:,1:M]), par_hopf)) * (integrator.iter > 1)
cb = ContinuousCallback(pSection, affect!)
solve(probsundials, Rodas4(), save_everystep = false, callback = cb)