text:
## Common Interface Solve Functions

function DiffEqBase.__solve(
    prob::Union{DiffEqBase.AbstractODEProblem, DiffEqBase.AbstractDAEProblem},
    alg::algType,
    timeseries = [],
    ts = [],
    ks = [],
    recompile::Type{Val{recompile_flag}} = Val{true};
    kwargs...,
) where {algType <: Union{SundialsODEAlgorithm, SundialsDAEAlgorithm}, recompile_flag}

    integrator = DiffEqBase.__init(prob, alg, timeseries, ts, ks; kwargs...)
    if integrator.sol.retcode == :Default
        solve!(integrator)
    end
    integrator.sol
end

function DiffEqBase.__solve(
    prob::Union{DiffEqBase.AbstractSteadyStateProblem{uType,isinplace},DiffEqBase.AbstractNonlinearProblem{uType,isinplace}},
    alg::algType,
    timeseries=[],
    ts=[],
    ks=[],
    recompile::Type{Val{recompile_flag}}=Val{true};
    kwargs...,
) where {algType <: SundialsNonlinearSolveAlgorithm,recompile_flag,uType,isinplace}

    if typeof(prob.u0) <: Number
        u0 = [prob.u0]
    else
        u0 = deepcopy(prob.u0)
    end

    sizeu = size(prob.u0)
    p = prob.p
    userdata = alg.userdata
    linsolve = linear_solver(alg)
    jac_upper = alg.jac_upper
    jac_lower = alg.jac_lower

    ### Fix the more general function to Sundials allowed style
    if typeof(prob.f) <: ODEFunction
        t = Inf
        if !isinplace && typeof(prob.u0) <: Number
            f! = (du, u) -> (du .= prob.f(first(u), p, t); Cint(0))
        elseif !isinplace && typeof(prob.u0) <: Vector{Float64}
            f! = (du, u) -> (du .= prob.f(u, p, t); Cint(0))
        elseif !isinplace && typeof(prob.u0) <: AbstractArray
            f! = (du, u) -> (du .= vec(prob.f(reshape(u, sizeu), p, t)); Cint(0))
        elseif typeof(prob.u0) <: Vector{Float64}
            f! = (du, u) -> prob.f(du, u, p, t)
        else # Then it's an in-place function on an abstract array
            f! = (du, u) -> (prob.f(reshape(du, sizeu), reshape(u, sizeu), p, t);
            du = vec(du);
            0)
        end
    elseif typeof(prob.f) <: NonlinearFunction
        if !isinplace && typeof(prob.u0) <: Number
            f! = (du, u) -> (du .= prob.f(first(u), p); Cint(0))
        elseif !isinplace && typeof(prob.u0) <: Vector{Float64}
            f! = (du, u) -> (du .= prob.f(u, p); Cint(0))
        elseif !isinplace && typeof(prob.u0) <: AbstractArray
            f! = (du, u) -> (du .= vec(prob.f(reshape(u, sizeu), p)); Cint(0))
        elseif typeof(prob.u0) <: Vector{Float64}
            f! = (du, u) -> prob.f(du, u, p)
        else # Then it's an in-place function on an abstract array
            f! = (du, u) -> (prob.f(reshape(du, sizeu), reshape(u, sizeu), p);
            du = vec(du);
            0)
        end
    end
    u = zero(u0)
    resid = similar(u)
    u = kinsol(f!, u0,
               userdata=userdata,
               linear_solver=linsolve,
               jac_upper=jac_upper,
               jac_lower=jac_lower)

    f!(resid, u)
    DiffEqBase.build_solution(prob, alg, u, resid;retcode=:Success)
end

function DiffEqBase.__init(
    prob::DiffEqBase.AbstractODEProblem{uType, tupType, isinplace},
    alg::SundialsODEAlgorithm{Method, LinearSolver},
    timeseries = [],
    ts = [],
    ks = [];
    verbose = true,
    callback = nothing,
    abstol = 1 / 10^6,
    reltol = 1 / 10^3,
    saveat = Float64[],
    tstops = Float64[],
    maxiters = Int(1e5),
    dt = nothing,
    dtmin = 0.0,
    dtmax = 0.0,
    timeseries_errors = true,
    dense_errors = false,
    save_everystep = isempty(saveat), save_idxs = nothing,
    save_on = true,
    save_start = save_everystep || isempty(saveat) || typeof(saveat) <: Number ? true :
                 prob.tspan[1] in saveat,
    save_end = save_everystep || isempty(saveat) || typeof(saveat) <: Number ? true :
               prob.tspan[2] in saveat,
    dense = save_everystep && isempty(saveat),
    progress = false,
    progress_name = "ODE",
    progress_message = DiffEqBase.ODE_DEFAULT_PROG_MESSAGE,
    save_timeseries = nothing,
    advance_to_tstop = false,
    stop_at_next_tstop = false,
    userdata = nothing,
    alias_u0 = false,
    kwargs...,
) where {uType, tupType, isinplace, Method, LinearSolver}

    tType = eltype(tupType)

    if verbose
        warned = !isempty(kwargs) && DiffEqBase.check_keywords(alg, kwargs, warnlist)
        warned && DiffEqBase.warn_compat()
    end

    if prob.f.mass_matrix != LinearAlgebra.I
        error("This solver is not able to use mass matrices.")
    end

    if typeof(reltol) <: AbstractArray
        error("Sundials only allows scalar reltol.")
    end

    progress && Logging.@logmsg(-1, progress_name, _id = _id = :Sundials, progress = 0)

    callbacks_internal = DiffEqBase.CallbackSet(callback)

    max_len_cb = DiffEqBase.max_vector_callback_length(callbacks_internal)
    if max_len_cb isa VectorContinuousCallback
        callback_cache = DiffEqBase.CallbackCache(max_len_cb.len, Float64, Float64)
    else
        callback_cache = nothing
    end

    tspan = prob.tspan
    t0 = tspan[1]

    tdir = sign(tspan[2] - tspan[1])

    tstops_internal, saveat_internal =
        tstop_saveat_disc_handling(tstops, saveat, tdir, tspan, tType)

    if typeof(prob.u0) <: Number
        u0 = [prob.u0]
    else
        if alias_u0
            u0 = vec(prob.u0)
        else
            u0 = vec(copy(prob.u0))
        end
    end

    sizeu = size(prob.u0)

    ### Fix the more general function to Sundials allowed style
    if !isinplace && typeof(prob.u0) <: Number
        f! = (du, u, p, t) -> (du .= prob.f(first(u), p, t); Cint(0))
    elseif !isinplace && typeof(prob.u0) <: Vector{Float64}
        f! = (du, u, p, t) -> (du .= prob.f(u, p, t); Cint(0))
    elseif !isinplace && typeof(prob.u0) <: AbstractArray
        f! = (du, u, p, t) -> (du .= vec(prob.f(reshape(u, sizeu), p, t)); Cint(0))
    elseif typeof(prob.u0) <: Vector{Float64}
        f! = prob.f
    else # Then it's an in-place function on an abstract array
        f! = (du, u, p, t) -> (prob.f(reshape(du, sizeu), reshape(u, sizeu), p, t);
        du = vec(du);
        0)
    end

    if typeof(alg) <: CVODE_BDF
        alg_code = CV_BDF
    elseif typeof(alg) <: CVODE_Adams
        alg_code = CV_ADAMS
    end

    #if Method == :Newton
    #    method_code = CV_NEWTON
    #elseif Method ==  :Functional
    #    method_code = CV_FUNCTIONAL
    #end

    mem_ptr = CVodeCreate(alg_code)
    (mem_ptr == C_NULL) && error("Failed to allocate CVODE solver object")
    mem = Handle(mem_ptr)

    !verbose && CVodeSetErrHandlerFn(
        mem,
        @cfunction(null_error_handler, Nothing, (Cint, Char, Char, Ptr{Cvoid})),
        C_NULL,
    )

    save_start ? ts = [t0] : ts = Float64[]

    u0nv = NVector(u0)
    _u0 = copy(u0)
    utmp = NVector(_u0)

    userfun = FunJac(
        f!,
        prob.f.jac,
        prob.p,
        nothing,
        prob.f.jac_prototype,
        alg.prec,
        alg.psetup,
        u0,
        _u0,
    )

    function getcfunf(::T) where {T}
        @cfunction(cvodefunjac, Cint, (realtype, N_Vector, N_Vector, Ref{T}))
    end

    flag = CVodeInit(mem, getcfunf(userfun), t0, convert(N_Vector, utmp))

    dt != nothing && (flag = CVodeSetInitStep(mem, dt))
    flag = CVodeSetMinStep(mem, dtmin)
    flag = CVodeSetMaxStep(mem, dtmax)
    flag = CVodeSetUserData(mem, userfun)
    if typeof(abstol) <: Array
        flag = CVodeSVtolerances(mem, reltol, abstol)
    else
        flag = CVodeSStolerances(mem, reltol, abstol)
    end
    flag = CVodeSetMaxNumSteps(mem, maxiters)
    flag = CVodeSetMaxOrd(mem, alg.max_order)
    flag = CVodeSetMaxHnilWarns(mem, alg.max_hnil_warns)
    flag = CVodeSetStabLimDet(mem, alg.stability_limit_detect)
    flag = CVodeSetMaxErrTestFails(mem, alg.max_error_test_failures)
    flag = CVodeSetMaxNonlinIters(mem, alg.max_nonlinear_iters)
    flag = CVodeSetMaxConvFails(mem, alg.max_convergence_failures)

    nojacobian = true

    if Method == :Newton # Only use a linear solver if it's a Newton-based method
        if LinearSolver in (:Dense, :LapackDense)
            nojacobian = false
            A = SUNDenseMatrix(length(u0), length(u0))
            _A = MatrixHandle(A, DenseMatrix())
            if LinearSolver === :Dense
                LS = SUNLinSol_Dense(u0, A)
                _LS = LinSolHandle(LS, Dense())
            else
                LS = SUNLinSol_LapackDense(u0, A)
                _LS = LinSolHandle(LS, LapackDense())
            end
        elseif LinearSolver in (:Band, :LapackBand)
            nojacobian = false
            A = SUNBandMatrix(length(u0), alg.jac_upper, alg.jac_lower)
            _A = MatrixHandle(A, BandMatrix())
            if LinearSolver === :Band
                LS = SUNLinSol_Band(u0, A)
                _LS = LinSolHandle(LS, Band())
            else
                LS = SUNLinSol_LapackBand(u0, A)
                _LS = LinSolHandle(LS, LapackBand())
            end
        elseif LinearSolver == :Diagonal
            nojacobian = false
            flag = CVDiag(mem)
            _A = nothing
            _LS = nothing
        elseif LinearSolver == :GMRES
            LS = SUNLinSol_SPGMR(u0, alg.prec_side, alg.krylov_dim)
            _A = nothing
            _LS = Sundials.LinSolHandle(LS, Sundials.SPGMR())
        elseif LinearSolver == :FGMRES
            LS = SUNLinSol_SPFGMR(u0, alg.prec_side, alg.krylov_dim)
            _A = nothing
            _LS = LinSolHandle(LS, SPFGMR())
        elseif LinearSolver == :BCG
            LS = SUNLinSol_SPBCGS(u0, alg.prec_side, alg.krylov_dim)
            _A = nothing
            _LS = LinSolHandle(LS, SPBCGS())
        elseif LinearSolver == :PCG
            LS = SUNLinSol_PCG(u0, alg.prec_side, alg.krylov_dim)
            _A = nothing
            _LS = LinSolHandle(LS, PCG())
        elseif LinearSolver == :TFQMR
            LS = SUNLinSol_SPTFQMR(u0, alg.prec_side, alg.krylov_dim)
            _A = nothing
            _LS = LinSolHandle(LS, PTFQMR())
        elseif LinearSolver == :KLU
            nojacobian = false
            nnz = length(SparseArrays.nonzeros(prob.f.jac_prototype))
            A = SUNSparseMatrix(length(u0), length(u0), nnz, CSC_MAT)
            LS = SUNLinSol_KLU(u0, A)
            _A = MatrixHandle(A, SparseMatrix())
            _LS = LinSolHandle(LS, KLU())
        end
        if LinearSolver !== :Diagonal
            flag = CVodeSetLinearSolver(mem, LS, _A === nothing ? C_NULL : A)
        end
        NLS = SUNNonlinSol_Newton(u0)
    else
        _A = nothing
        _LS = nothing
        # TODO: Anderson Acceleration
        anderson_m = 0
        NLS = SUNNonlinSol_FixedPoint(u0, anderson_m)
    end
    CVodeSetNonlinearSolver(mem, NLS)

    if DiffEqBase.has_jac(prob.f) && Method == :Newton
        function getcfunjac(::T) where {T}
            @cfunction(
                cvodejac,
                Cint,
                (
                    realtype,
                    N_Vector,
                    N_Vector,
                    SUNMatrix,
                    Ref{T},
                    N_Vector,
                    N_Vector,
                    N_Vector,
                )
            )
        end
        jac = getcfunjac(userfun)
        flag = CVodeSetUserData(mem, userfun)
        nojacobian || (flag = CVodeSetJacFn(mem, jac))
    else
        jac = nothing
    end

    if typeof(prob.f.jac_prototype) <: DiffEqBase.AbstractDiffEqLinearOperator
        function getcfunjtimes(::T) where {T}
            @cfunction(
                jactimes,
                Cint,
                (N_Vector, N_Vector, realtype, N_Vector, N_Vector, Ref{T}, N_Vector)
            )
        end
        jtimes = getcfunjtimes(userfun)
        CVodeSetJacTimes(mem, C_NULL, jtimes)
    end

    if alg.prec !== nothing
        function getpercfun(::T) where {T}
            @cfunction(
                precsolve,
                Cint,
                (
                    Float64,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    Float64,
                    Float64,
                    Int,
                    Ref{T},
                )
            )
        end
        precfun = getpercfun(userfun)

        function getpsetupfun(::T) where {T}
            @cfunction(
                precsetup,
                Cint,
                (Float64, N_Vector, N_Vector, Int, Ptr{Int}, Float64, Ref{T})
            )
        end
        psetupfun = alg.psetup === nothing ? C_NULL : getpsetupfun(userfun)

        CVodeSetPreconditioner(mem, psetupfun, precfun)
    end

    callbacks_internal == nothing ? tmp = nothing : tmp = similar(u0)
    callbacks_internal == nothing ? uprev = nothing : uprev = similar(u0)
    tout = [tspan[1]]

    if save_start
        if save_idxs === nothing
            ures = Vector{uType}()
            dures = Vector{uType}()
            save_value!(ures, u0, uType, sizeu, save_idxs)
            if dense
                f!(_u0, u0, prob.p, tspan[1])
                save_value!(dures, utmp, uType, sizeu, save_idxs)
            end
        else
            ures = u0[save_idxs]
            if dense
                f!(_u0, u0, prob.p, tspan[1])
                dures = _u0[save_idxs]
            end
        end
    else
        ures = Vector{uType}()
        dures = Vector{uType}()
    end

    sol = DiffEqBase.build_solution(
        prob,
        alg,
        ts,
        ures,
        dense = dense,
        interp = dense ? DiffEqBase.HermiteInterpolation(ts, ures, dures) :
                     DiffEqBase.LinearInterpolation(ts, ures),
        timeseries_errors = timeseries_errors,
        destats = DiffEqBase.DEStats(0),
        calculate_error = false,
    )
    opts = DEOptions(
        saveat_internal,
        tstops_internal,
        save_everystep,save_idxs,
        dense,
        timeseries_errors,
        dense_errors,
        save_on,
        save_end,
        callbacks_internal,
        abstol,
        reltol,
        verbose,
        advance_to_tstop,
        stop_at_next_tstop,
        progress,
        progress_name,
        progress_message,
        maxiters,
    )
    integrator = CVODEIntegrator(
        u0,
        prob.p,
        t0,
        t0,
        mem,
        _LS,
        _A,
        sol,
        alg,
        f!,
        userfun,
        jac,
        opts,
        tout,
        tdir,
        sizeu,
        false,
        tmp,
        uprev,
        Cint(flag),
        false,
        0,
        1,
        callback_cache,
        0.0,
    )

    initialize_callbacks!(integrator)
    integrator
