using Test
using OrdinaryDiffEq, OrdinaryDiffEqCore, ForwardDiff, FiniteDiff, LinearAlgebra, ADTypes,
    StaticArrays
import DifferentiationInterface as DI

# Version-dependent AD backend selection via DifferentiationInterface
# Zygote: Julia <= 1.11 only
# Enzyme: Julia <= 1.11 only (see https://github.com/EnzymeAD/Enzyme.jl/issues/2699)
# Mooncake: all versions
# ForwardDiff: all versions

# Define which backends are available based on Julia version
const JULIA_VERSION_ALLOWS_ENZYME_ZYGOTE = VERSION < v"1.12" && isempty(VERSION.prerelease)
# Enzyme segfaults on Julia 1.10 (compiler bug), but throws catchable errors on 1.11+
# Use this for @test_broken tests that would segfault on 1.10
const JULIA_VERSION_ENZYME_NO_SEGFAULT = VERSION >= v"1.11" && VERSION < v"1.12" && isempty(VERSION.prerelease)

# Load version-dependent packages and define helpers
# Note: Mooncake gradient support for ODE solves is currently broken (see discrete_adjoints.jl)
# Note: Enzyme requires set_runtime_activity for functions with internal closures that capture
# external variables (throws EnzymeRuntimeActivityError otherwise)
if JULIA_VERSION_ALLOWS_ENZYME_ZYGOTE
    using Enzyme
    get_gradient_backends() = [AutoForwardDiff(), AutoEnzyme(mode = Enzyme.set_runtime_activity(Enzyme.Reverse))]
    get_jacobian_backends() = [AutoForwardDiff(), AutoEnzyme(mode = Enzyme.set_runtime_activity(Enzyme.Forward))]
else
    get_gradient_backends() = [AutoForwardDiff()]
    get_jacobian_backends() = [AutoForwardDiff()]
end

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

# Callback tests with mutable closures
# These tests capture and mutate a `called` variable in the callback affect function.
# Enzyme cannot differentiate mutable closures (throws EnzymeMutabilityException).
# ForwardDiff works because it doesn't trace through closure mutation.
for x in 0:0.001:5
    called = false
    if x in [1.0, 2.0, 3.0, 4.0, 5.0]
        print("AD Ping $x")
    end
    function test_f(p)
        cb = ContinuousCallback(
            (u, t, i) -> u[1],
            (integrator) -> (called = true; integrator.p[2] = zero(integrator.p[2]))
        )
        prob = ODEProblem(f, eltype(p).([1.0, 0.0]), eltype(p).((0.0, 1.0)), copy(p))
        integrator = init(prob, Tsit5(), abstol = 1.0e-14, reltol = 1.0e-14, callback = cb)
        step!(integrator)
        return solve!(integrator).u[end]
    end
    p = [2.0, x]
    called = false
    findiff = FiniteDiff.finite_difference_jacobian(test_f, p)
    @test called
    called = false
    dijac = DI.jacobian(test_f, AutoForwardDiff(), p)
    @test called
    @test findiff ≈ dijac rtol = 1.0e-5
end

# Enzyme broken test for mutable closure callbacks (representative case)
# Enzyme throws EnzymeMutabilityException on mutable closures
# Note: On Julia 1.10, Enzyme segfaults instead of throwing, so we skip entirely
if JULIA_VERSION_ENZYME_NO_SEGFAULT
    @testset "Enzyme mutable closure limitation" begin
        called = false
        function test_f_enzyme(p)
            cb = ContinuousCallback(
                (u, t, i) -> u[1],
                (integrator) -> (called = true; integrator.p[2] = zero(integrator.p[2]))
            )
            prob = ODEProblem(f, eltype(p).([1.0, 0.0]), eltype(p).((0.0, 1.0)), copy(p))
            integrator = init(prob, Tsit5(), abstol = 1.0e-14, reltol = 1.0e-14, callback = cb)
            step!(integrator)
            return solve!(integrator).u[end]
        end
        p = [2.0, 3.0]
        findiff = FiniteDiff.finite_difference_jacobian(test_f_enzyme, p)
        # EnzymeMutabilityException: Cannot differentiate mutable closures
        @test_broken (
            dijac = DI.jacobian(test_f_enzyme, AutoEnzyme(mode = Enzyme.set_runtime_activity(Enzyme.Forward)), p);
            isapprox(dijac, findiff, rtol = 1.0e-5)
        )
    end
