using OrdinaryDiffEqVerner, OrdinaryDiffEqCore, OrdinaryDiffEqExplicitRK
import OrdinaryDiffEqExplicitTableaus
using DiffEqDevTools, Test
import ODEProblemLibrary: prob_ode_linear, prob_ode_2Dlinear, prob_ode_bigfloatlinear, prob_ode_bigfloat2Dlinear

# Problem mappings
const probbig = prob_ode_bigfloat2Dlinear
const probnumbig = prob_ode_bigfloatlinear
const probnum = prob_ode_linear
const prob = prob_ode_2Dlinear

testTol = 2.0

# Custom ODE functions for testing
function f!(du, u, p, t)
    return du[1] = -u[1]
end

function f(u, p, t)
    return -u
end

t_end = 64.0
setprecision(256)
prob_oop = ODEProblem(ODEFunction(f; analytic = (u0, p, t) -> exp(-t)), 1.0, (0.0, t_end))
prob_iip = ODEProblem(ODEFunction(f!; analytic = (u0, p, t) -> [exp(-t)]), [1.0], (0.0, t_end))

# -------------------------------------------------------------
# Helper for convergence test consistency
function check_convergence(dts, prob, alg, order_expected)
    # Use DiffEqDevTools for overall convergence estimate
    sim = test_convergence(dts, prob, alg)
    return @test (sim.𝒪est[:final] > order_expected) || (abs(sim.𝒪est[:final] - order_expected) < testTol)

end
# -------------------------------------------------------------

### Vern6
println("Vern6")
dts = (1 / 2) .^ (8:-1:5)
check_convergence(dts, probnumbig, Vern6(), 6)
check_convergence(dts, probbig, Vern6(), 6)

tabalg = ExplicitRK(tableau = OrdinaryDiffEqExplicitTableaus.VernerEfficient6(BigFloat))
sol1 = solve(probnumbig, Vern6(); dt = 1 / 2^6, adaptive = false, save_everystep = false)
sol2 = solve(probnumbig, tabalg; dt = 1 / 2^6, adaptive = false, save_everystep = false)
@test sol1.u[end] - sol2.u[end] < 1.0e-10

sol1 = solve(probbig, Vern6(); dt = 1 / 2^3, adaptive = false, save_everystep = false)
sol2 = solve(probbig, tabalg; dt = 1 / 2^3, adaptive = false, save_everystep = false)
@test minimum(sol1.u[end] - sol2.u[end] .< 1.0e-10)

sol1 = solve(probbig, tabalg; dt = 1 / 2^6)
sol2 = solve(probbig, Vern6(); dt = 1 / 2^6)
@test length(sol1.t) == length(sol2.t)
@test SciMLBase.successful_retcode(sol1)
@test SciMLBase.successful_retcode(sol2)

# -------------------------------------------------------------
### Vern7
println("Vern7")
dts = (1 / 2) .^ (6:-1:3)
check_convergence(dts, probnumbig, Vern7(), 7)
check_convergence(dts, probbig, Vern7(), 7)

tabalg = ExplicitRK(tableau = OrdinaryDiffEqExplicitTableaus.Verner7(BigFloat))
sol1 = solve(probnumbig, Vern7(); dt = 1 / 2^6, adaptive = false, save_everystep = false)
sol2 = solve(probnumbig, tabalg; dt = 1 / 2^6, adaptive = false, save_everystep = false)
@test sol1.u[end] - sol2.u[end] < 1.0e-10

sol1 = solve(probbig, Vern7(); dt = 1 / 2^3, adaptive = false, save_everystep = false)
sol2 = solve(probbig, tabalg; dt = 1 / 2^3, adaptive = false, save_everystep = false)
@test minimum(sol1.u[end] - sol2.u[end] .< 1.0e-10)

sol1 = solve(probbig, tabalg; dt = 1 / 2^6)
sol2 = solve(probbig, Vern7(); dt = 1 / 2^6)
@test length(sol1.t) == length(sol2.t)
@test SciMLBase.successful_retcode(sol1)
@test SciMLBase.successful_retcode(sol2)

# -------------------------------------------------------------
### Vern8
println("Vern8")
dts = (1 / 2) .^ (6:-1:3)
check_convergence(dts, probnumbig, Vern8(), 8)
check_convergence(dts, probbig, Vern8(), 8)

tabalg = ExplicitRK(tableau = OrdinaryDiffEqExplicitTableaus.Verner8(BigFloat))
sol1 = solve(probnumbig, Vern8(); dt = 1 / 2^6, adaptive = false, save_everystep = false)
sol2 = solve(probnumbig, tabalg; dt = 1 / 2^6, adaptive = false, save_everystep = false)
@test sol1.u[end] - sol2.u[end] < 1.0e-10

sol1 = solve(probbig, Vern8(); dt = 1 / 2^3, adaptive = false, save_everystep = false)
sol2 = solve(probbig, tabalg; dt = 1 / 2^3, adaptive = false, save_everystep = false)
@test minimum(sol1.u[end] - sol2.u[end] .< 1.0e-10)

