using OrdinaryDiffEq, OrdinaryDiffEqCore, DiffEqBase, Test, ADTypes
using Random
using OrdinaryDiffEqDefault
using ElasticArrays, LinearSolve
using OrdinaryDiffEqBDF, OrdinaryDiffEqExtrapolation, OrdinaryDiffEqFeagin, OrdinaryDiffEqHighOrderRK, OrdinaryDiffEqLowOrderRK, OrdinaryDiffEqLowStorageRK, OrdinaryDiffEqRosenbrock, OrdinaryDiffEqSDIRK, OrdinaryDiffEqSSPRK
using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
using NonlinearSolve: NewtonRaphson
Random.seed!(213)
CACHE_TEST_ALGS = [
    Euler(), Midpoint(), RK4(), SSPRK22(), SSPRK33(), SSPRK43(), SSPRK104(),
    CarpenterKennedy2N54(), SHLDDRK64(), ORK256(), DGLDDRK73_C(),
    CFRLDDRK64(), TSLDDRK74(), CKLLSRK43_2(), ParsaniKetchesonDeconinck3S32(),
    BS3(), BS5(), DP5(), DP8(), Feagin10(), Feagin12(), Feagin14(), TanYam7(),
    Tsit5(), TsitPap8(), Vern6(), Vern7(), Vern8(), Vern9(), OwrenZen3(), OwrenZen4(),
    OwrenZen5(), AutoTsit5(Rosenbrock23()), TRBDF2(), KenCarp4(), ABDF2(),
    OrdinaryDiffEqDefault.DefaultODEAlgorithm(),
]
broken_CACHE_TEST_ALGS = [
    QNDF(),
    ExtrapolationMidpointHairerWanner(),
    ImplicitEulerExtrapolation(),
    ImplicitDeuflhardExtrapolation(),
]
# AitkenNeville(threading=false) fails Elastic but not normal case

using InteractiveUtils

NON_IMPLICIT_ALGS = filter(
    (x) -> isconcretetype(x) && !OrdinaryDiffEqCore.isimplicit(x()),
    union(
        subtypes(OrdinaryDiffEqCore.OrdinaryDiffEqAlgorithm),
        subtypes(OrdinaryDiffEqCore.OrdinaryDiffEqAdaptiveAlgorithm)
    )
)

f = function (du, u, p, t)
    for i in 1:length(u)
        du[i] = (0.3 / length(u)) * u[i]
    end
    return
end

condition = function (u, t, integrator)
    return 1 - maximum(u)
end

affect! = function (integrator)
    u = integrator.u
    maxidx = findmax(u)[2]
    resize!(integrator, length(u) + 1)
    Θ = rand() / 5 + 0.25
    u[maxidx] = Θ
    u[end] = 1 - Θ
    return nothing
end

callback = ContinuousCallback(condition, affect!)

u0 = [0.2]
tspan = (0.0, 10.0)
prob = ODEProblem(f, u0, tspan)

println("Check for stochastic errors")
for i in 1:10
    @test_nowarn sol = solve(prob, Tsit5(), callback = callback)
end

println("Check some other integrators")
sol = solve(prob, Rosenbrock23(), callback = callback, dt = 1 / 2)
@test length(sol.u[end]) > 1
sol = solve(prob, Rosenbrock32(), callback = callback, dt = 1 / 2)
@test length(sol.u[end]) > 1
sol = solve(prob, KenCarp4(), callback = callback, dt = 1 / 2)
@test length(sol.u[end]) > 1
sol = solve(prob, TRBDF2(), callback = callback, dt = 1 / 2)
@test length(sol.u[end]) > 1
sol = solve(
    prob, TRBDF2(linsolve = LinearSolve.KrylovJL_GMRES()),
    callback = callback
)
@test length(sol.u[end]) > 1
sol = solve(
    prob, TRBDF2(nlsolve = NonlinearSolveAlg(NewtonRaphson())),
    callback = callback, dt = 1 / 2
)
@test length(sol.u[end]) > 1
sol = solve(
    prob, KenCarp4(nlsolve = NonlinearSolveAlg(NewtonRaphson())),
    callback = callback, dt = 1 / 2
)
@test length(sol.u[end]) > 1
sol = solve(
    prob, ImplicitEuler(nlsolve = NonlinearSolveAlg(NewtonRaphson())),
    callback = callback, dt = 1 / 2
)
@test length(sol.u[end]) > 1

for alg in CACHE_TEST_ALGS
    @show alg
    local sol = solve(prob, alg, callback = callback, dt = 1 / 2)
    @test length(sol.u[end]) > 1
end

