using OrdinaryDiffEq, Test, Random, StaticArrays, DiffEqCallbacks
using OrdinaryDiffEqCore: add_tstop!
import ODEProblemLibrary: prob_ode_linear
using OrdinaryDiffEqLowOrderRK

Random.seed!(100)

@testset "Tstops Tests on the Interval [0, 1]" begin
    prob = prob_ode_linear

    sol = solve(prob, Tsit5(), dt = 1 // 2^(6), tstops = [1 / 2])
    @test 1 // 2 ∈ sol.t

    sol = solve(prob, RK4(), dt = 1 // 3, tstops = [1 / 2], adaptive = false)
    @test sol.t == [0, 1 / 3, 1 / 2, 1 / 3 + 1 / 2, 1]

    sol = solve(
        prob, RK4(), dt = 1 // 3, tstops = [1 / 2],
        d_discontinuities = [-1 / 2, 1 / 2, 3 / 2], adaptive = false
    )
    @test sol.t == [0, 1 / 3, 1 / 2, 1 / 3 + 1 / 2, 1]

    # TODO
    integrator = init(
        prob, RK4(), tstops = [1 / 5, 1 / 4, 1 / 3, 1 / 2, 3 / 4],
        adaptive = false
    )

    sol = solve(prob, RK4(), tstops = [1 / 5, 1 / 4, 1 / 3, 1 / 2, 3 / 4], adaptive = false)
    @test sol.t == [0, 1 / 5, 1 / 4, 1 / 3, 1 / 2, 3 / 4, 1]

    sol = solve(
        prob, RK4(), tstops = [0, 1 / 5, 1 / 4, 1 / 3, 1 / 2, 3 / 4, 1],
        adaptive = false
    )
    @test sol.t == [0, 1 / 5, 1 / 4, 1 / 3, 1 / 2, 3 / 4, 1]

    sol = solve(prob, RK4(), tstops = 0:(1 // 16):1, adaptive = false)
    @test sol.t == collect(0:(1 // 16):1)

    sol = solve(prob, RK4(), tstops = range(0, stop = 1, length = 100), adaptive = false)
    @test sol.t == collect(range(0, stop = 1, length = 100))
end

@testset "Integrator Tstops Tests on the Interval $(["[-1, 0]", "[0, 1]"][i])" for (i, tdir) in enumerate(
        [
            -1.0;
            1.0
        ]
    )
    prob2 = remake(prob_ode_linear, tspan = (0.0, tdir * 1.0))
    integrator = init(prob2, Tsit5())
    tstops = tdir .* [0, 1 / 5, 1 / 4, 1 / 3, 1 / 2, 3 / 4, 1]
    for tstop in tstops
        add_tstop!(integrator, tstop)
    end
    @test_throws ErrorException add_tstop!(integrator, -0.1 * tdir)
    solve!(integrator)
    for tstop in tstops
        @test tstop ∈ integrator.sol.t
    end
end

@testset "Tstops Eps" begin
    function de(du, u, p, t) # specific DE does not impact the issue
        a, b = p
        du[1] = a * u[1]
        du[2] = b * u[2]
    end

    saveat = [0.0, 0.0094777, 1.5574]
    tstop = 0.010823

    affect!(integrator) = integrator.u[1] += 1.0
    condition(u, t, integrator) = t == tstop
    callback = DiscreteCallback(condition, affect!)

    prob = ODEProblem(de, zeros(2), (-1, 3.0), rand(2))
    sol = solve(prob, Tsit5(), saveat = saveat, tstops = tstop, callback = callback)
    @test sol.t[end] == 1.5574
end

@testset "Tstops Type Conversion" begin
    called = Ref(false)
    tval = rand()
    ff(du, u, p, t) = du .= 0
    cb = DiscreteCallback(
        (u, t, integrator) -> t == Float32(tval),
        integrator -> (called[] = true)
    )
    prob = ODEProblem(ff, [0.0], (0.0f0, 1.0f0))
    sol = solve(prob, Tsit5(), tstops = [tval], callback = cb)
end

@testset "Late binding tstops" begin
    function rhs(u, p, t)
        u * p + t
    end
    prob = ODEProblem(rhs, 1.0, (0.0, 1.0), 0.1; tstops = (p, tspan) -> tspan[1]:p:tspan[2])
    sol = solve(prob, Tsit5())
    @test 0.0:0.1:1.0 ⊆ sol.t
    prob2 = remake(prob; p = 0.07)
    sol2 = solve(prob2, Tsit5())
    @test 0.0:0.07:1.0 ⊆ sol2.t
end

# Tests for issue #2752: tstop overshoot errors with StaticArrays

@testset "StaticArrays vs Arrays with extreme precision" begin
    function precise_dynamics(u, p, t)
        x = @view u[1:2]
        v = @view u[3:4]
        dv = -0.01 * x + 1.0e-6 * sin(100 * t) * SVector{2}(1, 1)
        return SVector{4}(v[1], v[2], dv[1], dv[2])
    end

    function precise_dynamics_array!(du, u, p, t)
        x = @view u[1:2]
        v = @view u[3:4]
        dv = -0.01 * x + 1.0e-6 * sin(100 * t) * [1, 1]
        du[1] = v[1]
        du[2] = v[2]
        du[3] = dv[1]
        du[4] = dv[2]
    end

    u0_static = SVector{4}(1.0, -0.5, 0.01, 0.01)
    u0_array = [1.0, -0.5, 0.01, 0.01]
    tspan = (0.0, 2.0)
    tstops = [0.5, 1.0, 1.5]

    prob_static = ODEProblem(precise_dynamics, u0_static, tspan)
    sol_static = solve(
        prob_static, Vern9(); reltol = 1.0e-12, abstol = 1.0e-15,
        tstops = tstops
    )
    @test successful_retcode(sol_static)
    for tstop in tstops
        @test tstop ∈ sol_static.t
    end

    prob_array = ODEProblem(precise_dynamics_array!, u0_array, tspan)
    sol_array = solve(
        prob_array, Vern9(); reltol = 1.0e-12, abstol = 1.0e-15,
        tstops = tstops
    )
    @test successful_retcode(sol_array)
    for tstop in tstops
        @test tstop ∈ sol_array.t
    end

    @test isapprox(sol_static(2.0), sol_array(2.0), rtol = 1.0e-10)
end

@testset "Backward integration with tstop flags" begin
    function decay_ode(u, p, t)
        SA[-0.1 * u[1]]
    end

    u0 = SVector{1}(1.0)
    tspan = (2.0, 0.0)
    tstops = [1.5, 1.0, 0.5]

    prob = ODEProblem(decay_ode, u0, tspan)
    sol = solve(prob, Vern9(); tstops = tstops, reltol = 1.0e-12, abstol = 1.0e-15)
    @test successful_retcode(sol)
    for tstop in tstops
        @test tstop ∈ sol.t
    end
end

@testset "PresetTimeCallback with tstop flags" begin
    callback_times = Float64[]

    function affect_preset!(integrator)
        push!(callback_times, integrator.t)
        integrator.u += 0.1 * integrator.u
    end

    function simple_growth(u, p, t)
        SA[0.1 * u[1]]
    end

    u0 = SA[1.0]
    tspan = (0.0, 3.0)
    critical_times = [0.5, 1.0, 1.5, 2.0, 2.5]

    preset_cb = PresetTimeCallback(critical_times, affect_preset!)

    prob = ODEProblem(simple_growth, u0, tspan)
    sol = solve(
        prob, Vern9(); tstops = critical_times, callback = preset_cb,
        reltol = 1.0e-10, abstol = 1.0e-12
    )

    @test successful_retcode(sol)
    for time in critical_times
        @test any(abs.(sol.t .- time) .< 1.0e-10)
    end
    @test length(callback_times) == length(critical_times)
end

@testset "Multiple close tstops with StaticArrays" begin
    function oscillator(u, p, t)
        SVector{2}(u[2], -u[1])
    end

    u0 = SVector{2}(1.0, 0.0)
    tspan = (0.0, 4.0)
    close_tstops = [
        1.0, 1.0 + 1.0e-14, 1.0 + 2.0e-14, 1.0 + 5.0e-14,
        2.0, 2.0 + 1.0e-15, 2.0 + 1.0e-14,
        3.0, 3.0 + 1.0e-13,
    ]

    prob = ODEProblem(oscillator, u0, tspan)
    sol = solve(prob, Vern9(); tstops = close_tstops, reltol = 1.0e-12, abstol = 1.0e-15)
    @test successful_retcode(sol)
    for time in [1.0, 2.0, 3.0]
        @test any(abs.(sol.t .- time) .< 1.0e-10)
    end
end

# Tests for the `t > t_d` convention on `d_discontinuities`:
# v7 advances `integrator.t` by one ULP past a d_discontinuity so the
# post-discontinuity side of `t`-dependent branches in `f` is reached.
@testset "d_discontinuities: interior discontinuity with t > t_d" begin
    # f has slope 0 for t < 5 and slope 1 for t > 5.
    # Exact solution: u(t) = 0 for t ≤ 5, u(t) = t - 5 for t > 5. So u(10) = 5.
    # With a fixed-step non-adaptive integrator, a stale pre-discontinuity FSAL
    # would cause the first post-discontinuity step to use f(u, 5) = 0 and drop
    # accumulated value by exactly one step (0.5 with dt = 0.5).
    f(u, p, t) = t > 5.0 ? 1.0 : 0.0
    prob = ODEProblem(f, 0.0, (0.0, 10.0))
    sol = solve(
        prob, Euler(); dt = 0.5, d_discontinuities = [5.0],
        adaptive = false
    )
    @test sol.u[end] ≈ 5.0 atol = 1.0e-10

    # Same setup with Tsit5 (FSAL, adaptive) at tight tolerance.
    sol_tsit5 = solve(
        prob, Tsit5(); d_discontinuities = [5.0],
        reltol = 1.0e-12, abstol = 1.0e-14
    )
    @test sol_tsit5.u[end] ≈ 5.0 atol = 1.0e-10
    @test 5.0 ∈ sol_tsit5.t
end

@testset "d_discontinuities: starting-time discontinuity" begin
    # f has slope 0 at t = 0 and slope 1 for t > 0. Without handling of
    # t_d == t0, the starting-time discontinuity would be silently dropped
    # (filtered out of the tstops heap by the strict-inequality tspan check)
    # and fsalfirst would be stuck at 0.
    f(u, p, t) = t > 0.0 ? 1.0 : 0.0
    prob = ODEProblem(f, 0.0, (0.0, 5.0))
    sol = solve(
        prob, Euler(); dt = 0.5, d_discontinuities = [0.0],
        adaptive = false
    )
    @test sol.u[end] ≈ 5.0 atol = 1.0e-10
end

@testset "d_discontinuities: backward integration shifts with prevfloat" begin
    # Backwards integration with discontinuity at t = 5.
    # u(10) = 5. For t > 5 the slope is 1, so integrating backward from t = 10
    # to t = 5 decreases u by 5 → u(5) = 0. For t ≤ 5, slope is 0, so u stays
    # at 0 back to t = 0. Expected u(0) = 0.
    f(u, p, t) = t > 5.0 ? 1.0 : 0.0
    prob = ODEProblem(f, 5.0, (10.0, 0.0))
    sol = solve(
        prob, Tsit5(); d_discontinuities = [5.0],
        reltol = 1.0e-12, abstol = 1.0e-14
    )
    @test sol.u[end] ≈ 0.0 atol = 1.0e-10
end

@testset "Fixed-dt: no spurious tiny final step from accumulated drift" begin
    # Regression test: with fixed dt = 0.1 and tspan = (0.0, 1.0) the
    # accumulated `t + dt + dt + ...` drifts and used to produce a spurious
    # 12th step at 0.9999999999999999 followed by 1.0. A tolerance in
    # modify_dt_for_tstops! ensures the final tstop is hit cleanly.
    f!(du, u, p, t) = (du .= u; nothing)
    u0 = [1.0]
    prob = ODEProblem(f!, u0, (0.0, 1.0))
    sol = solve(prob, Euler(); dt = 0.1)
    @test length(sol.t) == 11
    @test sol.t[end] == 1.0

    # Reverse direction
    prob_rev = ODEProblem(f!, u0, (1.0, 0.0))
    sol_rev = solve(prob_rev, Euler(); dt = -0.1)
    @test length(sol_rev.t) == 11
    @test sol_rev.t[end] == 0.0

    # dt that does not evenly divide tspan still hits tspan[end] cleanly
    sol_uneven = solve(prob, Euler(); dt = 0.3)
    @test sol_uneven.t[end] == 1.0
end

@testset "d_discontinuities: tprev advanced past discontinuity" begin
    # Directly verify the shift-past invariant using the integrator interface:
    # after stepping past the discontinuity, integrator.tprev should be
    # nextfloat(t_d), not t_d itself.
    f(u, p, t) = t > 5.0 ? 1.0 : 0.0
    prob = ODEProblem(f, 0.0, (0.0, 10.0))
    integrator = init(prob, Tsit5(); d_discontinuities = [5.0])
    # Advance until we've landed on the tstop at t = 5.
    while integrator.t < 5.0
        step!(integrator)
    end
    @test integrator.t == 5.0
    # The next step triggers update_fsal! → shift_past_discontinuity!
    step!(integrator)
    @test integrator.tprev == nextfloat(5.0)
end