sol1 = solve(prob, tabalg; dt = 1 / 2^6)
sol2 = solve(prob, Vern8(); dt = 1 / 2^6)
@test length(sol1.t) == length(sol2.t)
@test SciMLBase.successful_retcode(sol1)
@test SciMLBase.successful_retcode(sol2)

# -------------------------------------------------------------
### Vern9
println("Vern9")
dts = (1 / 2) .^ (6:-1:3)
check_convergence(dts, probnumbig, Vern9(), 9)
check_convergence(dts, probbig, Vern9(), 9)

tabalg = ExplicitRK(tableau = OrdinaryDiffEqExplicitTableaus.VernerEfficient9(BigFloat))
sol1 = solve(probnumbig, Vern9(); dt = 1 / 2^6, adaptive = false, save_everystep = false)
sol2 = solve(probnumbig, tabalg; dt = 1 / 2^6, adaptive = false, save_everystep = false)
@test abs(sol1.u[end] - sol2.u[end]) < 1.0e-15

sol1 = solve(probbig, Vern9(); dt = 1 / 2^3, adaptive = false, save_everystep = false)
sol2 = solve(probbig, tabalg; dt = 1 / 2^3, adaptive = false, save_everystep = false)
@test minimum(abs.(sol1.u[end] - sol2.u[end]) .< 1.0e-15)

sol1 = solve(probbig, tabalg; dt = 1 / 2^6)
sol2 = solve(probbig, Vern9(); dt = 1 / 2^6)
@test length(sol1.t) == length(sol2.t)
@test SciMLBase.successful_retcode(sol1)
@test SciMLBase.successful_retcode(sol2)

# -------------------------------------------------------------
### RKV76IIa
println("RKV76IIa")
# dts = [1, 0.5, 0.25, 0.125, 0.0625, 0.03125, 0.015625, 0.0078125, 0.00390625, 0.001953125, 0.0009765625, 0.00048828125, 0.000244140625, 0.0001220703125, 0.00006103515625]
dts = [2, 1, 0.5, 0.25]

check_convergence(dts, prob_oop, RKV76IIa(), 7)

# -------------------------------------------------------------
### Backward solve lazy interpolation dt sign regression test
# Tests that lazy interpolation during a backward solve produces correct results.
# Previously, modify_dt_for_tstops! stored dtpropose without the sign, causing
# addsteps! to evaluate at times outside the step interval on the last step.
println("Backward solve lazy interpolation")

function backward_ode!(du, u, p, t)
    du[1] = -u[1]
    return du[2] = -2 * u[2]
end

prob_back = ODEProblem(backward_ode!, [1.0, 1.0], (1.0, 0.0))

interp_lazy = Float64[]
interp_nolazy = Float64[]

cb_lazy = DiscreteCallback(
    (u, t, int) -> true,
    integrator -> begin
        t_mid = (integrator.tprev + integrator.t) / 2
        curu = similar(integrator.u)
        integrator(curu, t_mid)
        push!(interp_lazy, curu[1])
        derivative_discontinuity!(integrator, false)
    end,
    save_positions = (false, false)
)

cb_nolazy = DiscreteCallback(
    (u, t, int) -> true,
    integrator -> begin
        t_mid = (integrator.tprev + integrator.t) / 2
        curu = similar(integrator.u)
        integrator(curu, t_mid)
        push!(interp_nolazy, curu[1])
        derivative_discontinuity!(integrator, false)
    end,
    save_positions = (false, false)
)

solve(
    prob_back, Vern9(lazy = Val{true}()), abstol = 1.0e-12, reltol = 1.0e-12,
    callback = cb_lazy, save_everystep = false
)
solve(
    prob_back, Vern9(lazy = Val{false}()), abstol = 1.0e-12, reltol = 1.0e-12,
    callback = cb_nolazy, save_everystep = false
)

@test length(interp_lazy) == length(interp_nolazy)
@test maximum(abs.(interp_lazy .- interp_nolazy)) < 1.0e-10

### Interpolation inside a step truncated by a ContinuousCallback
# The extra interpolation stages of the non-lazy interpolants live inside
# kshortsize, so a shortened step must explicitly force them to be recomputed
# rather than reuse the ones built for the original dt.
println("Interpolation in a callback-truncated step")

exp_prob = ODEProblem((du, u, p, t) -> (du[1] = u[1]; nothing), [1.0], (0.0, 5.0))
root_cb = ContinuousCallback((u, t, integ) -> u[1] - 50.0, integ -> nothing)

function worst_interior_error(sol, lo, hi)
    return maximum(
        abs(sol(t)[1] - exp(t)) / exp(t)
            for t in range(lo, hi, length = 21)[2:(end - 1)]
    )
end

for alg in (
        Vern6(lazy = false), Vern7(lazy = false),
        Vern8(lazy = false), Vern9(lazy = false),
        Vern6(), Vern7(), Vern8(), Vern9(),
    )
    sol = solve(
        exp_prob, alg; abstol = 1.0e-10, reltol = 1.0e-10, callback = root_cb
    )
    i = findfirst(≈(log(50.0)), sol.t)
    @test worst_interior_error(sol, sol.t[i - 1], sol.t[i]) < 1.0e-8
end
