using Test
using Sbplib.SimpleTimeSteppers


@testset "rk4" begin
    ode_solutions = @NamedTuple{F,v}.([
        ((v,t) ->  0., (v₀,t₀,t) -> v₀),
        ((v,t) ->  1., (v₀,t₀,t) -> v₀+t-t₀),
        ((v,t) ->   t, (v₀,t₀,t) -> v₀+(t^2-t₀^2)/2),
        ((v,t) ->   v, (v₀,t₀,t) -> v₀*exp(t-t₀)),
        ((v,t) -> t*v, (v₀,t₀,t) -> v₀*exp((t^2-t₀^2)/2)),
    ])

    @testset "Exact" begin
        @testset for i ∈ eachindex(ode_solutions[1:3])
            F, v = ode_solutions[i]
            cases = @NamedTuple{v₀,t₀}.([
                (0.,0.),
                (0.,1.),
                (1.,0.),
                (1.,1.),
                (2.,0.),
                (0.,2.),
                (2.,2.),
            ])

            @testset for (v₀,t₀) ∈ cases
                @testset for Δt ∈ [0., 1., .5, 2.]
                    @test rk4(F, v₀, t₀, Δt) == v(v₀, t₀, t₀+Δt)
                end
            end
        end
    end

    @testset "Approximate" begin
        @testset for i ∈ eachindex(ode_solutions)
            F, v = ode_solutions[i]

            @testset for v₀ ∈ [1., .5] ,t₀ ∈ [0., 1.]
                Δt = 0.1
                @test rk4(F, v₀, t₀, Δt) ≈ v(v₀, t₀, t₀+Δt) rtol=1e-6
            end
        end
    end
end

@testset "rkn4" begin

    function damped_oscillator(v₀, v̇₀,t₀,t; ζ)
        ω = sqrt(1-ζ^2)

        C = v̇₀ + ζ*v₀

        V = v₀*cos(ω*(t-t₀)) + C*sin(ω*(t-t₀))
        V̇ = C*cos(ω*(t-t₀)) - v₀*sin(ω*(t-t₀))
        v =    exp(-ζ*(t-t₀))*V
        v̇ = -ζ*exp(-ζ*(t-t₀))*V + exp(-ζ*(t-t₀))*V̇
        return v, v̇
    end

    ode_solutions = @NamedTuple{F,vv̇}.([
        ((v,v̇,t) ->  0.,       (v₀,v̇₀,t₀,t) -> (v₀ + (t-t₀)*v̇₀,                 v̇₀)),
        ((v,v̇,t) ->  1.,       (v₀,v̇₀,t₀,t) -> (v₀ + (t-t₀)*v̇₀ + (t-t₀)^2/2,    v̇₀ + t - t₀)),
        ((v,v̇,t) ->   t,       (v₀,v̇₀,t₀,t) -> (v₀ + (t-t₀)*v̇₀ + (t-t₀)^3/6+ t₀/2*(t-t₀)^2,  v̇₀ + (t^2-t₀^2)/2)),
        ((v,v̇,t) ->  -v,       (v₀,v̇₀,t₀,t) -> (v₀*cos(t-t₀) + v̇₀*sin(t-t₀),   v̇₀*cos(t-t₀) - v₀*sin(t-t₀))),
        ((v,v̇,t) ->  -v-6/5*v̇, (v₀,v̇₀,t₀,t) -> damped_oscillator(v₀,v̇₀,t₀,t; ζ=3/5)),
    ])

    @testset "Exact" begin
        @testset for i ∈ eachindex(ode_solutions[1:3])
            F, vv̇ = ode_solutions[i]
            @testset for v₀ ∈ [0., 1., 2.], v̇₀ ∈ [0., 1.], t₀ ∈ [0.,1.,2.]
                @testset for Δt ∈ [0., 1., .5, 2.]
                    V,V̇ = rkn4(F, v₀, v̇₀, t₀, Δt)
                    v,v̇ = vv̇(v₀, v̇₀, t₀, t₀+Δt)
                    @test V ≈ v rtol=1e-15
                    @test V̇ ≈ v̇ rtol=1e-15
                end
            end
        end
    end

    @testset "Approximate" begin
        @testset for i ∈ eachindex(ode_solutions)
            F, vv̇ = ode_solutions[i]
            @testset for v̇₀ ∈ [0., 0.5], v₀ ∈ [1., 0.5], t₀∈[0.,1.]
                Δt = 0.1
                V,V̇ = rkn4(F, v₀, v̇₀, t₀, Δt)
                v,v̇ = vv̇(v₀, v̇₀, t₀, t₀+Δt)

                if i != 5
                    @test V ≈ v rtol=1e-6
                    @test V̇ ≈ v̇ rtol=1e-6
                else
                    # TODO: Is this a bug or reasonable?
                    @test_broken V ≈ v rtol=1e-6
                    @test_broken V̇ ≈ v̇ rtol=1e-6
                end
            end
        end
    end
end

@testset "cdiff" begin
    ode_solutions = @NamedTuple{F,vv̇}.([
        ((v,t) ->  0., (v₀,v̇₀,t₀,t) -> v₀ + (t-t₀)*v̇₀),
        ((v,t) ->  1., (v₀,v̇₀,t₀,t) -> v₀ + (t-t₀)*v̇₀ + (t-t₀)^2/2),
        ((v,t) ->   t, (v₀,v̇₀,t₀,t) -> v₀ + (t-t₀)*v̇₀ + (t-t₀)^3/6+ t₀/2*(t-t₀)^2),
        ((v,t) ->  -v, (v₀,v̇₀,t₀,t) -> v₀*cos(t-t₀) + v̇₀*sin(t-t₀)),
    ])

    @testset "Exact" begin
        @testset for i ∈ eachindex(ode_solutions[1:3])
            F, v = ode_solutions[i]
            @testset for v₀ ∈ [1., 2.], v̇₀ ∈ [0., 1.], t₀ ∈ [0.,1.,2.]
                @testset for Δt ∈ [0., 1., .5]
                    v₋₁ = v(v₀, v̇₀, t₀, t₀-Δt)
                    v₁  = v(v₀, v̇₀, t₀, t₀+Δt)

                    V₁, V₀ = centered_2nd_order_fd(F, v₀, v₋₁, t₀, Δt)

                    @test V₁ ≈ v₁ rtol=1e-15
                    @test V₀ ≈ v₀ rtol=1e-15
                end
            end
        end
    end

    @testset "Approximate" begin
        @testset for i ∈ eachindex(ode_solutions)
            F, v = ode_solutions[i]
            @testset for v̇₀ ∈ [0., 0.5], v₀ ∈ [1., 0.5], t₀ ∈ [0.,1.]
                Δt = 0.1
                v₋₁ = v(v₀, v̇₀, t₀, t₀-Δt)
                v₁  = v(v₀, v̇₀, t₀, t₀+Δt)

                V₁, V₀ = centered_2nd_order_fd(F, v₀, v₋₁, t₀, Δt)


                @test V₁ ≈ v₁ rtol=1e-5
                @test V₀ ≈ v₀ rtol=1e-5
            end
        end
    end
end