end # function solve

function DiffEqBase.__init(
    prob::DiffEqBase.AbstractODEProblem{uType, tupType, isinplace},
    alg::ARKODE{Method, LinearSolver, MassLinearSolver},
    timeseries = [],
    ts = [],
    ks = [];
    verbose = true,
    callback = nothing,
    abstol = 1 / 10^6,
    reltol = 1 / 10^3,
    saveat = Float64[],
    tstops = Float64[],
    maxiters = Int(1e5),
    dt = nothing,
    dtmin = 0.0,
    dtmax = 0.0,
    timeseries_errors = true,
    dense_errors = false,
    save_everystep = isempty(saveat), save_idxs = nothing,
    dense = save_everystep,
    save_on = true,
    save_start = true,
    save_end = true,
    save_timeseries = nothing,
    progress = false,
    progress_name = "ODE",
    progress_message = DiffEqBase.ODE_DEFAULT_PROG_MESSAGE,
    advance_to_tstop = false,
    stop_at_next_tstop = false,
    userdata = nothing,
    alias_u0 = false,
    kwargs...,
) where {uType, tupType, isinplace, Method, LinearSolver, MassLinearSolver}

    tType = eltype(tupType)

    if verbose
        warned = !isempty(kwargs) && DiffEqBase.check_keywords(alg, kwargs, warnlist)
        warned && DiffEqBase.warn_compat()
    end

    if typeof(reltol) <: AbstractArray
        error("Sundials only allows scalar reltol.")
    end

    progress && Logging.@logmsg(-1, progress_name, _id = _id = :Sundials, progress = 0)

    callbacks_internal = DiffEqBase.CallbackSet(callback)

    max_len_cb = DiffEqBase.max_vector_callback_length(callbacks_internal)
    if max_len_cb isa VectorContinuousCallback
        callback_cache = DiffEqBase.CallbackCache(max_len_cb.len, Float64, Float64)
    else
        callback_cache = nothing
    end

    tspan = prob.tspan
    t0 = tspan[1]

    tdir = sign(tspan[2] - tspan[1])

    tstops_internal, saveat_internal =
        tstop_saveat_disc_handling(tstops, saveat, tdir, tspan, tType)

    if typeof(prob.u0) <: Number
        u0 = [prob.u0]
    else
        if alias_u0
            u0 = vec(prob.u0)
        else
            u0 = vec(copy(prob.u0))
        end
    end

    sizeu = size(prob.u0)
    save_start ? ts = [t0] : ts = Float64[]
    u0nv = NVector(u0)
    _u0 = copy(u0)
    utmp = NVector(_u0)

    function arkodemem(; fe = C_NULL, fi = C_NULL, t0 = t0, u0 = convert(N_Vector, u0nv))
        mem_ptr = ARKStepCreate(fe, fi, t0, u0)
        (mem_ptr == C_NULL) && error("Failed to allocate ARKODE solver object")
        mem = Handle(mem_ptr)

        !verbose && ARKStepSetErrHandlerFn(
            mem,
            @cfunction(null_error_handler, Nothing, (Cint, Char, Char, Ptr{Cvoid})),
            C_NULL,
        )
        return mem
    end

    ### Fix the more general function to Sundials allowed style
    if !isinplace && typeof(prob.u0) <: Number
        f! = (du, u, p, t) -> (du .= prob.f(first(u), p, t); Cint(0))
    elseif !isinplace && typeof(prob.u0) <: Vector{Float64}
        f! = (du, u, p, t) -> (du .= prob.f(u, p, t); Cint(0))
    elseif !isinplace && typeof(prob.u0) <: AbstractArray
        f! = (du, u, p, t) -> (du .= vec(prob.f(reshape(u, sizeu), p, t)); Cint(0))
    elseif typeof(prob.u0) <: Vector{Float64}
        f! = prob.f
    else # Then it's an in-place function on an abstract array
        f! = (du, u, p, t) -> (prob.f(reshape(du, sizeu), reshape(u, sizeu), p, t);
        du = vec(du);
        Cint(0))
    end

    if typeof(prob.problem_type) <: SplitODEProblem

        ### Fix the more general function to Sundials allowed style
        if !isinplace && typeof(prob.u0) <: Number
            f1! = (du, u, p, t) -> (du .= prob.f.f1(first(u), p, t); Cint(0))
            f2! = (du, u, p, t) -> (du .= prob.f.f2(first(u), p, t); Cint(0))
        elseif !isinplace && typeof(prob.u0) <: Vector{Float64}
            f1! = (du, u, p, t) -> (du .= prob.f.f1(u, p, t); Cint(0))
            f2! = (du, u, p, t) -> (du .= prob.f.f2(u, p, t); Cint(0))
        elseif !isinplace && typeof(prob.u0) <: AbstractArray
            f1! = (du, u, p, t) -> (du .= vec(prob.f.f1(reshape(u, sizeu), p, t)); Cint(0))
            f2! = (du, u, p, t) -> (du .= vec(prob.f.f2(reshape(u, sizeu), p, t)); Cint(0))
        elseif typeof(prob.u0) <: Vector{Float64}
            f1! = prob.f.f1
            f2! = prob.f.f2
        else # Then it's an in-place function on an abstract array
            f1! =
                (du, u, p, t) -> (prob.f.f1(reshape(du, sizeu), reshape(u, sizeu), p, t);
                du = vec(du);
                Cint(0))
            f2! =
                (du, u, p, t) -> (prob.f.f2(reshape(du, sizeu), reshape(u, sizeu), p, t);
                du = vec(du);
                Cint(0))
        end

        userfun = FunJac(
            f1!,
            f2!,
            prob.f.f1.jac,
            prob.p,
            prob.f.mass_matrix,
            prob.f.f1.jac_prototype,
            alg.prec,
            alg.psetup,
            u0,
            _u0,
            nothing,
        )

        function getcfunjac(::T) where {T}
            @cfunction(cvodefunjac, Cint, (realtype, N_Vector, N_Vector, Ref{T}))
        end
        function getcfunjac2(::T) where {T}
            @cfunction(cvodefunjac2, Cint, (realtype, N_Vector, N_Vector, Ref{T}))
        end
        cfj1 = getcfunjac(userfun)
        cfj2 = getcfunjac2(userfun)

        mem = arkodemem(fi = cfj1, fe = cfj2)
    else
        userfun = FunJac(
            f!,
            prob.f.jac,
            prob.p,
            prob.f.mass_matrix,
            prob.f.jac_prototype,
            alg.prec,
            alg.psetup,
            u0,
            _u0,
        )
        if alg.stiffness == Explicit()
            function getcfun1(::T) where {T}
                @cfunction(cvodefunjac, Cint, (realtype, N_Vector, N_Vector, Ref{T}))
            end
            cfj1 = getcfun1(userfun)
            mem = arkodemem(fe = cfj1)
        elseif alg.stiffness == Implicit()
            function getcfun2(::T) where {T}
                @cfunction(cvodefunjac, Cint, (realtype, N_Vector, N_Vector, Ref{T}))
            end
            cfj2 = getcfun2(userfun)
            mem = arkodemem(fi = cfj2)
        end
    end

    dt != nothing && (flag = ARKStepSetInitStep(mem, dt))
    flag = ARKStepSetMinStep(mem, dtmin)
    flag = ARKStepSetMaxStep(mem, dtmax)
    flag = ARKStepSetUserData(mem, userfun)
    if typeof(abstol) <: Array
        flag = ARKStepSVtolerances(mem, reltol, abstol)
    else
        flag = ARKStepSStolerances(mem, reltol, abstol)
    end
    flag = ARKStepSetMaxNumSteps(mem, maxiters)
    flag = ARKStepSetMaxHnilWarns(mem, alg.max_hnil_warns)
    flag = ARKStepSetMaxErrTestFails(mem, alg.max_error_test_failures)
    flag = ARKStepSetMaxConvFails(mem, alg.max_convergence_failures)
    flag = ARKStepSetPredictorMethod(mem, alg.predictor_method)
    flag = ARKStepSetNonlinConvCoef(mem, alg.nonlinear_convergence_coefficient)
    flag = ARKStepSetDenseOrder(mem, alg.dense_order)

    #=
    Reference from Manual on ARKODE
    To choose an explicit table, set itable to a negative value. This automatically calls ARKStepSetExplicit(). However, if the problem is posed in explicit form, i.e. 𝑦 ̇ = 𝑓 (𝑡, 𝑦), then we recommend that the ERKStep time- stepper module be used instead of ARKStep.
    To select an implicit table, set etable to a negative value. This automatically calls ARKStepSetImplicit(). If both itable and etable are non-negative, then these should match an existing implicit/explicit pair, listed in the section Additive Butcher tables. This automatically calls ARKStepSetImEx().
    =#
    if alg.itable == nothing && alg.etable == nothing
        flag = ARKStepSetOrder(mem, alg.order)
    elseif alg.itable == nothing && alg.etable != nothing
        flag = ARKStepSetTableNum(mem, -1, alg.etable)
    elseif alg.itable != nothing && alg.etable == nothing
        flag = ARKStepSetTableNum(mem, alg.itable, -1)
    else
        flag = ARKStepSetTableNum(mem, alg.itable, alg.etable)
    end

    flag = ARKStepSetNonlinCRDown(mem, alg.crdown)
    flag = ARKStepSetNonlinRDiv(mem, alg.rdiv)
    flag = ARKStepSetDeltaGammaMax(mem, alg.dgmax)
    flag = ARKStepSetMaxStepsBetweenLSet(mem, alg.msbp)
    #flag = ARKStepSetAdaptivityMethod(mem,alg.adaptivity_method,1,0)

    #flag = ARKStepSetFixedStep(mem,)
    alg.set_optimal_params && (flag = ARKStepSetOptimalParams(mem))

    if Method == :Newton && alg.stiffness !== Explicit() # Only use a linear solver if it's a Newton-based method
        if LinearSolver in (:Dense, :LapackDense)
            nojacobian = false
            A = SUNDenseMatrix(length(u0), length(u0))
            _A = MatrixHandle(A, DenseMatrix())
            if LinearSolver === :Dense
                LS = SUNLinSol_Dense(u0, A)
                _LS = LinSolHandle(LS, Dense())
            else
                LS = SUNLinSol_LapackDense(u0, A)
                _LS = LinSolHandle(LS, LapackDense())
            end
        elseif LinearSolver in (:Band, :LapackBand)
            nojacobian = false
            A = SUNBandMatrix(length(u0), alg.jac_upper, alg.jac_lower)
            _A = MatrixHandle(A, BandMatrix())
            if LinearSolver === :Band
                LS = SUNLinSol_Band(u0, A)
                _LS = LinSolHandle(LS, Band())
            else
                LS = SUNLinSol_LapackBand(u0, A)
                _LS = LinSolHandle(LS, LapackBand())
            end
        elseif LinearSolver == :GMRES
            LS = SUNLinSol_SPGMR(u0, alg.prec_side, alg.krylov_dim)
            _A = nothing
            _LS = Sundials.LinSolHandle(LS, Sundials.SPGMR())
        elseif LinearSolver == :FGMRES
            LS = SUNLinSol_SPFGMR(u0, alg.prec_side, alg.krylov_dim)
            _A = nothing
            _LS = LinSolHandle(LS, SPFGMR())
        elseif LinearSolver == :BCG
            LS = SUNLinSol_SPBCGS(u0, alg.prec_side, alg.krylov_dim)
            _A = nothing
            _LS = LinSolHandle(LS, SPBCGS())
        elseif LinearSolver == :PCG
            LS = SUNLinSol_PCG(u0, alg.prec_side, alg.krylov_dim)
            _A = nothing
            _LS = LinSolHandle(LS, PCG())
        elseif LinearSolver == :TFQMR
            LS = SUNLinSol_SPTFQMR(u0, alg.prec_side, alg.krylov_dim)
            _A = nothing
            _LS = LinSolHandle(LS, PTFQMR())
        elseif LinearSolver == :KLU
            nnz = length(SparseArrays.nonzeros(prob.f.jac_prototype))
            A = SUNSparseMatrix(length(u0), length(u0), nnz, CSC_MAT)
            LS = SUNLinSol_KLU(u0, A)
            _A = MatrixHandle(A, SparseMatrix())
            _LS = LinSolHandle(LS, KLU())
        end
        flag = ARKStepSetLinearSolver(mem, LS, _A === nothing ? C_NULL : A)
        flag = ARKStepSetMaxNonlinIters(mem, alg.max_nonlinear_iters)
    elseif Method == :Functional && alg.stiffness !== Explicit()
        ARKStepSetFixedPoint(mem, Clong(alg.krylov_dim))
    else
        _A = nothing
        _LS = nothing
    end

    if (
        typeof(prob.problem_type) <: SplitODEProblem &&
        typeof(prob.f.f1.jac_prototype) <: DiffEqBase.AbstractDiffEqLinearOperator
    ) || (
        !(typeof(prob.problem_type) <: SplitODEProblem) &&
        typeof(prob.f.jac_prototype) <: DiffEqBase.AbstractDiffEqLinearOperator
    ) && alg.stiffness !== Explicit()
        function getcfunjtimes(::T) where {T}
            @cfunction(
                jactimes,
                Cint,
                (N_Vector, N_Vector, realtype, N_Vector, N_Vector, Ref{T}, N_Vector)
            )
        end
        jtimes = getcfunjtimes(userfun)
        ARKStepSetJacTimes(mem, C_NULL, jtimes)
    end

    if prob.f.mass_matrix != LinearAlgebra.I && alg.stiffness !== Explicit()
        if MassLinearSolver in (:Dense, :LapackDense)
            nojacobian = false
            M = SUNDenseMatrix(length(u0), length(u0))
            _M = MatrixHandle(M, DenseMatrix())
            if MassLinearSolver === :Dense
                MLS = SUNLinSol_Dense(u0, M)
                _MLS = LinSolHandle(MLS, Dense())
            else
                MLS = SUNLinSol_LapackDense(u0, M)
                _MLS = LinSolHandle(MLS, LapackDense())
            end
        elseif MassLinearSolver in (:Band, :LapackBand)
            nojacobian = false
            M = SUNBandMatrix(length(u0), alg.jac_upper, alg.jac_lower)
            _M = MatrixHandle(M, BandMatrix())
            if MassLinearSolver === :Band
                MLS = SUNLinSol_Band(u0, M)
                _MLS = LinSolHandle(MLS, Band())
            else
                MLS = SUNLinSol_LapackBand(u0, M)
                _MLS = LinSolHandle(MLS, LapackBand())
            end
        elseif MassLinearSolver == :GMRES
            MLS = SUNLinSol_SPGMR(u0, alg.prec_side, alg.mass_krylov_dim)
            _M = nothing
            _MLS = LinSolHandle(MLS, SPGMR())
        elseif MassLinearSolver == :FGMRES
            MLS = SUNLinSol_SPGMR(u0, alg.prec_side, alg.mass_krylov_dim)
            _M = nothing
            _MLS = LinSolHandle(MLS, SPFGMR())
        elseif MassLinearSolver == :BCG
            MLS = SUNLinSol_SPGMR(u0, alg.prec_side, alg.mass_krylov_dim)
            _M = nothing
            _MLS = LinSolHandle(MLS, SPBCGS())
        elseif MassLinearSolver == :PCG
            MLS = SUNLinSol_SPGMR(u0, alg.prec_side, alg.mass_krylov_dim)
            _M = nothing
            _MLS = LinSolHandle(MLS, PCG())
        elseif MassLinearSolver == :TFQMR
            MLS = SUNLinSol_SPGMR(u0, alg.prec_side, alg.mass_krylov_dim)
            _M = nothing
            _MLS = LinSolHandle(MLS, PTFQMR())
        elseif MassLinearSolver == :KLU
            nnz = length(SparseArrays.nonzeros(prob.f.mass_matrix))
            M = SUNSparseMatrix(length(u0), length(u0), nnz, CSC_MAT)
            MLS = SUNLinSol_KLU(u0, M)
            _M = MatrixHandle(M, SparseMatrix())
            _MLS = LinSolHandle(MLS, KLU())
        end
        flag = ARKStepSetMassLinearSolver(mem, MLS, _M === nothing ? C_NULL : M, false)
        function getmatfun(::T) where {T}
            @cfunction(
                massmat,
                Cint,
                (realtype, SUNMatrix, Ref{T}, N_Vector, N_Vector, N_Vector)
            )
        end
        matfun = getmatfun(userfun)
        ARKStepSetMassFn(mem, matfun)
    else
        _M = nothing
        _MLS = nothing
    end

    if DiffEqBase.has_jac(prob.f) && alg.stiffness !== Explicit()
        function getfunjac(::T) where {T}
            @cfunction(
                cvodejac,
                Cint,
                (
                    realtype,
                    N_Vector,
                    N_Vector,
                    SUNMatrix,
                    Ref{T},
                    N_Vector,
                    N_Vector,
                    N_Vector,
                )
            )
        end
        jac = getfunjac(userfun)
        flag = ARKStepSetUserData(mem, userfun)
        flag = ARKStepSetJacFn(mem, jac)
    else
        jac = nothing
    end

    if alg.prec !== nothing && alg.stiffness !== Explicit()
        function getpercfun(::T) where {T}
            @cfunction(
                precsolve,
                Cint,
                (
                    Float64,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    Float64,
                    Float64,
                    Int,
                    Ref{T},
                )
            )
        end
        precfun = getpercfun(userfun)

        function getpsetupfun(::T) where {T}
            @cfunction(
                precsetup,
                Cint,
                (Float64, N_Vector, N_Vector, Int, Ptr{Int}, Float64, Ref{T})
            )
        end
        psetupfun = alg.psetup === nothing ? C_NULL : getpsetupfun(userfun)

        ARKStepSetPreconditioner(mem, psetupfun, precfun)
    end

    callbacks_internal == nothing ? tmp = nothing : tmp = similar(u0)
    callbacks_internal == nothing ? uprev = nothing : uprev = similar(u0)
    tout = [tspan[1]]

    if save_start
        if save_idxs === nothing
            ures = Vector{uType}()
            dures = Vector{uType}()
            save_value!(ures, u0, uType, sizeu, save_idxs)
            if dense
                f!(_u0, u0, prob.p, tspan[1])
                save_value!(dures, utmp, uType, sizeu, save_idxs)
            end
        else
            ures = u0[save_idxs]
            if dense
                f!(_u0, u0, prob.p, tspan[1])
                dures = _u0[save_idxs]
            end
        end
    else
        ures = Vector{uType}()
        dures = Vector{uType}()
    end

    sol = DiffEqBase.build_solution(
        prob,
        alg,
        ts,
        ures,
        dense = dense,
        interp = dense ? DiffEqBase.HermiteInterpolation(ts, ures, dures) :
                     DiffEqBase.LinearInterpolation(ts, ures),
        timeseries_errors = timeseries_errors,
        destats = DiffEqBase.DEStats(0),
        calculate_error = false,
    )
    opts = DEOptions(
        saveat_internal,
        tstops_internal,
        save_everystep,save_idxs,
        dense,
        timeseries_errors,
        dense_errors,
        save_on,
        save_end,
        callbacks_internal,
        abstol,
        reltol,
        verbose,
        advance_to_tstop,
        stop_at_next_tstop,
        progress,
        progress_name,
        progress_message,
        maxiters,
    )
    integrator = ARKODEIntegrator(
        utmp,
        prob.p,
        t0,
        t0,
        mem,
        _LS,
        _A,
        _MLS,
        _M,
        sol,
        alg,
        f!,
        userfun,
        jac,
        opts,
        tout,
        tdir,
        sizeu,
        false,
        tmp,
        uprev,
        Cint(flag),
        false,
        0,
        1,
        callback_cache,
        0.0,
    )

    initialize_callbacks!(integrator)
    integrator