end

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

# Callback tests with mutable closures (f2 variant)
# Same limitation as above: Enzyme cannot differentiate mutable closures.
for x in 2.1:0.001:5
    called = false
    if x in [3.0, 4.0, 5.0]
        print("AD Ping $x")
    end
    function test_f2(p)
        cb = ContinuousCallback(
            (u, t, i) -> u[1],
            (integrator) -> (called = true; integrator.p[2] = zero(integrator.p[2]))
        )
        prob = ODEProblem(f2, eltype(p).([1.0, 0.0]), eltype(p).((0.0, 1.0)), copy(p))
        integrator = init(prob, Tsit5(), abstol = 1.0e-12, reltol = 1.0e-12, callback = cb)
        step!(integrator)
        return solve!(integrator).u[end]
    end
    p = [2.0, x]
    findiff = FiniteDiff.finite_difference_jacobian(test_f2, p)
    @test called
    called = false
    dijac = DI.jacobian(test_f2, AutoForwardDiff(), p)
    @test called
    @test findiff ≈ dijac rtol = 1.0e-5
end

#=
#x = 2.0 is an interesting case

x = 2.0

function test_f2(p)
  cb = ContinuousCallback((u,t,i) -> u[1], (integrator)->(@show(x,integrator.t);called=true;integrator.p[2]=zero(integrator.p[2])))
  prob = ODEProblem(f2,eltype(p).([1.0,0.0]),eltype(p).((0.0,1.0)),copy(p))
  integrator = init(prob,Tsit5(),abstol=1e-12,reltol=1e-12,callback=cb)
  step!(integrator)
  solve!(integrator).u[end]
end

p = [2.0, x]
findiff = FiniteDiff.finite_difference_jacobian(test_f2,p)
@test called
called = false
fordiff = ForwardDiff.jacobian(test_f2,p)
@test called

# At that value, it shouldn't be called, but a small perturbation will make it called, so finite difference is wrong!
=#

# Lotka-Volterra callback tests with mutable closures
# Same limitation as above: Enzyme cannot differentiate mutable closures.
for x in 1.0:0.001:2.5
    if x in [1.5, 2.0, 2.5]
        print("AD Ping $x")
    end
    function lotka_volterra(du, u, p, t)
        x, y = u
        α, β, δ, γ = p
        du[1] = dx = α * x - β * x * y
        return du[2] = dy = -δ * y + γ * x * y
    end
    u0 = [1.0, 1.0]
    tspan = (0.0, 10.0)
    p = [x, 1.0, 3.0, 1.0]
    prob = ODEProblem(lotka_volterra, u0, tspan, p)
    sol = solve(prob, Tsit5())

    called = false
    function test_lotka(p)
        cb = ContinuousCallback(
            (u, t, i) -> u[1] - 2.5,
            (integrator) -> (called = true; integrator.p[4] = 1.5)
        )
        prob = ODEProblem(
            lotka_volterra, eltype(p).([1.0, 1.0]), eltype(p).((0.0, 10.0)),
            copy(p)
        )
        integrator = init(prob, Tsit5(), abstol = 1.0e-12, reltol = 1.0e-12, callback = cb)
        step!(integrator)
        return solve!(integrator).u[end]
    end

    findiff = FiniteDiff.finite_difference_jacobian(test_lotka, p)
    @test called
    called = false
    dijac = DI.jacobian(test_lotka, AutoForwardDiff(), p)
    @test called
    @test findiff ≈ dijac rtol = 1.0e-5
end

# Gradients and Hessians
# Note: These functions create ODEProblem inside the differentiated function.
# Enzyme fails on SciML's type introspection code (isinplace, num_types_in_tuple).
# ForwardDiff works, so we use that and mark Enzyme as @test_broken.

function myobj(θ)
    f(u, p, t) = -θ[1] * u
    u0, _ = promote(10.0, θ[1])
    prob = ODEProblem(f, u0, (0.0, 1.0))
    sol = solve(prob, Tsit5())
    diff = sol.u - 10 * exp.(-sol.t)
    return diff'diff
