using Sundials

## f routine. Compute function f(t,y).

function f(t, y_nv, ydot_nv)
    y = convert(Vector, y_nv)
    ydot = convert(Vector, ydot_nv)
    ydot[1] = -0.04 * y[1] + 1.0e4 * y[2] * y[3]
    ydot[3] = 3.0e7 * y[2] * y[2]
    ydot[2] = -ydot[1] - ydot[3]
    return Sundials.CV_SUCCESS
end

## g routine. Compute functions g_i(t,y) for i = 0,1.

function g(t, y_nv, gout_ptr, user_data)
    y = convert(Vector, y_nv)
    gout = Sundials.asarray(gout_ptr, (2,))
    gout[1] = y[1] - 0.0001
    gout[2] = y[3] - 0.01
    return Sundials.CV_SUCCESS
end

g_C = @cfunction(
    g, Cint,
    (Sundials.realtype, Sundials.N_Vector, Ptr{Sundials.realtype}, Ptr{Cvoid})
)

## Jacobian routine. Compute J(t,y) = df/dy.
# broken -- needs a wrapper from Sundials._DlsMat to Matrix and Jac user function wrapper
function Jac(N, t, ny, fy, Jptr, user_data, tmp1, tmp2, tmp3)
    y = convert(Vector, ny)
    dlsmat = unpack(
        IOString(
            unsafe_wrap(
                convert(Ptr{UInt8}, Jptr),
                (sum(map(sizeof, Sundials._DlsMat)) + 10,),
                false
            )
        ),
        Sundials._DlsMat
    )
    J = unsafe_wrap(unsafe_ref(dlsmat.cols), (Int(neq), Int(neq)), false)
    J[1, 1] = -0.04
    J[1, 2] = 1.0e4 * y[3]
    J[1, 3] = 1.0e4 * y[2]
    J[2, 1] = 0.04
    J[2, 2] = -1.0e4 * y[3] - 6.0e7 * y[2]
    J[2, 3] = -1.0e4 * y[2]
    J[3, 2] = 6.0e7 * y[2]
    return Sundials.CV_SUCCESS
end

neq = 3

t0 = 0.0
t1 = 0.4
tmult = 10.0
nout = 12
y0 = [1.0, 0.0, 0.0]
reltol = 1.0e-4
abstol = [1.0e-8, 1.0e-14, 1.0e-6]
userdata = nothing
ctx_ptr = Ref{Sundials.SUNContext}(C_NULL)
Sundials.SUNContext_Create(C_NULL, Base.unsafe_convert(Ptr{Sundials.SUNContext}, ctx_ptr))
ctx = ctx_ptr[]
mem_ptr = Sundials.CVodeCreate(Sundials.CV_BDF, ctx)
cvode_mem = Sundials.Handle(mem_ptr)
userfun = Sundials.UserFunctionAndData(f, userdata)
Sundials.CVodeSetUserData(cvode_mem, userfun)

function getcfunrob(userfun::T) where {T}
    return @cfunction(
        Sundials.cvodefun,
        Cint,
        (Sundials.realtype, Sundials.N_Vector, Sundials.N_Vector, Ref{T})
    )
end

# Create NVector before using it
y0_nvec = Sundials.NVector(y0, ctx)

Sundials.@checkflag Sundials.CVodeInit(
    cvode_mem, getcfunrob(userfun), t1,
    y0_nvec
)
Sundials.@checkflag Sundials.CVodeInit(cvode_mem, getcfunrob(userfun), t0, y0_nvec)
abstol_nvec = Sundials.NVector(abstol, ctx)
Sundials.@checkflag Sundials.CVodeSVtolerances(cvode_mem, reltol, abstol_nvec)
Sundials.@checkflag Sundials.CVodeRootInit(cvode_mem, 2, g_C)
A = Sundials.SUNDenseMatrix(neq, neq, ctx)
mat_handle = Sundials.MatrixHandle(A, Sundials.DenseMatrix())
LS = Sundials.SUNLinSol_Dense(y0_nvec, A, ctx)
LS_handle = Sundials.LinSolHandle(LS, Sundials.Dense())
Sundials.@checkflag Sundials.CVDlsSetLinearSolver(cvode_mem, LS, A)
#Sundials.@checkflag Sundials.CVDlsSetDenseJacFn(cvode_mem, Jac)

iout = 0
tout = t1
t = [t0]

while iout < nout
    y = similar(y0)
    y_nvec = Sundials.NVector(y, ctx)
    flag = Sundials.CVode(cvode_mem, tout, y_nvec, t, Sundials.CV_NORMAL)
    copyto!(y, y_nvec.v)
    println("T=", tout, ", Y=", y)
    if flag == Sundials.CV_ROOT_RETURN
        rootsfound = zeros(Cint, 2)
        Sundials.@checkflag Sundials.CVodeGetRootInfo(cvode_mem, rootsfound)
        println("roots=", rootsfound)
    elseif flag == Sundials.CV_SUCCESS
        global iout += 1
        global tout *= tmult
    end
end

empty!(cvode_mem)
empty!(mat_handle)
empty!(LS_handle)
Sundials.SUNContext_Free(ctx)
