using Sundials, Test, ForwardDiff
using Sundials: N_Vector, N_Vector_S
using LinearAlgebra

# Create context for tests
ctx_ptr = Ref{Sundials.SUNContext}(C_NULL)
Sundials.SUNContext_Create(C_NULL, Base.unsafe_convert(Ptr{Sundials.SUNContext}, ctx_ptr))
ctx = ctx_ptr[]

function mycopy!(pp, arr::Matrix)
    nj = size(arr, 2)
    ps = unsafe_wrap(Array, pp, nj)
    for j in 1:nj
        arr[:, j] = Sundials.asarray(ps[j])
    end
    return arr
end

function mycopy!(arr::Matrix, pp)
    nj = size(arr, 2)
    ps = unsafe_wrap(Array, pp, nj)
    for j in 1:nj
        Sundials.asarray(ps[j])[:] = arr[:, j]
    end
    return
end

f!(dy, t, y, p) = (dy[:] = y .* p)

function srhs(t, y, ydot, ys, ysdot)
    n = length(y)
    np = 2
    dyt = similar(y)
    chunk = min(n, 8)
    c1 = ForwardDiff.JacobianConfig(nothing, dyt, dyt, ForwardDiff.Chunk{chunk}())
    c2 = ForwardDiff.JacobianConfig(nothing, dyt, p, ForwardDiff.Chunk{chunk}())
    jac = ForwardDiff.jacobian((dy, y) -> f!(dy, t, y, p), dyt, y, c1, Val{false}())
    #jac = ReverseDiff.jacobian!(t1, y)
    ysdot[:] = jac * ys

    jac = ForwardDiff.jacobian((dy, p) -> f!(dy, t, y, p), dyt, p, c2, Val{false}())
    #jac = ReverseDiff.jacobian!(t2, p)
    return ysdot[:, 1:np] += jac
end

"""
    sens(f!, t0, y0, p, tout; reltol, abstol)
Compute the solution and sensitivities to the parametrized ODE problem defined by `f!(ẏ, t, y, p)`, starting at t0, y0, p.
Return the solutions `y` at times tout (>t0) as well as the corresponding sensitivity matrices `ys`.
y[i, t] is the solutions component i at timestep t.
ys[i, j, t] is the i-th component sensitivity wrt the j-th parameter at timestep t.
ys[i, np+j, t] the i-th component sensitivity wrt the j-th initial condition value.
"""
function sens(
        f!::Function,
        t0::Float64,
        y0::Vector{Float64},
        p::Vector{Float64},
        tout::Vector{Float64};
        reltol::Float64 = 1.0e-5,
        abstol::Float64 = 1.0e-5
    )
    n = length(y0)
    np = length(p)
    ys0 = zeros(n, np .+ n)
    ys0[:, np .+ (1:n)] = Matrix(1.0I, n, n)

    #t1 = ReverseDiff.JacobianTape((dy,y)->f!(dy,0,y,p), dyt, y0)
    #t2 = ReverseDiff.JacobianTape((dy,p)->f!(dy,0,y0,p), dyt, p)

    pbar = abs.(vcat(p, y0))
    return y, ys = cvodes(f!, srhs, t0, y0, ys0, p, reltol, abstol, pbar, tout)
end

### internals

## data structure dealing with the sundials callbacks

mutable struct CVSData
    f::Any  # f(t,y,dy)
    fs::Any # fs()
    p::Any
    jys::Any
    jdys::Any
end

function CVSData(f, fs, p, n::Int, nS::Int)
    return CVSData(
        f, fs, p, Array{Float64}(undef, n, nS),
        Array{Float64}(undef, n, nS)
    )
end

function cvrhsfn(t::Float64, y::N_Vector, dy::N_Vector, data::CVSData)
    data.f(convert(Vector, dy), t, convert(Vector, y), data.p)
    return Sundials.CV_SUCCESS
end