end

function myobj2(θ)
    f(du, u, p, t) = (du[1] = -θ[1] * u[1])
    u0, _ = promote(10.0, θ[1])
    prob = ODEProblem(f, [u0], (0.0, 1.0))
    sol = solve(prob, Tsit5())
    diff = sol[:, 1] .- 10 .* exp.(-sol.t)
    return diff'diff
end

function myobj3(θ)
    f(u, p, t) = -θ[1] * u
    u0, _ = promote(10.0, θ[1])
    tspan_end, _ = promote(1.0, θ[1])
    prob = ODEProblem(f, u0, (0.0, tspan_end))
    sol = solve(prob, Tsit5())
    diff = sol.u - 10 * exp.(-sol.t)
    return diff'diff
end

function myobj4(θ)
    f(du, u, p, t) = (du[1] = -θ[1] * u[1])
    u0, _ = promote(10.0, θ[1])
    tspan_end, _ = promote(1.0, θ[1])
    prob = ODEProblem(f, [u0], (0.0, tspan_end))
    sol = solve(prob, Tsit5())
    diff = sol[:, 1] .- 10 .* exp.(-sol.t)
    return diff'diff
end

# ForwardDiff gradient tests (these work)
DI.gradient(myobj, AutoForwardDiff(), [1.0])
DI.gradient(myobj2, AutoForwardDiff(), [1.0])
DI.gradient(myobj3, AutoForwardDiff(), [1.0])
DI.gradient(myobj4, AutoForwardDiff(), [1.0])

# Hessian tests
DI.hessian(myobj, AutoForwardDiff(), [1.0])
DI.hessian(myobj2, AutoForwardDiff(), [1.0])
DI.hessian(myobj3, AutoForwardDiff(), [1.0])
DI.hessian(myobj4, AutoForwardDiff(), [1.0])

# Enzyme fails on these due to SciML type introspection (isinplace, num_types_in_tuple)
# Note: On Julia 1.10, Enzyme segfaults instead of throwing, so we skip entirely
if JULIA_VERSION_ENZYME_NO_SEGFAULT
    @testset "Enzyme ODEProblem construction limitation" begin
        # Enzyme has trouble differentiating through ODEProblem constructor
        # due to SciML's type introspection code
        @test_broken (DI.gradient(myobj, AutoEnzyme(mode = Enzyme.set_runtime_activity(Enzyme.Reverse)), [1.0]); true)
    end
end

f1s = function (du, u, p, t)
    du[1] = p[2] * u[1]
    du[2] = p[3] * u[2]
    du[3] = p[4] * u[3]
    du[4] = p[5] * u[4]
    return nothing
end
f2s = function (du, u, p, t)
    du[1] = p[1] * u[2]
    du[2] = p[1] * u[3]
    du[3] = p[1] * u[4]
    du[4] = p[1] * u[1]
    return nothing
end
u0 = [3.4, 3.3, 3.2, 3.1]
params = [0.002, -0.005, -0.004, -0.003, -0.002]
tspan = (7.0, 84.0)
times = collect(minimum(tspan):0.5:maximum(tspan))
prob = SplitODEProblem(f1s, f2s, u0, tspan, params)
sol2 = solve(prob, KenCarp4(); dt = 0.5, saveat = times)

# Enzyme fails on SplitODEProblem jacobians (mixed activity for jl_new_struct)
function difffunc(p)
    tmp_prob = remake(prob, p = p)
    return vec(solve(tmp_prob, KenCarp4(), saveat = times))
end
DI.jacobian(difffunc, AutoForwardDiff(), ones(5))

# https://github.com/SciML/OrdinaryDiffEq.jl/issues/1221

f_a = function (du, u, p, t)
    return du[1] = -p[1] * u[1] + exp(-t)
end