end # function solve

function tstop_saveat_disc_handling(tstops, saveat, tdir, tspan, tType)

    if isempty(tstops) # TODO: Specialize more
        tstops_vec = [tspan[2]]
    else
        tstops_vec = vec(collect(
            tType,
            Iterators.filter(
                x -> tdir * tspan[1] < tdir * x ≤ tdir * tspan[end],
                Iterators.flatten((tstops, tspan[end])),
            ),
        ))
    end

    if tdir > 0
        tstops_internal = DataStructures.BinaryMinHeap(tstops_vec)
    else
        tstops_internal = DataStructures.BinaryMaxHeap(tstops_vec)
    end

    if typeof(saveat) <: Number
        if (tspan[1]:saveat:tspan[end])[end] == tspan[end]
            saveat_vec = convert(
                Vector{tType},
                collect(tType, (tspan[1] + saveat):saveat:tspan[end]),
            )
        else
            saveat_vec = convert(
                Vector{tType},
                collect(tType, (tspan[1] + saveat):saveat:(tspan[end] - saveat)),
            )
        end
    elseif isempty(saveat)
        saveat_vec = saveat
    else
        saveat_vec = vec(collect(
            tType,
            Iterators.filter(x -> tdir * tspan[1] < tdir * x < tdir * tspan[end], saveat),
        ))
    end

    if tdir > 0
        saveat_internal = DataStructures.BinaryMinHeap(saveat_vec)
    else
        saveat_internal = DataStructures.BinaryMaxHeap(saveat_vec)
    end

    tstops_internal, saveat_internal
end

## Solve for DAEs uses IDA