function cvsensrhsfn(
        ns::Cint,
        t::Float64,
        y::N_Vector,
        dy::N_Vector,
        ys::N_Vector_S,
        dys::N_Vector_S,
        data::CVSData,
        tmp1::N_Vector,
        tmp2::N_Vector
    )
    jys = data.jys
    jdys = data.jdys
    mycopy!(ys, data.jys)
    data.fs(t, convert(Vector, y), convert(Vector, dy), jys, jdys)
    mycopy!(jdys, dys)
    return Sundials.CV_SUCCESS
end

## cvodes wrapper

"Given the sensitivity problem, return (y,ys) where
y[i,t] is the solutions i-th component for timestep t and
ys[i,j,t] is the sensitivity of the i-th component wrt to the j-th parameter, where
the last parameter indices correspond to the initial conditions components."
function cvodes(f, fS, t0, y0, yS0, p, reltol, abstol, pbar, t::AbstractVector)
    N, Ns = size(yS0)
    y = zeros(N, length(t))
    ys = zeros(N, Ns, length(t))
    tret = [t0]
    yret = similar(y0)
    ysret = similar(yS0)
    y0n = Sundials.NVector(y0, ctx)
    yretn = Sundials.NVector(yret, ctx)
    yS0n = [Sundials.NVector(yS0[:, j], ctx) for j in 1:Ns]
    yS0nv = Sundials.NVectorArray(yS0n)
    crhs = Sundials.@cfunction(
        cvrhsfn,
        Cint,
        (Sundials.realtype, N_Vector, N_Vector, Ref{CVSData})
    )
    csensrhs = Sundials.@cfunction(
        cvsensrhsfn,
        Cint,
        (
            Cint,
            Sundials.realtype,
            N_Vector,
            N_Vector,
            N_Vector_S,
            N_Vector_S,
            Ref{CVSData},
            N_Vector,
            N_Vector,
        )
    )

    ##

    mem_ptr = Sundials.CVodeCreate(Sundials.CV_ADAMS, ctx)
    #mem_ptr = Sundials.CVodeCreate(Sundials.CV_BDF, ctx)
    cvode_mem = Sundials.Handle(mem_ptr)
    Sundials.CVodeInit(cvode_mem, crhs, t0, y0n)
    nls = Sundials.SUNNonlinSol_FixedPoint(yretn, Cint(0), ctx)
    nls_handle = Sundials.NonLinSolHandle(nls, Sundials.FixedPoint())
    Sundials.CVodeSetNonlinearSolver(cvode_mem, nls)
    Sundials.CVodeSStolerances(cvode_mem, reltol, abstol)
    data = CVSData(f, fS, p, size(yS0)...)
    Sundials.CVodeSetUserData(cvode_mem, data)

    Sundials.CVodeSensInit(cvode_mem, Ns, Sundials.CV_STAGGERED, csensrhs, yS0nv)
    nlsS = Sundials.SUNNonlinSol_FixedPointSens(Cint(Ns), yretn, Cint(0), ctx)
    nlsS_handle = Sundials.NonLinSolHandle(nlsS, Sundials.FixedPoint())
    Sundials.CVodeSetNonlinearSolverSensStg(cvode_mem, nlsS)
    Sundials.CVodeSetSensParams(cvode_mem, C_NULL, pbar, C_NULL)
    Sundials.CVodeSensEEtolerances(cvode_mem)
    GC.@preserve data y0n yretn yS0n yS0nv pbar nls_handle nlsS_handle begin
        for i in 1:length(t)
            Sundials.CVode(cvode_mem, t[i], yretn, tret, Sundials.CV_NORMAL)
            Sundials.CVodeGetSens(cvode_mem, tret, yS0nv)
            mycopy!(pointer(yS0nv), ysret)
            y[:, i] = yret
            ys[:, :, i] = ysret
        end
    end
    empty!(nlsS_handle)
    empty!(nls_handle)
    empty!(cvode_mem)
    return y, ys
end

t0 = 0.0
t = [1.0, 2.0]
y0 = [1.0, 2.0]
p = [3.0, 4.0]
y, ys = sens(f!, t0, y0, p, t)
@test isapprox(y[1, 1], 20.0856; rtol = 1.0e-3)
@test isapprox(ys[2, 2, 2], 11924.3; rtol = 1.0e-3)

# Clean up context
Sundials.SUNContext_Free(ctx)