# Enzyme fails on this due to ODEProblem construction inside differentiated function
of_a = p -> begin
    u0 = [0.0]
    tspan = (0.0, 5.0)
    prob = ODEProblem(f_a, u0, tspan, p)
    # sol = solve(prob, Tsit5())                      # works
    # sol = solve(prob, Rodas5(autodiff=false))       # works
    sol = solve(prob, Rodas5(autodiff = AutoForwardDiff()), abstol = 1.0e-14, reltol = 1.0e-14)          # fails
    return sum(t -> abs2(t[1]), sol([1.0, 2.0, 3.0]))
end

@test !iszero(DI.gradient(of_a, AutoForwardDiff(), [1.0]))
@test DI.gradient(of_a, AutoForwardDiff(), [1.0]) ≈
    FiniteDiff.finite_difference_gradient(t -> of_a(t), [1.0]) rtol = 1.0e-5

SOLVERS_FOR_AD = (
    (BS3, 1.0e-12),
    (Tsit5, 1.0e-12),
    (KenCarp4, 1.0e-11),
    (KenCarp47, 1.0e-11),
    (KenCarp5, 1.0e-11),
    (KenCarp58, 1.0e-12),
    (TRBDF2, 1.0e-7),
    (Rodas4, 1.0e-12),
    (Rodas5, 1.0e-12),
    (Rosenbrock23, 1.0e-10),
    (Rosenbrock32, 1.0e-11),
    (Vern6, 1.0e-11),
    (Vern7, 1.0e-11),
    (RadauIIA3, 1.0e-4),
    (RadauIIA5, 1.0e-12),
)

@testset "$alg can handle ForwardDiff.Dual in u0 with rtol=$rtol when iip=$iip" for (alg, rtol) in SOLVERS_FOR_AD,
        iip in (
            true,
            false,
        )

    if iip
        f = (du, u, p, t) -> du .= -0.5 * u
    else
        f = (u, p, t) -> -0.5 * u
    end

    g = u0 -> begin
        tspan = (0.0, 1.0)
        prob = ODEProblem(
            f,
            u0,
            tspan
        )
        solve(prob, alg(), abstol = 1.0e-14, reltol = 1.0e-14)(last(tspan))[1]
    end
    @test DI.gradient(g, AutoForwardDiff(), [10.0])[1] ≈ exp(-0.5) rtol = rtol
end

@testset "$alg can handle ForwardDiff.Dual in t0 with rtol=$rtol when iip=$iip" for (alg, rtol) in SOLVERS_FOR_AD,
        iip in (
            true,
            false,
        )

    if iip
        f = (du, u, p, t) -> du .= -0.5 * u
    else
        f = (u, p, t) -> -0.5 * u
    end

    _u0 = 10.0
    g = t0 -> begin
        tspan = (t0, 1.0)
        u0 = typeof(t0)[_u0]
        prob = ODEProblem(
            f,
            u0,
            tspan
        )
        solve(prob, alg(), abstol = 1.0e-14, reltol = 1.0e-14)(last(tspan))[1]
    end
    @test DI.derivative(g, AutoForwardDiff(), 0.0) ≈ _u0 / 2 * exp(-0.5) rtol = rtol
end

# https://github.com/SciML/DifferentialEquations.jl/issues/903

#ode function
function foo!(du, u, p, t)
    du .= p .* u
    return nothing
end

#least squares objective function
function objfun(x, prob, data, solver, reltol, abstol)
    prob = remake(prob, p = x)
    sol = solve(prob, solver, reltol = reltol, abstol = abstol)
    ofv = 0.0
    if any((s.retcode != ReturnCode.Success for s in sol))
        ofv = 1.0e12
    else
        ofv = sum((sol .- data) .^ 2)
    end
    return ofv
end

p0 = [1.1, 1.2, 1.3, 1.4]
tspan = (0.0, 1.0)
u0 = 2 * ones(4)
saveat = 0.0:0.01:1.0
reltol = 1.0e-14
abstol = 1.0e-14
prob = ODEProblem{true}(foo!, u0, tspan, p0, saveat = saveat)
data = solve(prob, Tsit5(), reltol = reltol, abstol = abstol)
fn(x, solver) = objfun(x, prob, data, solver, reltol, abstol)