function DiffEqBase.__init(
    prob::DiffEqBase.AbstractDAEProblem{uType, duType, tupType, isinplace},
    alg::SundialsDAEAlgorithm{LinearSolver},
    timeseries = [],
    ts = [],
    ks = [];
    verbose = true,
    dt = nothing,
    dtmax = 0.0,
    save_on = true,
    save_start = true,
    callback = nothing,
    abstol = 1 / 10^6,
    reltol = 1 / 10^3,
    saveat = Float64[],
    tstops = Float64[],
    maxiters = Int(1e5),
    timeseries_errors = true,
    dense_errors = false,
    save_everystep = isempty(saveat), save_idxs = nothing,
    dense = save_everystep,
    save_timeseries = nothing,
    save_end = true,
    progress = false,
    progress_name = "ODE",
    progress_message = DiffEqBase.ODE_DEFAULT_PROG_MESSAGE,
    advance_to_tstop = false,
    stop_at_next_tstop = false,
    userdata = nothing,
    kwargs...,
) where {uType, duType, tupType, isinplace, LinearSolver}

    tType = eltype(tupType)

    if verbose
        warned = !isempty(kwargs) && DiffEqBase.check_keywords(alg, kwargs, warnida)
        warned && DiffEqBase.warn_compat()
    end

    if typeof(reltol) <: AbstractArray
        error("Sundials only allows scalar reltol.")
    end

    progress && Logging.@logmsg(-1, progress_name, _id = _id = :Sundials, progress = 0)

    callbacks_internal = DiffEqBase.CallbackSet(callback)

    max_len_cb = DiffEqBase.max_vector_callback_length(callbacks_internal)
    if max_len_cb isa VectorContinuousCallback
        callback_cache = DiffEqBase.CallbackCache(max_len_cb.len, Float64, Float64)
    else
        callback_cache = nothing
    end

    tspan = prob.tspan
    t0 = tspan[1]

    tdir = sign(tspan[2] - tspan[1])

    tstops_internal, saveat_internal =
        tstop_saveat_disc_handling(tstops, saveat, tdir, tspan, tType)

    if typeof(prob.u0) <: Number
        u0 = [prob.u0]
    else
        u0 = vec(copy(prob.u0))
    end

    if typeof(prob.du0) <: Number
        du0 = [prob.du0]
    else
        du0 = vec(copy(prob.du0))
    end

    sizeu = size(prob.u0)
    sizedu = size(prob.du0)

    ### Fix the more general function to Sundials allowed style
    if !isinplace && typeof(prob.u0) <: Number
        f! = (out, du, u, p, t) -> (out .= prob.f(first(du), first(u), p, t); Cint(0))
    elseif !isinplace && typeof(prob.u0) <: Vector{Float64}
        f! = (out, du, u, p, t) -> (out .= prob.f(du, u, p, t); Cint(0))
    elseif !isinplace && typeof(prob.u0) <: AbstractArray
        f! =
            (out, du, u, p, t) ->
                (out .= vec(prob.f(reshape(du, sizedu), reshape(u, sizeu), p, t));
                Cint(0))
    elseif typeof(prob.u0) <: Vector{Float64}
        f! = prob.f
    else # Then it's an in-place function on an abstract array
        f! =
            (out, du, u, p, t) ->
                (prob.f(reshape(out, sizeu), reshape(du, sizedu), reshape(u, sizeu), p, t);
                Cint(0))
    end

    mem_ptr = IDACreate()
    (mem_ptr == C_NULL) && error("Failed to allocate IDA solver object")
    mem = Handle(mem_ptr)

    !verbose && IDASetErrHandlerFn(
        mem,
        @cfunction(null_error_handler, Nothing, (Cint, Char, Char, Ptr{Cvoid})),
        C_NULL,
    )

    ts = [t0]

    _u0 = copy(u0)
    utmp = NVector(_u0)
    _du0 = copy(du0)
    dutmp = NVector(_du0)
    rtest = zeros(length(u0))

    userfun = FunJac(
        f!,
        prob.f.jac,
        prob.p,
        nothing,
        prob.f.jac_prototype,
        alg.prec,
        alg.psetup,
        _u0,
        _du0,
        rtest,
    )

    u0nv = NVector(u0)

    function getcfun(::T) where {T}
        @cfunction(idasolfun, Cint, (realtype, N_Vector, N_Vector, N_Vector, Ref{T}))
    end
    cfun = getcfun(userfun)
    flag = IDAInit(mem, cfun, t0, convert(N_Vector, utmp), convert(N_Vector, dutmp))
    dt != nothing && (flag = IDASetInitStep(mem, dt))
    flag = IDASetUserData(mem, userfun)
    flag = IDASetMaxStep(mem, dtmax)
    if typeof(abstol) <: Array
        flag = IDASVtolerances(mem, reltol, abstol)
    else
        flag = IDASStolerances(mem, reltol, abstol)
    end
    flag = IDASetMaxNumSteps(mem, maxiters)
    flag = IDASetMaxOrd(mem, alg.max_order)
    flag = IDASetMaxErrTestFails(mem, alg.max_error_test_failures)
    flag = IDASetNonlinConvCoef(mem, alg.nonlinear_convergence_coefficient)
    flag = IDASetMaxNonlinIters(mem, alg.max_nonlinear_iters)
    flag = IDASetMaxConvFails(mem, alg.max_convergence_failures)
    flag = IDASetNonlinConvCoefIC(mem, alg.nonlinear_convergence_coefficient_ic)
    flag = IDASetMaxNumStepsIC(mem, alg.max_num_steps_ic)
    flag = IDASetMaxNumJacsIC(mem, alg.max_num_jacs_ic)
    flag = IDASetMaxNumItersIC(mem, alg.max_num_iters_ic)
    #flag = IDASetMaxBacksIC(mem,alg.max_num_backs_ic) # Needs newer version?
    flag = IDASetLineSearchOffIC(mem, alg.use_linesearch_ic)

    prec_side = isnothing(alg.prec) ? 0 : 1  # IDA only supports left preconditioning (prec_side = 1)
    if LinearSolver in (:Dense, :LapackDense)
        nojacobian = false
        A = SUNDenseMatrix(length(u0), length(u0))
        _A = MatrixHandle(A, DenseMatrix())
        if LinearSolver === :Dense
            LS = SUNLinSol_Dense(u0, A)
            _LS = LinSolHandle(LS, Dense())
        else
            LS = SUNLinSol_LapackDense(u0, A)
            _LS = LinSolHandle(LS, LapackDense())
        end
    elseif LinearSolver in (:Band, :LapackBand)
        nojacobian = false
        A = SUNBandMatrix(length(u0), alg.jac_upper, alg.jac_lower)
        _A = MatrixHandle(A, BandMatrix())
        if LinearSolver === :Band
            LS = SUNLinSol_Band(u0, A)
            _LS = LinSolHandle(LS, Band())
        else
            LS = SUNLinSol_LapackBand(u0, A)
            _LS = LinSolHandle(LS, LapackBand())
        end
    elseif LinearSolver == :GMRES
        LS = SUNLinSol_SPGMR(u0, prec_side, alg.krylov_dim)
        _A = nothing
        _LS = LinSolHandle(LS, SPGMR())
    elseif LinearSolver == :FGMRES
        LS = SUNLinSol_SPFGMR(u0, prec_side, alg.krylov_dim)
        _A = nothing
        _LS = LinSolHandle(LS, SPFGMR())
    elseif LinearSolver == :BCG
        LS = SUNLinSol_SPBCGS(u0, prec_side, alg.krylov_dim)
        _A = nothing
        _LS = LinSolHandle(LS, SPBCGS())
    elseif LinearSolver == :PCG
        LS = SUNLinSol_PCG(u0, prec_side, alg.krylov_dim)
        _A = nothing
        _LS = LinSolHandle(LS, PCG())
    elseif LinearSolver == :TFQMR
        LS = SUNLinSol_SPTFQMR(u0, prec_side, alg.krylov_dim)
        _A = nothing
        _LS = LinSolHandle(LS, PTFQMR())
    elseif LinearSolver == :KLU
        nnz = length(SparseArrays.nonzeros(prob.f.jac_prototype))
        A = SUNSparseMatrix(length(u0), length(u0), nnz, Sundials.CSC_MAT)
        LS = SUNLinSol_KLU(u0, A)
        _A = MatrixHandle(A, SparseMatrix())
        _LS = LinSolHandle(LS, KLU())
    end
    flag = IDASetLinearSolver(mem, LS, _A === nothing ? C_NULL : A)

    if typeof(prob.f.jac_prototype) <: DiffEqBase.AbstractDiffEqLinearOperator
        function getcfunjtimes(::T) where {T}
            @cfunction(
                idajactimes,
                Cint,
                (
                    realtype,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    realtype,
                    Ref{T},
                    N_Vector,
                    N_Vector,
                )
            )
        end
        jtimes = getcfunjtimes(userfun)
        IDASetJacTimes(mem, C_NULL, jtimes)
    end

    if alg.prec !== nothing
        function getprecfun(::T) where {T}
            @cfunction(
                idaprecsolve,
                Cint,
                (
                    Float64,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    Float64,
                    Float64,
                    Ref{T},
                )
            )
        end
        precfun = getprecfun(userfun)

        function getpsetupfun(::T) where {T}
            @cfunction(
                idaprecsetup,
                Cint,
                (Float64, N_Vector, N_Vector, N_Vector, Float64, Ref{T})
            )
        end
        psetupfun = alg.psetup === nothing ? C_NULL : getpsetupfun(userfun)

        IDASetPreconditioner(mem, psetupfun, precfun)
    end

    if DiffEqBase.has_jac(prob.f)
        function getcfunjacc(::T) where {T}
            @cfunction(
                idajac,
                Cint,
                (
                    realtype,
                    realtype,
                    N_Vector,
                    N_Vector,
                    N_Vector,
                    SUNMatrix,
                    Ref{T},
                    N_Vector,
                    N_Vector,
                    N_Vector,
                )
            )
        end
        jac = getcfunjacc(userfun)
        flag = IDASetUserData(mem, userfun)
        flag = IDASetJacFn(mem, jac)
    else
        jac = nothing
    end

    tout = [tspan[1]]

    f!(rtest, du0, u0, prob.p, t0)
    if any(abs.(rtest) .>= reltol)
        if prob.differential_vars === nothing && !alg.init_all
            error("Must supply differential_vars argument to DAEProblem constructor to use IDA initial value solver.")
        end
        prob.differential_vars != nothing &&
            (flag = IDASetId(mem, collect(Float64, prob.differential_vars)))

        if dt != nothing
            _t = float(dt)
        else
            _t = float(tspan[2])
        end
        if alg.init_all
            init_type = IDA_Y_INIT
        else
            init_type = IDA_YA_YDP_INIT
        end
        flag = IDACalcIC(mem, init_type, _t)
    end

    if save_start
        if save_idxs === nothing
            ures = Vector{uType}()
            dures = Vector{uType}()
            save_value!(ures, u0, uType, sizeu, save_idxs)
            if dense
                save_value!(dures, du0, uType, sizedu, save_idxs)
            end
        else
            ures = u0[save_idxs]
            if dense
                dures = du0[save_idxs]
            end
        end
    else
        ures = Vector{uType}()
        dures = Vector{uType}()
    end

    callbacks_internal == nothing ? tmp = nothing : tmp = similar(u0)
    callbacks_internal == nothing ? uprev = nothing : uprev = similar(u0)

    if flag >= 0
        retcode = :Default
    else
        retcode = :InitialFailure
    end

    sol = DiffEqBase.build_solution(
        prob,
        alg,
        ts,
        ures,
        dense = dense,
        interp = dense ? DiffEqBase.HermiteInterpolation(ts, ures, dures) :
                     DiffEqBase.LinearInterpolation(ts, ures),
        calculate_error = false,
        timeseries_errors = timeseries_errors,
        retcode = retcode,
        destats = DiffEqBase.DEStats(0),
        dense_errors = dense_errors,
    )

    opts = DEOptions(
        saveat_internal,
        tstops_internal,
        save_everystep,save_idxs,
        dense,
        timeseries_errors,
        dense_errors,
        save_on,
        save_end,
        callbacks_internal,
        abstol,
        reltol,
        verbose,
        advance_to_tstop,
        stop_at_next_tstop,
        progress,
        progress_name,
        progress_message,
        maxiters,
    )

    integrator = IDAIntegrator(
        utmp,
        dutmp,
        prob.p,
        t0,
        t0,
        mem,
        _LS,
        _A,
        sol,
        alg,
        f!,
        userfun,
        jac,
        opts,
        tout,
        tdir,
        sizeu,
        sizedu,
        false,
        tmp,
        uprev,
        Cint(flag),
        false,
        0,
        1,
        callback_cache,
        0.0,
    )

    initialize_callbacks!(integrator)
    integrator
end # function solve

## Common calls

function interpret_sundials_retcode(flag)
    flag >= 0 && return :Success
    flag == -1 && return :MaxIters
    (flag == -2 || flag == -3) && return :Unstable
    flag == -4 && return :ConvergenceFailure
    return :Failure
end

function solver_step(integrator::CVODEIntegrator, tstop)
    integrator.flag =
        CVode(integrator.mem, tstop, integrator.u, integrator.tout, CV_ONE_STEP)
    if integrator.opts.progress
        Logging.@logmsg(
            -1,
            integrator.opts.progress_name,
            _id = :Sundials,
            message = integrator.opts.progress_message(
                integrator.dt,
                integrator.u,
                integrator.p,
                integrator.t,
            ),
            progress = integrator.t / integrator.sol.prob.tspan[2]
        )
    end
end
function solver_step(integrator::ARKODEIntegrator, tstop)
    integrator.flag =
        ARKStepEvolve(integrator.mem, tstop, integrator.u, integrator.tout, ARK_ONE_STEP)
    if integrator.opts.progress
        Logging.@logmsg(
            -1,
            integrator.opts.progress_name,
            _id = :Sundials,
            message = integrator.opts.progress_message(
                integrator.dt,
                integrator.u,
                integrator.p,
                integrator.t,
            ),
            progress = integrator.t / integrator.sol.prob.tspan[2]
        )
    end