for alg in broken_CACHE_TEST_ALGS
    @show alg
    @test_broken length(solve(prob, alg, callback = callback, dt = 1 / 2)[end]) > 1
end

sol = solve(
    prob, Rodas4(autodiff = AutoForwardDiff(chunksize = 1)),
    callback = callback, dt = 1 / 2
)
@test length(sol.u[end]) > 1
sol = solve(
    prob, Rodas5(autodiff = AutoForwardDiff(chunksize = 1)),
    callback = callback, dt = 1 / 2
)
@test length(sol.u[end]) > 1

# cache tests resizing multidimensional arrays
println("Check resizing multidimensional arrays")
u0_matrix = ElasticArray(ones(2, 2))
f_matrix = (du, u, p, t) -> du .= u
prob_matrix = ODEProblem(f_matrix, u0_matrix, (0.0, 2.0))
condition_matrix = (u, t, integrator) -> t - 1
affect_matrix! = function (integrator)
    resize!(integrator, (2, 3))
    integrator.u .= 1
    return nothing
end
callback_matrix = ContinuousCallback(condition_matrix, affect_matrix!)

for alg in CACHE_TEST_ALGS
    OrdinaryDiffEqCore.isimplicit(alg) && continue # this restriction should be removed in the future
    @show alg
    local sol = solve(prob_matrix, alg, callback = callback_matrix, dt = 1 / 2)
    @test size(sol.u[end]) == (2, 3)
end

# additional cache tests to find more bugs
println("Additional resize! checks")
u0resize3 = ones(2)
fresize3 = (du, u, p, t) -> du .= u
prob_resize3 = ODEProblem(fresize3, u0resize3, (0.0, 2.0))
condition_resize3 = (u, t, integrator) -> t - 1
affect!_resize3 = function (integrator)
    resize!(integrator, 3)
    integrator.u .= 1
    return nothing
end
callback_resize3 = ContinuousCallback(condition_resize3, affect!_resize3)

for alg in CACHE_TEST_ALGS
    (OrdinaryDiffEqCore.isimplicit(alg) || OrdinaryDiffEqCore.alg_order(alg) < 2) &&
        continue
    @show alg
    local sol = solve(prob_resize3, alg, callback = callback_resize3, dt = 0.125)
    @test size(sol.u[end]) == (3,)
    @test all(sol.u[end] .== sol.u[end][1])
    @test sol.u[end][1] ≈ exp(1) atol = 1.0e-2
end

# additional cache tests to find more bugs
println("Enforced adaptive dt checks")
u0_adapt = zeros(2)
f_adapt = (du, u, p, t) -> du .= t
prob_adapt = ODEProblem(f_adapt, u0_adapt, (0.0, 2.0))
condition_adapt = (u, t, integrator) -> true
affect!_adapt = function (integrator)
    dt = rand((0.0625, 0.125, 0.25))
    set_proposed_dt!(integrator, dt)
    integrator.opts.dtmax = dt
    integrator.dtcache = dt
    derivative_discontinuity!(integrator, false)
    return nothing
end
callback_adapt = DiscreteCallback(
    condition_adapt, affect!_adapt,
    save_positions = (false, false)
)

for alg in CACHE_TEST_ALGS
    (OrdinaryDiffEqCore.isimplicit(alg) || OrdinaryDiffEqCore.alg_order(alg) < 2) &&
        continue
    @show alg
    local sol = solve(prob_adapt, alg, callback = callback_adapt, dt = 0.125)
    @test all(
        idx -> all(isapprox.(sol.u[idx], 0.5 * sol.t[idx]^2, atol = 1.0e-6)),
        eachindex(sol.t)
    )
end

# Force switching

function f2(du, u, p, t)
    @assert length(u) == length(du) "length(u) = $(length(u)), length(du) = $(length(du)) at time $(t)"
    for i in 1:length(u)
        if t > 10
            du[i] = -10000 * u[i]
        else
            du[i] = 0.3 * u[i]
        end
    end
    return du
end

function condition2(u, t, integrator)
    return 1 - maximum(u)
end

function affect2!(integrator)
    u = integrator.u
    resize!(integrator, length(u) + 1)
    maxidx = findmax(u)[2]
    Θ = rand()
    u[maxidx] = Θ
    u[end] = 1 - Θ
    return nothing
end

callback = ContinuousCallback(condition2, affect2!)
u0 = [0.2]
tspan = (0.0, 20.0)
prob = ODEProblem(f2, u0, tspan)
sol = solve(prob, AutoTsit5(Rosenbrock23()), callback = callback)
@test length(sol.u[end]) > 1