@test norm(DI.gradient(x -> fn(x, Tsit5()), AutoForwardDiff(), p0)) < 1.0e-9
@test norm(DI.gradient(x -> fn(x, Vern7()), AutoForwardDiff(), p0)) < 1.0e-9
@test norm(DI.gradient(x -> fn(x, Rodas4()), AutoForwardDiff(), p0)) < 1.0e-9
@test norm(DI.gradient(x -> fn(x, Rosenbrock23()), AutoForwardDiff(), p0)) < 1.0e-6
@test norm(DI.gradient(x -> fn(x, AutoTsit5(Rosenbrock23())), AutoForwardDiff(), p0)) < 1.0e-9
@test norm(DI.gradient(x -> fn(x, AutoVern9(Rodas4())), AutoForwardDiff(), p0)) < 1.0e-9

# test AD with LinearExponential()
function f_linexp(x)
    K = MatrixOperator(x)
    u0 = eltype(x).([1.0, 0.0])
    prob = ODEProblem(K, u0, (0.0, 10.0))
    return sol = solve(prob, LinearExponential(), tstops = [0.0, 10.0])[2, :]
end
K_ = [-1.0 0.0; 1.0 -1.0]
@test isapprox(DI.jacobian(f_linexp, AutoForwardDiff(), K_)[2], 0.00226999, atol = 1.0e-6)

implicit_algs = [
    FBDF,
    Rosenbrock23,
    TRBDF2,
]

@testset "deprecated AD keyword arguments still work with $alg" for alg in implicit_algs
    f = (du, u, p, t) -> du .= -0.5 * u
    alg1 = alg(autodiff = AutoForwardDiff())
    alg2 = @test_deprecated alg(autodiff = true)

    alg3 = alg(autodiff = AutoFiniteDiff())
    alg4 = @test_deprecated alg(autodiff = false)

    alg5 = alg(autodiff = AutoForwardDiff(chunksize = 5))
    alg6 = @test_deprecated alg(autodiff = true, chunk_size = 5)

    alg7 = alg(autodiff = AutoFiniteDiff(fdtype = Val(:central)))
    alg8 = @test_deprecated alg(autodiff = false, diff_type = Val(:central))

    alg9 = alg(autodiff = AutoForwardDiff(chunksize = 1))
    alg10 = @test_deprecated alg(chunk_size = 1)

    @test OrdinaryDiffEqCore.alg_autodiff(alg1) == OrdinaryDiffEqCore.alg_autodiff(alg2)
    @test OrdinaryDiffEqCore.alg_autodiff(alg3) == OrdinaryDiffEqCore.alg_autodiff(alg4)
    @test OrdinaryDiffEqCore.alg_autodiff(alg5) == OrdinaryDiffEqCore.alg_autodiff(alg6)
    @test OrdinaryDiffEqCore.alg_autodiff(alg7) == OrdinaryDiffEqCore.alg_autodiff(alg8)
    @test OrdinaryDiffEqCore.alg_autodiff(alg9) == OrdinaryDiffEqCore.alg_autodiff(alg10)
end

# https://github.com/SciML/OrdinaryDiffEq.jl/issues/2675
x0 = [0.1]
DI.gradient(AutoForwardDiff(), x0) do x
    prob = ODEProblem{true}((du, u, p, t) -> (du[1] = -u[1]), x, (0.0, 1.0))
    sol = solve(prob, DefaultODEAlgorithm(), reltol = 1.0e-6)
    sum(sum, sol.u)
end ≈ [6.765310476296564]

# Test with multiple AD backends
# Note: Enzyme reverse mode through ODE solve requires SciMLSensitivity for adjoint methods.
# ForwardDiff works without SciMLSensitivity, so we only test ForwardDiff here.
@testset "DifferentiationInterface multi-backend tests" begin
    # Simple ODE for testing
    function simple_ode!(du, u, p, t)
        du[1] = -p[1] * u[1]
        return nothing
    end

    function loss_fn(p)
        prob = ODEProblem(simple_ode!, [1.0], (0.0, 1.0), p)
        sol = solve(prob, Tsit5(), abstol = 1.0e-12, reltol = 1.0e-12)
        return sum(sol.u[end])
    end

    p_test = [0.5]
    ref_grad = FiniteDiff.finite_difference_gradient(loss_fn, p_test)

    # ForwardDiff works without SciMLSensitivity
    grad = DI.gradient(loss_fn, AutoForwardDiff(), p_test)
    @test grad ≈ ref_grad rtol = 1.0e-6