end
function solver_step(integrator::IDAIntegrator, tstop)
    integrator.flag = IDASolve(
        integrator.mem,
        tstop,
        integrator.tout,
        integrator.u,
        integrator.du,
        IDA_ONE_STEP,
    )
    if integrator.opts.progress
        Logging.@logmsg(
            -1,
            integrator.opts.progress_name,
            _id = :Sundials,
            message = integrator.opts.progress_message(
                integrator.dt,
                integrator.u,
                integrator.p,
                integrator.t,
            ),
            progress = integrator.t / integrator.sol.prob.tspan[2]
        )
    end
end

function set_stop_time(integrator::CVODEIntegrator, tstop)
    CVodeSetStopTime(integrator.mem, tstop)
end
function set_stop_time(integrator::ARKODEIntegrator, tstop)
    ARKStepSetStopTime(integrator.mem, tstop)
end
function set_stop_time(integrator::IDAIntegrator, tstop)
    IDASetStopTime(integrator.mem, tstop)
end

function get_iters!(integrator::CVODEIntegrator, iters)
    CVodeGetNumSteps(integrator.mem, iters)
end
function get_iters!(integrator::ARKODEIntegrator, iters)
    ARKStepGetNumSteps(integrator.mem, iters)
end
function get_iters!(integrator::IDAIntegrator, iters)
    IDAGetNumSteps(integrator.mem, iters)
end

function DiffEqBase.solve!(integrator::AbstractSundialsIntegrator)
    uType = eltype(integrator.sol.u)
    iters = Ref(Clong(-1))
    while !isempty(integrator.opts.tstops)
        # Sundials can have floating point issues approaching a tstop if
        # there is a modifying event each
        # The call to first is an overload of Base.first implemented in DataStructures
        while integrator.tdir * (integrator.t - first(integrator.opts.tstops)) < -1e6eps()
            tstop = first(integrator.opts.tstops)
            set_stop_time(integrator, tstop)
            integrator.tprev = integrator.t
            if !(typeof(integrator.opts.callback.continuous_callbacks) <: Tuple{})
                integrator.uprev .= integrator.u
            end
            integrator.userfun.p = integrator.p
            solver_step(integrator, tstop)
            integrator.t = first(integrator.tout)
            integrator.flag < 0 && break
            handle_callbacks!(integrator)
            integrator.flag < 0 && break
            if isempty(integrator.opts.tstops)
                break
            end
            get_iters!(integrator, iters)
            if iters[] + 1 > integrator.opts.maxiters
                integrator.flag = -1 # MaxIters: -1
                break
            end
        end
        integrator.flag < 0 && break
        handle_tstop!(integrator)
    end

    if integrator.opts.save_end &&
       (isempty(integrator.sol.t) || integrator.sol.t[end] != integrator.t)
        save_value!(integrator.sol.u, integrator.u, uType, integrator.sizeu, integrator.opts.save_idxs)
        push!(integrator.sol.t, integrator.t)
        if integrator.opts.dense
            integrator(integrator.u, integrator.t, Val{1})
            save_value!(integrator.sol.interp.du, integrator.u, uType, integrator.sizeu, integrator.opts.save_idxs)
        end
    end

    if integrator.opts.progress
        Logging.@logmsg(
            -1,
            integrator.opts.progress_name,
            _id = :Sundials,
            message = integrator.opts.progress_message(
                integrator.dt,
                integrator.u,
                integrator.p,
                integrator.t,
            ),
            progress = "done"
        )
    end

    fill_destats!(integrator)
    empty!(integrator.mem)
    integrator.A != nothing && empty!(integrator.A)
    integrator.LS != nothing && empty!(integrator.LS)

    if DiffEqBase.has_analytic(integrator.sol.prob.f)
        DiffEqBase.calculate_solution_errors!(
            integrator.sol;
            timeseries_errors = integrator.opts.timeseries_errors,
            dense_errors = integrator.opts.dense_errors,
        )
    end

    if integrator.sol.retcode === :Default
        integrator.sol = DiffEqBase.solution_new_retcode(
            integrator.sol,
            interpret_sundials_retcode(integrator.flag),
        )
    end

    return integrator.sol
end

function handle_tstop!(integrator::AbstractSundialsIntegrator)
    tstops = integrator.opts.tstops
    if !isempty(tstops)
        if integrator.tdir * (integrator.t - first(integrator.opts.tstops)) > -1e6eps()
            pop!(tstops)
            t = integrator.t
            integrator.just_hit_tstop = true
        end
    end
end

function fill_destats!(integrator::AbstractSundialsIntegrator) end

function fill_destats!(integrator::CVODEIntegrator)
    destats = integrator.sol.destats
    mem = integrator.mem
    tmp = Ref(Clong(-1))
    CVodeGetNumRhsEvals(mem, tmp)
    destats.nf = tmp[]
    CVodeGetNumLinSolvSetups(mem, tmp)
    destats.nw = tmp[]
    CVodeGetNumErrTestFails(mem, tmp)
    destats.nreject = tmp[]
    CVodeGetNumSteps(mem, tmp)
    destats.naccept = tmp[] - destats.nreject
    CVodeGetNumNonlinSolvIters(mem, tmp)
    destats.nnonliniter = tmp[]
    CVodeGetNumNonlinSolvConvFails(mem, tmp)
    destats.nnonlinconvfail = tmp[]
    if method_choice(integrator.alg) == :Newton
        CVodeGetNumJacEvals(mem, tmp)
        destats.njacs = tmp[]
    end
end

function fill_destats!(integrator::ARKODEIntegrator)
    destats = integrator.sol.destats
    mem = integrator.mem
    tmp = Ref(Clong(-1))
    tmp2 = Ref(Clong(-1))
    ARKStepGetNumRhsEvals(mem, tmp, tmp2)
    destats.nf = tmp[]
    destats.nf2 = tmp2[]
    integrator.alg.stiffness !== Explicit() && ARKStepGetNumLinSolvSetups(mem, tmp)
    destats.nw = tmp[]
    ARKStepGetNumErrTestFails(mem, tmp)
    destats.nreject = tmp[]
    ARKStepGetNumSteps(mem, tmp)
    destats.naccept = tmp[] - destats.nreject
    integrator.alg.stiffness !== Explicit() && ARKStepGetNumNonlinSolvIters(mem, tmp)
    destats.nnonliniter = tmp[]
    integrator.alg.stiffness !== Explicit() && ARKStepGetNumNonlinSolvConvFails(mem, tmp)
    destats.nnonlinconvfail = tmp[]
    if integrator.alg.stiffness !== Explicit() && method_choice(integrator.alg) == :Newton
        ARKStepGetNumJacEvals(mem, tmp)
        destats.njacs = tmp[]
    end
end

function fill_destats!(integrator::IDAIntegrator)
    destats = integrator.sol.destats
    mem = integrator.mem
    tmp = Ref(Clong(-1))
    IDAGetNumResEvals(mem, tmp)
    destats.nf = tmp[]
    IDAGetNumLinSolvSetups(mem, tmp)
    destats.nw = tmp[]
    IDAGetNumErrTestFails(mem, tmp)
    destats.nreject = tmp[]
    IDAGetNumSteps(mem, tmp)
    destats.naccept = tmp[] - destats.nreject
    IDAGetNumNonlinSolvIters(mem, tmp)
    destats.nnonliniter = tmp[]
    IDAGetNumNonlinSolvConvFails(mem, tmp)
    destats.nnonlinconvfail = tmp[]
    if method_choice(integrator.alg) == :Newton
        IDAGetNumJacEvals(mem, tmp)
        destats.njacs = tmp[]
    end
end

function initialize_callbacks!(integrator, initialize_save = true)
    t = integrator.t
    u = integrator.u
    callbacks = integrator.opts.callback
    integrator.u_modified = true

    u_modified = initialize!(callbacks, u, t, integrator)

    # if the user modifies u, we need to fix current values
    if u_modified

        handle_callback_modifiers!(integrator)

        if initialize_save && (
            any((c) -> c.save_positions[2], callbacks.discrete_callbacks) ||
            any((c) -> c.save_positions[2], callbacks.continuous_callbacks)
        )
            savevalues!(integrator, true)
        end
    end

    # reset this as it is now handled so the integrators should proceed as normal
    integrator.u_modified = false
end


text:
using Yao, YaoArrayRegister, YaoBlocks, BenchmarkTools
using PkgBenchmark

function bench(block)
    N = nqubits(block)
    r = rand_state(N)
    return @benchmarkable apply!($r, $block)
end

const SUITE = BenchmarkGroup()

SUITE["gates"] = BenchmarkGroup()
SUITE["random circuits"] = BenchmarkGroup()

for G in [:X, :Y, :Z, :S, :T, :Sdag, :Tdag]
    GT = Expr(:(.), :ConstGate, QuoteNode(Symbol(G, :Gate)))
    GN = Expr(:(.), :ConstGate, QuoteNode(G))

    SUITE["gates"][string(G)] = BenchmarkGroup()
    for n in 1:4:25
        SUITE["gates"][string(G)][n] = @eval bench(repeat($n, $GN))
    end
end

ext: jl


text:
using Flux: Conv, SkipConnection, relu
"""
    ResNetHP

Hyperparameters for the convolutional resnet architecture.

| Parameter                 | Type                | Default   |
|:--------------------------|:--------------------|:----------|
| `num_blocks`              | `Int`               |  -        |
| `num_filters`             | `Int`               |  -        |
| `conv_kernel_size`        | `Tuple{Int, Int}`   |  -        |
| `num_policy_head_filters` | `Int`               | `2`       |
| `num_value_head_filters`  | `Int`               | `1`       |
| `batch_norm_momentum`     | `Float32`           | `0.6f0`   |

The trunk of the two-head network consists of `num_blocks` consecutive blocks.
Each block features two convolutional layers with `num_filters` filters and
with kernel size `conv_kernel_size`. Note that both kernel dimensions must be
odd.

During training, the network is evaluated in training mode on the whole
dataset to compute the loss before it is switched to test model, using
big batches. Therefore, it makes sense to use a high batch norm momentum
(put a lot of weight on the latest measurement).

# AlphaGo Zero Parameters

The network in the original paper from Deepmind features 20 blocks with 256
filters per convolutional layer.
"""
@kwdef struct ResNetHP_
  indim :: Tuple{Int, Int, Int}
  outdim_head1 :: Int
  num_blocks :: Int
  num_filters :: Int
  conv_kernel_size :: Tuple{Int, Int}
  num_head1_filters :: Int = 2 # policy
  num_head2_filters :: Int = 1 # value
  batch_norm_momentum :: Float32 = 0.6f0
end

"""
    ResNet <: TwoHeadNetwork

The convolutional residual network architecture that is used
in the original AlphaGo Zero paper.
"""
mutable struct ResNet_
  gspec
  hyper
  common
  head1 #policy/state (vector/tensor like)
  head2 #value/reward (scalar like)
end

function ResNetBlock(size, n, bnmom)
  pad = size .÷ 2
  layers = Chain(
    Conv(size, n=>n, pad=pad, bias=false),
    BatchNorm(n, relu, momentum=bnmom),
    Conv(size, n=>n, pad=pad, bias=false),
    BatchNorm(n, momentum=bnmom))
  return Chain(
    SkipConnection(layers, +),
    x -> relu.(x))
    # relu)
end

function ResNet_(gspec::AbstractGameSpec, hyper::ResNetHP_)
  indim = hyper.indim
  outdim = hyper.outdim_head1
  ksize = hyper.conv_kernel_size
  @assert all(ksize .% 2 .== 1)
  pad = ksize .÷ 2
  nf = hyper.num_filters
  npf = hyper.num_head1_filters
  nvf = hyper.num_head2_filters
  bnmom = hyper.batch_norm_momentum
  common = Chain(
    Conv(ksize, indim[3]=>nf, pad=pad, bias=false),
    BatchNorm(nf, relu, momentum=bnmom),
    [ResNetBlock(ksize, nf, bnmom) for i in 1:hyper.num_blocks]...)
  phead = Chain(
    Conv((1, 1), nf=>npf, bias=false),
    BatchNorm(npf, relu, momentum=bnmom),
    flatten,
    Dense(indim[1] * indim[2] * npf, outdim),
    softmax)
  vhead = Chain(
    Conv((1, 1), nf=>nvf, bias=false),
    BatchNorm(nvf, relu, momentum=bnmom),
    flatten,
    Dense(indim[1] * indim[2] * nvf, nf, relu),
    Dense(nf, 1, tanh))
  ResNet_(gspec, hyper, common, phead, vhead)
end

### Representation ###

@kwdef struct RepresentationResnetHP
  num_blocks :: Int
  num_filters :: Int
  conv_kernel_size :: Tuple{Int, Int}
  batch_norm_momentum :: Float32 = 0.6f0
end

mutable struct RepresentationResnetNetwork <: AbstractRepresentation
  gspec
  hyper
  common
end

