using OrdinaryDiffEq, DiffEqBase, DiffEqCallbacks, Test
using Random, DiffEqOperators
using ElasticArrays
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()]
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) && !OrdinaryDiffEq.isimplicit(x()),union(subtypes(OrdinaryDiffEq.OrdinaryDiffEqAlgorithm),subtypes(OrdinaryDiffEq.OrdinaryDiffEqAdaptiveAlgorithm)))

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

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

affect! = function (integrator)
  u = integrator.u
  resize!(integrator,length(u)+1)
  maxidx = findmax(u)[2]
  Θ = rand()/5 + 0.25
  u[maxidx] = Θ
  u[end] = 1-Θ
  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[end]) > 1
sol = solve(prob,Rosenbrock32(),callback=callback,dt=1/2)
@test length(sol[end]) > 1
sol = solve(prob,KenCarp4(),callback=callback,dt=1/2)
@test length(sol[end]) > 1
sol = solve(prob,TRBDF2(),callback=callback,dt=1/2)
@test length(sol[end]) > 1

Jv = JacVecOperator(f,u0,nothing,0.0)
ff2 = ODEFunction(f;jac_prototype=Jv)
prob2 = ODEProblem(ff2,u0,tspan)
sol = solve(prob2,TRBDF2(linsolve=LinSolveGMRES()),callback=callback)

for alg in CACHE_TEST_ALGS
  @show alg
  local sol = solve(prob,alg,callback=callback,dt=1/2)
  @test length(sol[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(chunk_size=1),callback=callback,dt=1/2)
@test length(sol[end]) > 1
sol = solve(prob,Rodas5(chunk_size=1),callback=callback,dt=1/2)
@test length(sol[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
  nothing
end
callback_matrix = ContinuousCallback(condition_matrix, affect_matrix!)

for alg in CACHE_TEST_ALGS
  OrdinaryDiffEq.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[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
  nothing
end
callback_resize3 = ContinuousCallback(condition_resize3, affect!_resize3)

for alg in CACHE_TEST_ALGS
  (OrdinaryDiffEq.isimplicit(alg) || OrdinaryDiffEq.alg_order(alg) < 2) && continue
  @show alg
  local sol = solve(prob_resize3, alg, callback=callback_resize3, dt=0.125)
  @test size(sol[end]) == (3,)
  @test all(sol[end] .== sol[end][1])
  @test sol[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
  u_modified!(integrator, false)
  nothing
end
callback_adapt = DiscreteCallback(condition_adapt, affect!_adapt, save_positions=(false,false))

for alg in CACHE_TEST_ALGS
  (OrdinaryDiffEq.isimplicit(alg) || OrdinaryDiffEq.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)
  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-Θ
  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[end]) > 1