end

# Tests migrated from DiffEqBase downstream to cover complex numbers, StaticArrays,
# and ensemble AD scenarios (previously tested via SciMLSensitivity integration).

@testset "Complex number ODE AD" begin
    # Complex ODE: verify ForwardDiff gradient through complex-valued solve matches
    # the same problem reformulated in reals.
    pd = 3
    H0 = rand(ComplexF64, pd, pd)
    A = rand(ComplexF64, pd, pd)
    function f_complex!(du, u, p, t)
        a, b, c = p
        du .= (A * u) * (a * cos(b * t + c))
        du .+= H0 * u
        return nothing
    end
    u0_c = hcat(normalize(rand(ComplexF64, pd)), normalize(rand(pd)))
    prob_c = ODEProblem(
        f_complex!, u0_c, (0.0, 1.0), rand(3);
        saveat = range(0.0, 1.0, length = 3), reltol = 1.0e-6, alg = Tsit5()
    )
    cost_c(u) = abs2(tr(first(u)' * u[2])) - abs2(tr(first(u)' * last(u)))

    function loss_complex(p)
        prob = remake(prob_c; p)
        sol = solve(prob)
        return cost_c(sol.u) + sum(p) / 10
    end

    # Same problem via reals
    function real_f!(du, u, p, t)
        complex_u = complex.(selectdim(u, 3, 1), selectdim(u, 3, 2))
        complex_du = copy(complex_u)
        prob_c.f(complex_du, complex_u, p, t)
        selectdim(du, 3, 1) .= real(complex_du)
        selectdim(du, 3, 2) .= imag(complex_du)
        return nothing
    end
    prob_real = remake(
        prob_c; f = real_f!,
        u0 = cat(real(prob_c.u0), imag(prob_c.u0); dims = 3)
    )

    function loss_real(p)
        prob = remake(prob_real; p)
        sol = solve(prob)
        u = [complex.(selectdim(u, 3, 1), selectdim(u, 3, 2)) for u in sol.u]
        return cost_c(u) + sum(p) / 10
    end

    p0 = rand(3)
    grad_real = DI.gradient(loss_real, AutoForwardDiff(), p0)
    grad_complex = DI.gradient(loss_complex, AutoForwardDiff(), p0)
    @test all(isfinite, grad_complex)
    @test grad_complex ≈ grad_real rtol = 1.0e-4
end

@testset "StaticArrays AD through solve" begin
    f_sa(u, p, t) = copy(u)

    du1 = DI.derivative(AutoForwardDiff(), 5.0) do x
        prob = ODEProblem(f_sa, [x], (0.0, 1.0), nothing)
        sol = solve(prob, Tsit5())
        sol.u[end][1]
    end

    du2 = DI.derivative(AutoForwardDiff(), 5.0) do x
        prob = ODEProblem(f_sa, SVector(x), (0.0, 1.0), nothing)
        sol = solve(prob, Tsit5())
        sol.u[end][1]
    end

    @test du1 ≈ du2
end

@testset "Ensemble AD" begin
    function lv!(du, u, p, t)
        du[1] = p[1] * u[1] - p[2] * u[1] * u[2]
        return du[2] = -p[3] * u[2] + p[4] * u[1] * u[2]
    end

    p_lv = [1.5, 1.0, 3.0, 1.0]
    u0_lv = [1.0, 1.0]
    prob_lv = ODEProblem(lv!, u0_lv, (0.0, 10.0), p_lv)

    function sum_of_solution_ens(x)
        _prob = remake(prob_lv, u0 = x[1:2], p = x[3:end])
        return sum(solve(_prob, Tsit5(), saveat = 0.1))
    end

    gs = DI.gradient(sum_of_solution_ens, AutoForwardDiff(), [u0_lv; p_lv])
    ref = FiniteDiff.finite_difference_gradient(sum_of_solution_ens, [u0_lv; p_lv])
    # Finite differences lose accuracy over long integrations; ForwardDiff is exact.
    @test gs ≈ ref rtol = 0.1
end