function RepresentationNetwork(gspec::AbstractGameSpec, hyper::RepresentationResnetHP)
  indim = GI.state_dim(gspec)
  resnethyper = ResNetHP_(
    indim,
    1, # outdim not used
    hyper.num_blocks,
    hyper.num_filters,
    hyper.conv_kernel_size,
    1, # num_head1_filters not used
    1, # num_head2_filters not used
    hyper.batch_norm_momentum
  )
  resnet = ResNet_(gspec, resnethyper)
  common = Chain(resnet.common,
    Conv((1,1), hyper.num_filters=>indim[3], relu)
  )
  return RepresentationResnetNetwork(gspec, hyper, common)
end

### Dynamics ###

@kwdef struct DynamicsResnetHP
  num_blocks :: Int
  num_filters :: Int
  conv_kernel_size :: Tuple{Int, Int}
  num_reward_head2_filters :: Int = 1 # head2
  batch_norm_momentum :: Float32 = 0.6f0
end

mutable struct DynamicsResnetNetwork <: AbstractDynamics
  gspec
  hyper
  common
  statehead
  rewardhead
end

# Dynamics Network with identity as state head
function DynamicsNetwork(gspec::AbstractGameSpec, hyper::DynamicsResnetHP)
  state_dim = GI.state_dim(gspec)
  # indim = (state_dim[1], state_dim[2], hyper.num_filters+1)
  indim = (state_dim[1], state_dim[2], state_dim[3]+1)
  resnethyper = ResNetHP_(
    indim, # indim not used
    1, # outdim_head1 not used
    hyper.num_blocks,
    hyper.num_filters,
    hyper.conv_kernel_size,
    1, # num_head1_filters not used
    hyper.num_reward_head2_filters,
    hyper.batch_norm_momentum
  )
  resnet = ResNet_(gspec, resnethyper)
  small_state = Conv((1,1), hyper.num_filters=>state_dim[3], relu) # a lot quicker hash
  return DynamicsResnetNetwork(gspec, hyper,
    resnet.common,
    small_state, # state head as identity
    resnet.head2)
end

@kwdef struct PredictionResnetHP
  num_blocks :: Int
  num_filters :: Int
  conv_kernel_size :: Tuple{Int, Int}
  num_policy_head1_filters :: Int = 2 # head1
  num_value_head2_filters :: Int = 1 # head2
  batch_norm_momentum :: Float32 = 0.6f0
end

mutable struct PredictionResnetNetwork <: AbstractPrediction
  gspec
  hyper
  common
  policyhead
  valuehead
end

function PredictionNetwork(gspec::AbstractGameSpec, hyper::PredictionResnetHP)
  indim = GI.state_dim(gspec)
  outdim_head1 = GI.num_actions(gspec)
  resnethyper = ResNetHP_(
    indim, # indim not used
    outdim_head1,
    hyper.num_blocks,
    hyper.num_filters,
    hyper.conv_kernel_size,
    hyper.num_policy_head1_filters,
    hyper.num_value_head2_filters,
    hyper.batch_norm_momentum
  )
  resnet = ResNet_(gspec, resnethyper)
  return PredictionResnetNetwork(gspec, hyper,
    resnet.common, # leaving only residual layers
    resnet.head1,
    resnet.head2)
end


text:
### A Pluto.jl notebook ###
# v0.17.3

using Markdown
using InteractiveUtils

# ╔═╡ 78bdae60-f576-4359-8826-fc669cb65e94
# hideall

using Genie

# ╔═╡ 8528f242-491d-11ec-32a0-8d13960b85b1
md"""
# Using Genie in an interactive environment (Jupyter/IJulia, REPL, etc)

Genie can be used for ad-hoc exploratory programming, to quickly whip up a web server
and expose your Julia functions.

Once you have `Genie` into scope, you can define a new `route`.
A `route` maps a URL to a function.
"""

# ╔═╡ 1f1ed8b2-3d42-4ae5-9731-db80f6ea7c21
md"""

```julia

julia> using Pkg
julia> Pkg.add("Genie")

julia> using Genie
```
"""

# ╔═╡ de738176-e311-48b0-9cd9-72195e8e4f4e
md"""

```julia
julia> route("/") do
         "Hi there!"
       end
```
"""

# ╔═╡ c230793f-f0e4-4f8e-8786-62ff8b770fd4
# hideall

route("/") do
	"Hi there!"
end

# ╔═╡ 3a7f699e-1f1b-458b-9714-bcede1cd2c3b
md"""

```julia
julia> Genie.startup()
```
"""

# ╔═╡ 1dcb137a-6527-4b0f-831d-0cd190f2e09a
# hideall

Genie.startup(async=true);

# ╔═╡ 23aab64c-cf08-4c26-a2bd-b0e6f9609f66
md"""
Finally, now navigate to <http://localhost:8000> – you should see the message "Hi there!".

We can define more complex URIs which can also map to previously defined functions:
"""

# ╔═╡ 36fc192e-69be-427e-9262-51d843deab1d
md"""
```julia
julia> function hello_world()
         "Hello World! NEW"
       end
```
"""

# ╔═╡ 71e60edf-c682-49a8-9a5b-72004b80119d
# hideall

function hello_world()
         "Hello World! NEW"
end

# ╔═╡ 25778267-a2b6-459a-a0e0-225d0dd7c360
md"""

```julia
julia> route("/hello/world", hello_world)
```
"""

# ╔═╡ 1d6a7afd-6507-405e-a703-54e01b63fe6d
# hideall

route("/hello/world", hello_world)

# ╔═╡ 47689ca7-2ce3-444c-94e1-84284c548f5c
md"""
Obviously, the functions can be defined anywhere (in any other module) as long as they are accessible in the current scope.

You can now visit <http://localhost:8000/hello/world> in the browser.

Of course we can access GET params:
"""

# ╔═╡ 095f83a8-24cc-402d-a1c6-1a56010606fc
md"""

```julia
julia> route("/echo/:message") do
         params(:message)
       end
```
"""

# ╔═╡ 190216f8-fda7-49a8-bf8f-a472d60aee89
# hideall

route("/echo/:message") do
         params(:message)
end

# ╔═╡ 733a1371-9fb9-4221-bab3-51657fe597cd
md"""
Accessing <http://localhost:8000/echo/ciao> should echo "ciao".

And we can even match by types:
"""

# ╔═╡ ed8c5f05-59ac-4c46-b403-e20cd499b2ec
md"""

```julia

julia> route("/sum/:x::Int/:y::Int") do
         params(:x) + params(:y)
       end
```
"""

# ╔═╡ 7fb97a9e-f92c-49a9-a2cf-6e0b68cdadd9
# hideall

route("/sum/:x::Int/:y::Int") do
         params(:x) + params(:y)
end

# ╔═╡ 043b9ca3-fa9c-4b88-a2fb-9f40f22890f3
md"""

By default, GET params are extracted as `SubString` (more exactly, `SubString{String}`).
If type constraints are added, Genie will attempt to convert the `SubString` to the indicated type.

For the above to work, we also need to tell Genie how to perform the conversion:
"""

# ╔═╡ 728d9da4-e24f-4a0e-9253-4a22bb885fca
md"""
```julia
julia> import Base.convert

julia> convert(::Type{Int}, s::AbstractString) = parse(Int, s)
```"""

# ╔═╡ d5726bee-c104-4181-93f8-c8a6c0be8c2b
# hideall

convert(::Type{Int}, s::AbstractString) = parse(Int, s)

# ╔═╡ d3c7b6f6-cf3a-44f5-bd1a-bd66729208a8
md"""

Now if we access <http://localhost:8000/sum/2/3> we should see `5`

## Handling query string params

Query string params, which look like `...?foo=bar&baz=2` are automatically unpacked by Genie and placed into the `params` collection. For example:
"""

# ╔═╡ 1136b778-e5eb-4e12-b5ba-b4cb8a48d2ad
md"""

```julia
julia> route("/sum/:x::Int/:y::Int") do
         params(:x) + params(:y) + parse(Int, get(params, :initial_value, "0"))
       end
```
"""

# ╔═╡ e6964ee4-02b0-4161-b85a-44adabab49be
# hideall

route("/sum/:x::Int/:y::Int") do
         params(:x) + params(:y) + parse(Int, get(params, :initial_value, "0"))
end

# ╔═╡ 63177c34-6987-480a-9e8e-e8328eb5d195
md"""
Accessing <http://localhost:8000/sum/2/3?initial_value=10> will now output `15`.
"""

# ╔═╡ 00000000-0000-0000-0000-000000000001
PLUTO_PROJECT_TOML_CONTENTS = """
[deps]
Genie = "c43c736e-a2d1-11e8-161f-af95117fbd1e"

[compat]
Genie = "~4.2.0"
"""

# ╔═╡ 00000000-0000-0000-0000-000000000002
PLUTO_MANIFEST_TOML_CONTENTS = """
# This file is machine-generated - editing it directly is not advised

[[ArgParse]]
deps = ["Logging", "TextWrap"]
git-tree-sha1 = "3102bce13da501c9104df33549f511cd25264d7d"
uuid = "c7e460c6-2fb9-53a9-8c5b-16f535851c63"
version = "1.1.4"

[[ArgTools]]
uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"

[[Artifacts]]
uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"

[[Base64]]
uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

[[CSTParser]]
deps = ["Tokenize"]
git-tree-sha1 = "f9a6389348207faf5e5c62cbc7e89d19688d338a"
uuid = "00ebfdb7-1f24-5e51-bd34-a7502290713f"
version = "3.3.0"

[[CodeTracking]]
deps = ["InteractiveUtils", "UUIDs"]
git-tree-sha1 = "9aa8a5ebb6b5bf469a7e0e2b5202cf6f8c291104"
uuid = "da1fd8a2-8d9e-5ec2-8556-3022fb5608a2"
version = "1.0.6"

[[CommonMark]]
deps = ["Crayons", "JSON", "URIs"]
git-tree-sha1 = "393ac9df4eb085c2ab12005fc496dae2e1da344e"
uuid = "a80b9123-70ca-4bc0-993e-6e3bcb318db6"
version = "0.8.3"

[[Compat]]
deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
git-tree-sha1 = "dce3e3fea680869eaa0b774b2e8343e9ff442313"
uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
version = "3.40.0"

[[Crayons]]
git-tree-sha1 = "3f71217b538d7aaee0b69ab47d9b7724ca8afa0d"
uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
version = "4.0.4"

[[DataStructures]]
deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
git-tree-sha1 = "7d9d316f04214f7efdbb6398d545446e246eff02"
uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
version = "0.18.10"

[[Dates]]
deps = ["Printf"]
uuid = "ade2ca70-3891-5945-98fb-dc099432e06a"

[[DelimitedFiles]]
deps = ["Mmap"]
uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"

[[Distributed]]
deps = ["Random", "Serialization", "Sockets"]
uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"

[[Downloads]]
deps = ["ArgTools", "LibCURL", "NetworkOptions"]
uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6"

[[EzXML]]
deps = ["Printf", "XML2_jll"]
git-tree-sha1 = "0fa3b52a04a4e210aeb1626def9c90df3ae65268"
uuid = "8f5d6c58-4d21-5cfd-889c-e3ad7ee6a615"
version = "1.1.0"

[[FilePathsBase]]
deps = ["Dates", "Mmap", "Printf", "Test", "UUIDs"]
git-tree-sha1 = "618835ab81e4a40acf215c98768978d82abc5d97"
uuid = "48062228-2e41-5def-b9a4-89aafe57970f"
version = "0.9.16"

[[FileWatching]]
uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"

[[GMP_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "781609d7-10c4-51f6-84f2-b8444358ff6d"

[[Genie]]
deps = ["ArgParse", "Dates", "Distributed", "EzXML", "FilePathsBase", "HTTP", "HttpCommon", "Inflector", "JSON3", "JuliaFormatter", "Logging", "Markdown", "MbedTLS", "Millboard", "Nettle", "OrderedCollections", "Pkg", "REPL", "Reexport", "Revise", "SHA", "Serialization", "Sockets", "UUIDs", "Unicode", "VersionCheck", "YAML"]
git-tree-sha1 = "9ec6a1b969aaf8aa6d6b9513055fba3effa7671e"
uuid = "c43c736e-a2d1-11e8-161f-af95117fbd1e"
version = "4.2.0"

[[HTTP]]
deps = ["Base64", "Dates", "IniFile", "Logging", "MbedTLS", "NetworkOptions", "Sockets", "URIs"]
git-tree-sha1 = "0fa77022fe4b511826b39c894c90daf5fce3334a"
uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3"
version = "0.9.17"

[[HttpCommon]]
deps = ["Dates", "Nullables", "Test", "URIParser"]
git-tree-sha1 = "46313284237aa6ca67a6bce6d6fbd323d19cff59"
uuid = "77172c1b-203f-54ac-aa54-3f1198fe9f90"
version = "0.5.0"

[[Inflector]]
deps = ["Unicode"]
git-tree-sha1 = "8555b54ddf27806b070ce1d1cf623e1feb13750c"
uuid = "6d011eab-0732-4556-8808-e463c76bf3b6"
version = "1.0.1"

[[IniFile]]
deps = ["Test"]
git-tree-sha1 = "098e4d2c533924c921f9f9847274f2ad89e018b8"
uuid = "83e8ac13-25f8-5344-8a64-a9f2b223428f"
version = "0.5.0"

[[InteractiveUtils]]
deps = ["Markdown"]
uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"

[[JLLWrappers]]
deps = ["Preferences"]
git-tree-sha1 = "642a199af8b68253517b80bd3bfd17eb4e84df6e"
uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
version = "1.3.0"

[[JSON]]
deps = ["Dates", "Mmap", "Parsers", "Unicode"]
git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37"
uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
version = "0.21.2"

[[JSON3]]
deps = ["Dates", "Mmap", "Parsers", "StructTypes", "UUIDs"]
git-tree-sha1 = "7d58534ffb62cd947950b3aa9b993e63307a6125"
uuid = "0f8b85d8-7281-11e9-16c2-39a750bddbf1"
version = "1.9.2"

[[JuliaFormatter]]
deps = ["CSTParser", "CommonMark", "DataStructures", "Pkg", "Tokenize"]
git-tree-sha1 = "10c95cebcfa37c1f510a726c90886db4745e1238"
uuid = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
version = "0.15.11"

[[JuliaInterpreter]]
deps = ["CodeTracking", "InteractiveUtils", "Random", "UUIDs"]
git-tree-sha1 = "e273807f38074f033d94207a201e6e827d8417db"
uuid = "aa1ae85d-cabe-5617-a682-6adf51b2e16a"
version = "0.8.21"

[[LibCURL]]
deps = ["LibCURL_jll", "MozillaCACerts_jll"]
uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21"

[[LibCURL_jll]]
deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"]
uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0"

[[LibGit2]]
deps = ["Base64", "NetworkOptions", "Printf", "SHA"]
uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"

[[LibSSH2_jll]]
deps = ["Artifacts", "Libdl", "MbedTLS_jll"]
uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8"

[[Libdl]]
uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb"

[[Libiconv_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "42b62845d70a619f063a7da093d995ec8e15e778"
uuid = "94ce4f54-9a6c-5748-9c1c-f9c7231a4531"
version = "1.16.1+1"

[[LinearAlgebra]]
deps = ["Libdl"]
uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

[[Logging]]
uuid = "56ddb016-857b-54e1-b83d-db4d58db5568"

[[LoweredCodeUtils]]
deps = ["JuliaInterpreter"]
git-tree-sha1 = "491a883c4fef1103077a7f648961adbf9c8dd933"
uuid = "6f1432cf-f94c-5a45-995e-cdbf5db27b0b"
version = "2.1.2"

[[Markdown]]
deps = ["Base64"]
uuid = "d6f4376e-aef5-505a-96c1-9c027394607a"

[[MbedTLS]]
deps = ["Dates", "MbedTLS_jll", "Random", "Sockets"]
git-tree-sha1 = "1c38e51c3d08ef2278062ebceade0e46cefc96fe"
uuid = "739be429-bea8-5141-9913-cc70e7f3736d"
version = "1.0.3"

[[MbedTLS_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1"

[[Millboard]]
git-tree-sha1 = "ea6a5b7e56e76d8051023faaa11d91d1d881dac3"
uuid = "39ec1447-df44-5f4c-beaa-866f30b4d3b2"
version = "0.2.5"

[[Mmap]]
uuid = "a63ad114-7e13-5084-954f-fe012c677804"

[[MozillaCACerts_jll]]
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"

[[Nettle]]
deps = ["Libdl", "Nettle_jll"]
git-tree-sha1 = "a68340b9edfd98d0ed96aee8137cb716ea3b6dea"
uuid = "49dea1ee-f6fa-5aa6-9a11-8816cee7d4b9"
version = "0.5.1"

[[Nettle_jll]]
deps = ["Artifacts", "GMP_jll", "JLLWrappers", "Libdl", "Pkg"]
git-tree-sha1 = "eca63e3847dad608cfa6a3329b95ef674c7160b4"
uuid = "4c82536e-c426-54e4-b420-14f461c4ed8b"
version = "3.7.2+0"

[[NetworkOptions]]
uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"

[[Nullables]]
git-tree-sha1 = "8f87854cc8f3685a60689d8edecaa29d2251979b"
uuid = "4d1e1d77-625e-5b40-9113-a560ec7a8ecd"
version = "1.0.0"

[[OrderedCollections]]
git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
version = "1.4.1"

[[Parsers]]
deps = ["Dates"]
git-tree-sha1 = "ae4bbcadb2906ccc085cf52ac286dc1377dceccc"
uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0"
version = "2.1.2"

[[Pkg]]
deps = ["Artifacts", "Dates", "Downloads", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"]
uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"

[[Preferences]]
deps = ["TOML"]
git-tree-sha1 = "00cfd92944ca9c760982747e9a1d0d5d86ab1e5a"
uuid = "21216c6a-2e73-6563-6e65-726566657250"
version = "1.2.2"

[[Printf]]
deps = ["Unicode"]
uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"

[[ProgressMeter]]
deps = ["Distributed", "Printf"]
git-tree-sha1 = "afadeba63d90ff223a6a48d2009434ecee2ec9e8"
uuid = "92933f4c-e287-5a05-a399-4b506db050ca"
version = "1.7.1"

[[REPL]]
deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"]
uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"

[[Random]]
deps = ["Serialization"]
uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

[[Reexport]]
git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b"
uuid = "189a3867-3050-52da-a836-e630ba90ab69"
version = "1.2.2"

[[Requires]]
deps = ["UUIDs"]
git-tree-sha1 = "4036a3bd08ac7e968e27c203d45f5fff15020621"
uuid = "ae029012-a4dd-5104-9daa-d747884805df"
version = "1.1.3"

[[Revise]]
deps = ["CodeTracking", "Distributed", "FileWatching", "JuliaInterpreter", "LibGit2", "LoweredCodeUtils", "OrderedCollections", "Pkg", "REPL", "Requires", "UUIDs", "Unicode"]
git-tree-sha1 = "41deb3df28ecf75307b6e492a738821b031f8425"
uuid = "295af30f-e4ad-537b-8983-00126c2a3abe"
version = "3.1.20"

[[SHA]]
uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce"

[[Scratch]]
deps = ["Dates"]
git-tree-sha1 = "0b4b7f1393cff97c33891da2a0bf69c6ed241fda"
uuid = "6c6a2e73-6563-6170-7368-637461726353"
version = "1.1.0"

[[Serialization]]
uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b"

[[SharedArrays]]
deps = ["Distributed", "Mmap", "Random", "Serialization"]
uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383"

[[Sockets]]
uuid = "6462fe0b-24de-5631-8697-dd941f90decc"

[[SparseArrays]]
deps = ["LinearAlgebra", "Random"]
uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

[[Statistics]]
deps = ["LinearAlgebra", "SparseArrays"]
uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"

[[StringEncodings]]
deps = ["Libiconv_jll"]
git-tree-sha1 = "50ccd5ddb00d19392577902f0079267a72c5ab04"
uuid = "69024149-9ee7-55f6-a4c4-859efe599b68"
version = "0.3.5"

[[StructTypes]]
deps = ["Dates", "UUIDs"]
git-tree-sha1 = "d24a825a95a6d98c385001212dc9020d609f2d4f"
uuid = "856f2bd8-1eba-4b0a-8007-ebc267875bd4"
version = "1.8.1"

[[TOML]]
deps = ["Dates"]
uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76"

[[Tar]]
deps = ["ArgTools", "SHA"]
uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e"

[[Test]]
deps = ["InteractiveUtils", "Logging", "Random", "Serialization"]
uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

[[TextWrap]]
git-tree-sha1 = "9250ef9b01b66667380cf3275b3f7488d0e25faf"
uuid = "b718987f-49a8-5099-9789-dcd902bef87d"
version = "1.0.1"

[[Tokenize]]
git-tree-sha1 = "0952c9cee34988092d73a5708780b3917166a0dd"
uuid = "0796e94c-ce3b-5d07-9a54-7f471281c624"
version = "0.5.21"

[[URIParser]]
deps = ["Unicode"]
git-tree-sha1 = "53a9f49546b8d2dd2e688d216421d050c9a31d0d"
uuid = "30578b45-9adc-5946-b283-645ec420af67"
version = "0.4.1"

[[URIs]]
git-tree-sha1 = "97bbe755a53fe859669cd907f2d96aee8d2c1355"
uuid = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4"
version = "1.3.0"

[[UUIDs]]
deps = ["Random", "SHA"]
uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"

[[Unicode]]
uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"

[[UrlDownload]]
deps = ["HTTP", "ProgressMeter"]
git-tree-sha1 = "05f86730c7a53c9da603bd506a4fc9ad0851171c"
uuid = "856ac37a-3032-4c1c-9122-f86d88358c8b"
version = "1.0.0"

[[VersionCheck]]
deps = ["Dates", "JSON3", "Logging", "Pkg", "Random", "Scratch", "UrlDownload"]
git-tree-sha1 = "89ef2431dd59344ebaf052d0737205854ded0c62"
uuid = "a637dc6b-bca1-447e-a4fa-35264c9d0580"
version = "0.2.0"

[[XML2_jll]]
deps = ["Artifacts", "JLLWrappers", "Libdl", "Libiconv_jll", "Pkg", "Zlib_jll"]
git-tree-sha1 = "1acf5bdf07aa0907e0a37d3718bb88d4b687b74a"
uuid = "02c8fc9c-b97f-50b9-bbe4-9be30ff0a78a"
version = "2.9.12+0"

[[YAML]]
deps = ["Base64", "Dates", "Printf", "StringEncodings"]
git-tree-sha1 = "3c6e8b9f5cdaaa21340f841653942e1a6b6561e5"
uuid = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
version = "0.4.7"

[[Zlib_jll]]
deps = ["Libdl"]
uuid = "83775a58-1f1d-513f-b197-d71354ab007a"

[[nghttp2_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d"

[[p7zip_jll]]
deps = ["Artifacts", "Libdl"]
uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0"
"""

# ╔═╡ Cell order:
# ╟─8528f242-491d-11ec-32a0-8d13960b85b1
# ╟─1f1ed8b2-3d42-4ae5-9731-db80f6ea7c21
# ╠═78bdae60-f576-4359-8826-fc669cb65e94
# ╟─de738176-e311-48b0-9cd9-72195e8e4f4e
# ╠═c230793f-f0e4-4f8e-8786-62ff8b770fd4
# ╟─3a7f699e-1f1b-458b-9714-bcede1cd2c3b
# ╠═1dcb137a-6527-4b0f-831d-0cd190f2e09a
# ╟─23aab64c-cf08-4c26-a2bd-b0e6f9609f66
# ╟─36fc192e-69be-427e-9262-51d843deab1d
# ╠═71e60edf-c682-49a8-9a5b-72004b80119d
# ╟─25778267-a2b6-459a-a0e0-225d0dd7c360
# ╠═1d6a7afd-6507-405e-a703-54e01b63fe6d
# ╟─47689ca7-2ce3-444c-94e1-84284c548f5c
# ╟─095f83a8-24cc-402d-a1c6-1a56010606fc
# ╠═190216f8-fda7-49a8-bf8f-a472d60aee89
# ╟─733a1371-9fb9-4221-bab3-51657fe597cd
# ╟─ed8c5f05-59ac-4c46-b403-e20cd499b2ec
# ╠═7fb97a9e-f92c-49a9-a2cf-6e0b68cdadd9
# ╟─043b9ca3-fa9c-4b88-a2fb-9f40f22890f3
# ╟─728d9da4-e24f-4a0e-9253-4a22bb885fca
# ╠═d5726bee-c104-4181-93f8-c8a6c0be8c2b
# ╟─d3c7b6f6-cf3a-44f5-bd1a-bd66729208a8
# ╟─1136b778-e5eb-4e12-b5ba-b4cb8a48d2ad
# ╠═e6964ee4-02b0-4161-b85a-44adabab49be
# ╟─63177c34-6987-480a-9e8e-e8328eb5d195
# ╟─00000000-0000-0000-0000-000000000001
# ╟─00000000-0000-0000-0000-000000000002


text:
include("imports.jl")

##

@testset ExtendedTestSet "FastAI.jl" begin
    @testset ExtendedTestSet "datablock.jl" begin
        include("datablock.jl")
    end

    @testset ExtendedTestSet "fasterai.jl" begin
        include("fasterai.jl")
    end

    @testset ExtendedTestSet "encodings/" begin
        @testset ExtendedTestSet "projective.jl" begin
            include("encodings/projective.jl")
        end
        @testset ExtendedTestSet "imagepreprocessing.jl" begin
            include("encodings/imagepreprocessing.jl")
        end
        @testset ExtendedTestSet "keypointpreprocessing.jl" begin
            include("encodings/keypointpreprocessing.jl")
        end
        @testset ExtendedTestSet "many.jl" begin
            include("encodings/many.jl")
        end
        @testset ExtendedTestSet "tabularpreprocessing.jl" begin
            include("encodings/tabularpreprocessing.jl")
        end
    end

    @testset ExtendedTestSet "methods/" begin
        @testset ExtendedTestSet "imageclassification.jl" begin
            include("methods/imageclassification.jl")
        end
        @testset ExtendedTestSet "imagesegmentation.jl" begin
            include("methods/imagesegmentation.jl")
        end
        @testset ExtendedTestSet "imagekeypointregression.jl" begin
            include("methods/imagekeypointregression.jl")
        end
        @testset ExtendedTestSet "tabularclassification.jl" begin
            include("methods/tabularclassification.jl")
        end
        @testset ExtendedTestSet "tabularregression.jl" begin
            include("methods/tabularregression.jl")
        end
    end

    @testset ExtendedTestSet "datasets/" begin
        @testset ExtendedTestSet "transformations.jl" begin
            include("datasets/transformations.jl")
        end
        @testset ExtendedTestSet "containers.jl" begin
            include("datasets/containers.jl")
        end
        @testset ExtendedTestSet "recipes.jl" begin
            include("datasets/recipes.jl")
        end
        @testset ExtendedTestSet "registry.jl" begin
            include("datasets/registry.jl")
        end
    end

    @testset ExtendedTestSet "learner.jl" begin
        include("learner.jl")
    end

    @testset ExtendedTestSet "training/" begin
        @testset ExtendedTestSet "paramgroups.jl" begin
            include("training/paramgroups.jl")
        end
        @testset ExtendedTestSet "discriminativelrs.jl" begin
            include("training/discriminativelrs.jl")
        end
        @testset ExtendedTestSet "fitonecycle.jl" begin
            include("training/fitonecycle.jl")
        end
        @testset ExtendedTestSet "finetune.jl" begin
            include("training/finetune.jl")
        end
        @testset ExtendedTestSet "lrfind.jl" begin
            include("training/lrfind.jl")
        end
        # TODO: test learning rate finder
    end

    @testset ExtendedTestSet "models/" begin
        @testset ExtendedTestSet "tabularmodel.jl" begin
            include("models/tabularmodel.jl")
        end
    end
end


text:
using BenchmarkTools, Plots, Plots.PlotMeasures, LaTeXStrings
using BenchmarkTools: minimum, median

SUITE = BenchmarkGroup()
model = "PLATOONING"
cases = ["PLAD01-BND42", "PLAD01-BND42-discrete",
         "PLAD01-BND30", "PLAD01-BND30-discrete"]

SUITE[model] = BenchmarkGroup()

include("Platoon.jl")
validation = []

boxdirs = BoxDirections{Float64, Vector{Float64}}(10)
octdirs = CustomDirections([Vector(vi) for vi in OctDirections(10)])

# ----------------------------------------
#  PLAD01-BND42 (dense time)
# ----------------------------------------

prob_PLAD01_BND42 = platoon(; deterministic_switching=true)
sol_PLAD01_BND42 = solve(prob_PLAD01_BND42, alg=BOX(δ=0.01),
                         clustering_method=BoxClustering(1),
                         intersection_method=TemplateHullIntersection(boxdirs),
                         intersect_source_invariant=false,
                         tspan = (0.0 .. 20.0))
property = dmin_specification(sol_PLAD01_BND42, -42.0)
push!(validation, Int(property))
SUITE[model][cases[1]] = @benchmarkable solve($prob_PLAD01_BND42, alg=BOX(δ=0.01),
                               clustering_method=BoxClustering(1),
                               intersection_method=TemplateHullIntersection($boxdirs),
                               intersect_source_invariant=false,
                               tspan = (0.0 .. 20.0))

# ----------------------------------------
#  PLAD01-BND42 (discrete time)
# ----------------------------------------

prob_PLAD01_BND42 = platoon(; deterministic_switching=true)
sol_PLAD01_BND42_d = solve(prob_PLAD01_BND42,
                           alg=BOX(δ=0.1, approx_model=NoBloating()),
                           clustering_method=BoxClustering(1, [3,1,1,1,1,1,1,1,1,1]),
                           intersection_method=TemplateHullIntersection(boxdirs),
                           intersect_source_invariant=false,
                           tspan = (0.0 .. 20.0))
property = dmin_specification(sol_PLAD01_BND42_d, -42.0)
push!(validation, Int(property))
SUITE[model][cases[2]] = @benchmarkable solve($prob_PLAD01_BND42,
                               alg=BOX(δ=0.1, approx_model=NoBloating()),
                               clustering_method=BoxClustering(1, [3,1,1,1,1,1,1,1,1,1]),
                               intersection_method=TemplateHullIntersection($boxdirs),
                               intersect_source_invariant=false,
                               tspan = (0.0 .. 20.0))

# ----------------------------------------
#  PLAD01-BND30 (dense time)
# ----------------------------------------
prob_PLAD01_BND30 = platoon(; deterministic_switching=true)
imethod = TemplateHullIntersection(octdirs)
cmethod = LazyClustering(1)
alg = LGG09(δ=0.03, template=octdirs, approx_model=Forward(setops=octdirs))
sol_PLAD01_BND30 = solve(prob_PLAD01_BND30,
                         alg=alg,
                         clustering_method=cmethod,
                         intersection_method=imethod,
                         intersect_source_invariant=false,
                         tspan = (0.0 .. 20.0))
property = dmin_specification(sol_PLAD01_BND30, -30.0)
push!(validation, Int(property))
SUITE[model][cases[3]] = @benchmarkable solve($prob_PLAD01_BND30,
                            alg=$alg,
                            clustering_method=$cmethod,
                            intersection_method=$imethod,
                            intersect_source_invariant=false,
                            tspan = (0.0 .. 20.0))

# ----------------------------------------
#  PLAD01-BND30 (discrete time)
# ----------------------------------------
prob_PLAD01_BND30 = platoon(; deterministic_switching=true)
imethod = TemplateHullIntersection(octdirs)
cmethod = LazyClustering(1)
alg = LGG09(δ=0.1, template=octdirs, approx_model=NoBloating())
sol_PLAD01_BND30_d = solve(prob_PLAD01_BND30,
                           alg=alg,
                           clustering_method=cmethod,
                           intersection_method=imethod,
                           intersect_source_invariant=false,
                           tspan = (0.0 .. 20.0))
property = dmin_specification(sol_PLAD01_BND30_d, -30.0)
push!(validation, Int(property))
SUITE[model][cases[4]] = @benchmarkable solve($prob_PLAD01_BND30,
                            alg=$alg,
                            clustering_method=$cmethod,
                            intersection_method=$imethod,
                            intersect_source_invariant=false,
                            tspan = (0.0 .. 20.0))

# ==============================================================================
# Execute benchmarks and save benchmark results
# ==============================================================================

# tune parameters
tune!(SUITE)

# run the benchmarks
results = run(SUITE, verbose=true)

# return the sample with the smallest time value in each test
println("minimum time for each benchmark:\n", minimum(results))

# return the median for each test
println("median time for each benchmark:\n", median(results))

# export runtimes
runtimes = Dict()
for (i, c) in enumerate(cases)
    tm = median(results[model][c]).time * 1e-9
    runtimes[c] = round(tm, digits=4)
end

for (i, c) in enumerate(cases)
    print(io, "JuliaReach, $model, $c, $(validation[i]), $(runtimes[c])\n")
end

# ==============================================================================
# Plot
# ==============================================================================

fig = Plots.plot()

Plots.plot!(fig, sol_PLAD01_BND30, vars=(0, 1), linecolor=:blue, color=:blue, alpha=0.8, lw=1.0,
    tickfont=font(30, "Times"), guidefontsize=45,
    xlab=L"t",
    ylab=L"x_{1}",
    xtick=[0, 5, 10, 15, 20.], ytick=[-30, -20, -10, 0],
    xlims=(0., 20.), ylims=(-31, 7),
    bottom_margin=6mm, left_margin=2mm, right_margin=4mm, top_margin=3mm,
    size=(1000, 1000))

savefig(fig, joinpath(@__DIR__, "ARCH-COMP21-JuliaReach-Platoon-PLAD01-BND30"))

# release memory
sol_PLAD01_BND42 = nothing
sol_PLAD01_BND42_d = nothing
sol_PLAD01_BND30 = nothing
sol_PLAD01_BND30_d = nothing
GC.gc()


text:
using LinearAlgebra
function backtracking_line_search(objective,x,alpha,beta,step,arc,gradient)
    demanding_step = step
    while objective(x + demanding_step * arc) > objective(x) + alpha * demanding_step * gradient' * arc
        demanding_step = beta * demanding_step
    end
    return demanding_step
end

max_stars_repo_name: zhuhs-q/basic-optimal-method

max_issues_repo_name: Jacob-jinxy/basic-optimal-method

max_forks_repo_name: Jacob-jinxy/basic-optimal-method


text:
using Avalon
import Avalon.fit!
# import Yota.@nodiff
using Distributions
using MLDataUtils

# @nodiff eps(x) x

function binarycrossentropy(ŷ, y; agg=mean, ϵ=eps(eltype(ŷ)))
    agg(@.(-y * log(ŷ + ϵ) - (1 - y) * log(1 - ŷ + ϵ)))
end

function kldiv_normal(mu, log_s2; agg=mean)
    -1/2 * agg(@. 1 + log_s2 - mu ^ 2f0 - exp(log_s2))
end

mutable struct VAE
    encoder::Sequential
    enc2mu::Linear
    enc2s::Linear
    decoder::Sequential
    beta::Real
end

Avalon.ignored_field(m::VAE, ::Val{(:beta,)}) = true

VAE(encoder::Sequential, enc2mu::Any, enc2s::Any, decoder::Sequential; beta=1) =
    VAE(encoder,
        enc2mu,
        enc2s,
        decoder,
        beta)

VAE(encoder::Sequential, z_len::Int, decoder::Sequential; beta=1) =
    VAE(encoder,
        Linear(z_len => z_len),
        Linear(z_len => z_len),
        decoder,
        beta)

function Base.show(io::IO, m::VAE)
    print(io, "VAE()")
end

function (m::VAE)(x)
    h = m.encoder(x)
    mu = m.enc2mu(h)
    return mu
end

function encode(m::VAE, x)
    h = m.encoder(x)
    mu = m.enc2mu(h)
    log_s2 = m.enc2s(h)
    return mu, log_s2
end

decode(m::VAE, z) = m.decoder(z)

function loss_function(m::VAE, ϵ, x)
    mu, log_s2 = encode(m, x)
    z = mu .+ sqrt.(exp.(log_s2)) .* ϵ
    x_rec = decode(m, z)
    BCE = binarycrossentropy(x_rec, x; agg=sum)
    KLD = kldiv_normal(mu, log_s2; agg=sum)
    return BCE + m.beta * KLD
end

function fit!(m::VAE, X::AbstractMatrix{T};
              n_epochs=50, batch_size=100, opt=Adam(; lr=1e-5), device=CPU()) where T
    println("Starting training")
    for epoch in 1:n_epochs
        epoch_cost = 0
        t = @elapsed for (i, x) in enumerate(eachbatch(X, size=batch_size))
            x = device(x)
            ϵ = typeof(x)(rand(Normal(0, 1), size(m.enc2mu.W, 1), batch_size))
            cost, g = grad(loss_function, m, ϵ, x)
            update!(opt, m, g[2])
            epoch_cost += cost
        end
        println("Epoch $epoch: avg_cost=$(epoch_cost / (size(X,2) / batch_size)), elapsed=$t")
    end
    return m
end


text:
using PolyChaos, LinearAlgebra

function mytest(mOP::MultiOrthoPoly)
    @time P = [ computeSP([i,j,k],mOP) for i=0:mOP.dim-1,j=0:mOP.dim-1,k=0:mOP.dim-1 ];
    @time T = Tensor(dim,mOP);
    P_ = [ T.get([i,j,k]) for i=0:mOP.dim-1,j=0:mOP.dim-1,k=0:mOP.dim-1 ];
    @show norm(P-P_)
    return P_
end

##
deg, dim = 10, 3
α, β = 1.23, 3.23
## univariate
op = OrthoPoly("gaussian",deg;Nrec=100)
opq = OrthoPolyQ(op)
mOP = MultiOrthoPoly([opq],deg)
mytest(mOP)
## multivariate
op2 = OrthoPoly("beta01",deg,Dict(:shape_a=>α,:shape_b=>β);Nrec=100)
opq2 = OrthoPolyQ(op2)
mOP2 = MultiOrthoPoly([opq;opq2],deg)
mytest(mOP2)

mOP3 = MultiOrthoPoly([opq,opq,opq,opq],deg)
@time T = Tensor(dim,mOP3)

##
b = [He(deg)]
Gal = initBasis(b)
@time tensor(dim,Gal);

b2 = [He(deg),Ja(deg,α,β)]
Gal2 = initBasis(b2)
@time tensor(dim,Gal2);

b3 = [ He(deg),He(deg),He(deg),He(deg),He(deg),He(deg) ];
Gal3 = initBasis(b3);
@time tensor(dim,Gal3)