using StochasticDiffEq, DiffEqNoiseProcess, Test, Random
@testset "SDE Stratonovich Reversal Tests" begin
    Random.seed!(100)
    α = 1.01
    β = 0.87

    dt = 1.0e-3
    tspan = (0.0, 1.0)
    u₀ = 1 / 2

    tarray = collect(tspan[1]:dt:tspan[2])

    f!(du, u, p, t) = du .= α * u
    g!(du, u, p, t) = du .= β * u

    prob = SDEProblem(f!, g!, [u₀], tspan)
    sol = solve(prob, EulerHeun(), dt = dt, save_noise = true, adaptive = false)

    _sol = deepcopy(sol) # to make sure the plot is correct
    W1 = NoiseGrid(reverse!(_sol.t), reverse!(_sol.W.W))
    prob1 = SDEProblem(f!, g!, sol.u[end], reverse(tspan), noise = W1)
    sol1 = solve(prob1, EulerHeun(), dt = dt)

    _sol = deepcopy(sol)
    W2 = NoiseWrapper(_sol.W, reverse = true)
    prob2 = SDEProblem(f!, g!, sol.u[end], reverse(tspan), noise = W2)
    sol2 = solve(prob2, EulerHeun(), dt = dt, save_noise = false)

    # same time steps

    @test sol.u ≈ reverse(sol1.u) atol = 1.0e-2
    @test sol.u ≈ reverse(sol2.u) atol = 1.0e-2
    @test sol1.u ≈ sol2.u atol = 1.0e-5

    # 1/10 size time steps

    _sol = deepcopy(sol) # to make sure the plot is correct
    W1 = NoiseGrid(reverse!(_sol.t), reverse!(_sol.W.W))
    prob1 = SDEProblem(f!, g!, sol.u[end], reverse(tspan), noise = W1)
    sol1 = solve(prob1, EulerHeun(), dt = 0.1 * dt)

    _sol = deepcopy(sol)
    W2 = NoiseWrapper(_sol.W, reverse = true)
    prob2 = SDEProblem(f!, g!, sol.u[end], reverse(tspan), noise = W2)
    sol2 = solve(prob2, EulerHeun(), dt = 0.1 * dt, save_noise = false)

    @test sol.u ≈ sol1(tarray).u atol = 2.0e-2
    @test sol.u ≈ sol2(tarray).u atol = 2.0e-2

    # diagonal noise

    prob = SDEProblem(f!, g!, [u₀, u₀], tspan)
    sol = solve(prob, EulerHeun(), dt = dt, save_noise = true)

    _sol = deepcopy(sol) # to make sure the plot is correct
    W1 = NoiseGrid(reverse!(_sol.t), reverse!(_sol.W.W))
    prob1 = SDEProblem(f!, g!, sol.u[end], reverse(tspan), noise = W1)
    sol1 = solve(prob1, EulerHeun(), dt = dt)

    _sol = deepcopy(sol)
    W2 = NoiseWrapper(_sol.W, reverse = true)
    prob2 = SDEProblem(f!, g!, sol.u[end], reverse(tspan), noise = W2)
    sol2 = solve(prob2, EulerHeun(), dt = dt)

    @test sol.u ≈ reverse(sol1.u) atol = 5.0e-2
    @test sol.u ≈ reverse(sol2.u) atol = 5.0e-2
    @test sol1.u ≈ sol2.u atol = 1.0e-5

    # non-diagonal noise

    function gnd!(du, u, p, t)
        du[1, 1] = 0.3u[1]
        du[1, 2] = 0.6u[1]
        du[1, 3] = 0.9u[1]
        du[1, 4] = 0.12u[2]
        du[2, 1] = 1.2u[1]
        du[2, 2] = 0.2u[2]
        du[2, 3] = 0.3u[2]
        du[2, 4] = 1.8u[2]
    end
    prob = SDEProblem(f!, gnd!, [u₀, u₀], tspan, noise_rate_prototype = zeros(2, 4))
    sol = solve(prob, EulerHeun(), dt = dt, save_noise = true)

    _sol = deepcopy(sol) # to make sure the plot is correct
    W1 = NoiseGrid(reverse!(_sol.t), reverse!(_sol.W.W))
    prob1 = SDEProblem(
        f!, gnd!, sol.u[end], reverse(tspan), noise = W1,
        noise_rate_prototype = zeros(2, 4)
    )
    sol1 = solve(prob1, EulerHeun(), dt = dt)

    _sol = deepcopy(sol)
    W2 = NoiseWrapper(_sol.W, reverse = true)
    prob2 = SDEProblem(
        f!, gnd!, sol.u[end], reverse(tspan), noise = W2,
        noise_rate_prototype = zeros(2, 4)
    )
    sol2 = solve(prob2, EulerHeun(), dt = dt)

    @test sol.u ≈ reverse(sol1.u) atol = 5.0e-1
    @test sol.u ≈ reverse(sol2.u) atol = 5.0e-1
    @test sol1.u ≈ sol2.u atol = 1.0e-5

    ###
    ### OOP
    ###

    f(u, p, t) = α * u
    g(u, p, t) = β * u

    prob = SDEProblem(f, g, u₀, tspan)
    sol = solve(prob, EulerHeun(), dt = dt, save_noise = true)

    _sol = deepcopy(sol) # to make sure the plot is correct
    W1 = NoiseGrid(reverse!(_sol.t), reverse!(_sol.W.W))
    prob1 = SDEProblem(f, g, sol.u[end], reverse(tspan), noise = W1)
    sol1 = solve(prob1, EulerHeun(), dt = dt)

    _sol = deepcopy(sol)
    W2 = NoiseWrapper(_sol.W, reverse = true)
    prob2 = SDEProblem(f, g, sol.u[end], reverse(tspan), noise = W2)
    sol2 = solve(prob2, EulerHeun(), dt = dt)

    @test sol.u ≈ reverse(sol1.u) atol = 2.0e-2
    @test sol.u ≈ reverse(sol2.u) atol = 2.0e-2
    @test sol1.u ≈ sol2.u atol = 1.0e-5

    # diagonal noise

    prob = SDEProblem(f, g, [u₀, u₀], tspan)
    sol = solve(prob, EulerHeun(), dt = dt, save_noise = true)

    _sol = deepcopy(sol) # to make sure the plot is correct
    W1 = NoiseGrid(reverse!(_sol.t), reverse!(_sol.W.W))
    prob1 = SDEProblem(f, g, sol.u[end], reverse(tspan), noise = W1)
    sol1 = solve(prob1, EulerHeun(), dt = dt)

    _sol = deepcopy(sol)
    W2 = NoiseWrapper(_sol.W, reverse = true)
    prob2 = SDEProblem(f, g, sol.u[end], reverse(tspan), noise = W2)
    sol2 = solve(prob2, EulerHeun(), dt = dt)

    @test sol.u ≈ reverse(sol1.u) atol = 5.0e-2
    @test sol.u ≈ reverse(sol2.u) atol = 5.0e-2
    @test sol1.u ≈ sol2.u atol = 1.0e-5
end

@testset "Reverse a given NoiseProcess " begin
    # Noise Wrapper
    _W = WienerProcess(0.0, 0.0, 0.0)

    dt = 0.1
    calculate_step!(_W, dt, nothing, nothing)

    for i in 1:100
        accept_step!(_W, dt, nothing, nothing)
    end

    W1 = NoiseWrapper(_W, reverse = true)
    dt = -0.1
    calculate_step!(W1, dt, nothing, nothing)
    for i in 1:99
        accept_step!(W1, dt, nothing, nothing)
    end

    W2 = reverse(_W)
    dt = -0.1
    calculate_step!(W2, dt, nothing, nothing)
    for i in 1:99
        accept_step!(W2, dt, nothing, nothing)
    end

    @test isapprox(_W.W[2:end], reverse(W1.W), atol = 1.0e-12)
    @test isapprox(_W.W[2:end], reverse(W2.W), atol = 1.0e-12)
    @test isapprox(_W.Z[2:end], reverse(W1.Z), atol = 1.0e-12)
    @test isapprox(_W.Z[2:end], reverse(W2.Z), atol = 1.0e-12)

    # Noise Grid
    dt = 0.001
    t = 0:dt:1
    brownian_values = cumsum([0; [sqrt(dt) * randn() for i in 1:(length(t) - 1)]])
    brownian_values2 = cumsum([0; [sqrt(dt) * randn() for i in 1:(length(t) - 1)]])
    _W = NoiseGrid(t, brownian_values, brownian_values2)

    prob = NoiseProblem(_W, (0.0, 1.0))
    sol = solve(prob; dt = dt)

    W1 = NoiseGrid(reverse(sol.t), reverse(sol.W), reverse(sol.Z))
    prob = NoiseProblem(W1, (1.0, 0.0))
    sol1 = solve(prob; dt = -dt)

    W2 = reverse(_W)
    prob = NoiseProblem(W2, (1.0, 0.0))
    sol2 = solve(prob; dt = -dt)

    @test isapprox(sol.W, reverse(sol1.W), atol = 1.0e-12)
    @test isapprox(sol.W, reverse(sol2.W), atol = 1.0e-12)
    @test isapprox(sol.Z, reverse(sol1.Z), atol = 1.0e-12)
    @test isapprox(sol.Z, reverse(sol2.Z), atol = 1.0e-12)
end

@testset "SDE Ito Basic Reversal Tests" begin
    n = 100000
    T = 2.0
    dt = T / n
    x0 = 0.3

    b(u, p, t) = sin(t) + cos(u)
    σ(u, p, t) = pi + atan(u)

    seed = 10
    Random.seed!(seed)
    W = [0.0; cumsum(sqrt(dt) * randn(n + 1))]
    ts = collect(0.0:dt:(T + dt))
    ts2 = ts[1:(end - 1)]
    p = nothing

    """
    For ito integrals
    """

    x = x0 # starting point
    xs = [x]
    t = 0.0
    for i in 1:n
        t, x
        x += b(x, p, t) * dt + σ(x, p, t) * (W[i + 1] - W[i]) # this is an ito integral
        t += dt
        push!(xs, x)
    end

    z = xs[end]

    # Now reverse... t

    dσ_dx(u, p, t) = 1 / (1 + u^2) # d(arctan(x))/dx
    z = xs[end]
    @show "starting point" z
    ys = [z]
    t = T
    for i in n:-1:1
        t, z
        cor = 1 / 2 * dσ_dx(z, p, t) * σ(z, p, t)
        z -= (b(z, p, t) - 0 * 2 * cor) * dt + σ(z, p, t) * (W[i + 1] - W[i]) # reverse ito integral
        t -= dt
        push!(ys, z)
    end

    W1 = NoiseGrid(ts, W)
    prob1 = SDEProblem(b, σ, x0, (0.0, 2.0), noise = W1)
    sol1 = solve(prob1, EM(false), dt = dt, adaptive = false)

    @test isapprox(xs, sol1(ts2).u, rtol = 1.0e-8)

    W1rev = NoiseGrid(reverse(ts), reverse(W))
    prob1 = SDEProblem(b, σ, sol1.u[end], (sol1.t[end], sol1.t[1]), noise = W1rev)
    sol2 = solve(prob1, EM(false), dt = dt, adaptive = false)

    @test isapprox(ys, sol2(reverse(ts2)).u, atol = 1.0e-3)
    @test !isapprox(sol1.u, reverse(sol2.u), atol = 1.0e-0)

    bwrong(u, p, t) = b(u, p, t) - 1 // 2 * dσ_dx(u, p, t) * σ(u, p, t)
    W1rev = NoiseGrid(reverse(ts), reverse(W))
    prob1 = SDEProblem(bwrong, σ, sol1.u[end], (sol1.t[end], sol1.t[1]), noise = W1rev)
    sol2 = solve(prob1, EM(false), dt = dt, adaptive = false)

    @test !isapprox(ys, sol2(reverse(ts2)).u, atol = 1.0e-0)
    @test !isapprox(sol1.u, reverse(sol2.u), atol = 1.0e-0)

    bcorrected(u, p, t) = b(u, p, t) - 2 * 1 // 2 * dσ_dx(u, p, t) * σ(u, p, t)
    W1rev = NoiseGrid(reverse(ts), reverse(W))
    prob1 = SDEProblem(bcorrected, σ, sol1.u[end], (sol1.t[end], sol1.t[1]), noise = W1rev)
    sol2 = solve(prob1, EM(false), dt = dt, adaptive = false)

    #plot(ts, sol1.u - reverse(sol2.u))
    @test !isapprox(ys, sol2(reverse(ts2)).u, atol = 1.0e-6)
    @test isapprox(sol1(ts2).u, sol2(ts2).u, atol = 1.0e-0)
end

"""
Some more Ito reversals
"""

@testset "SDE Ito Reversal Tests" begin
    Random.seed!(100)
    α = 1.0
    β = 0.3

    dt = 1.0e-3
    tspan = (0.0, 1.0)
    u₀ = 1 / 2

    tarray = collect(tspan[1]:dt:tspan[2])

    f!(du, u, p, t) = du .= α * u
    g!(du, u, p, t) = du .= β * u

    prob = SDEProblem(f!, g!, [u₀], tspan)
    sol = solve(prob, EM(), dt = dt, save_noise = true, adaptive = false)

    fcorrected!(du, u, p, t) = du .= (α - β^2) * u

    _sol = deepcopy(sol) # to make sure the plot is correct
    W1 = NoiseGrid(
        reverse(_sol.t),
        reverse(_sol.W.W)        # ,reverse(_sol.W.Z)
    )
    prob1 = SDEProblem(fcorrected!, g!, sol.u[end], reverse(tspan), noise = W1)
    sol1 = solve(prob1, EM(), dt = dt, adaptive = false)

    _sol = deepcopy(sol)
    W2 = NoiseWrapper(_sol.W, reverse = true)
    prob2 = SDEProblem(fcorrected!, g!, sol.u[end], reverse(tspan), noise = W2)
    sol2 = solve(prob2, EM(), dt = dt, save_noise = false, adaptive = false)

    # same time steps

    @test sol.u ≈ reverse(sol1.u) rtol = 1.0e-2
    @test sol.u ≈ reverse(sol2.u) rtol = 1.0e-2
    @test sol1.u ≈ sol2.u atol = 1.0e-8

    f(u, p, t) = α * u
    g(u, p, t) = β * u

    prob = SDEProblem(f, g, u₀, tspan)
    sol = solve(prob, EM(), dt = dt, save_noise = true, adaptive = false)

    fcorrected(u, p, t) = (α - β^2) * u

    _sol = deepcopy(sol) # to make sure the plot is correct
    W1 = NoiseGrid(
        reverse(_sol.t),
        reverse(_sol.W.W)        # ,reverse(_sol.W.Z)
    )
    prob1 = SDEProblem(fcorrected, g, sol.u[end], reverse(tspan), noise = W1)
    sol1 = solve(prob1, EM(), dt = dt, adaptive = false)

    _sol = deepcopy(sol)
    W2 = NoiseWrapper(_sol.W, reverse = true)
    prob2 = SDEProblem(fcorrected, g, sol.u[end], reverse(tspan), noise = W2)
    sol2 = solve(prob2, EM(), dt = dt, adaptive = false)

    @test sol.u ≈ reverse(sol1.u) rtol = 1.0e-2
    @test sol.u ≈ reverse(sol2.u) rtol = 1.0e-2
    @test sol1.u ≈ sol2.u atol = 1.0e-7

    # using Plots; plt = plot(sol)
    # plot!(reverse(sol1.t),reverse(sol1[1,:]))
    # plot(vcat(sol.u - reverse(sol1.u) ...))
end

@testset "SDE Ito additive noise Reversal Tests" begin
    using SDEProblemLibrary: prob_sde_additive, prob_sde_additivesystem
    Random.seed!(100)
    dt = 1.0e-3

    @testset "OOP 1D" begin
        # OOP 1D
        prob = prob_sde_additive

        # non-adaptive

        ### SOSRA2
        sol = solve(prob, SOSRA2(), dt = dt, save_noise = true, adaptive = false)
        W1 = reverse(sol.W)

        prob1 = SDEProblem(
            prob.f, prob.g, sol.u[end], reverse(prob.tspan), prob.p,
            noise = W1
        )
        sol1 = solve(prob1, SOSRA2(), dt = dt, adaptive = false)

        @test sol.u ≈ reverse(sol1.u) rtol = 1.0e-4
        @test sol.t ≈ reverse(sol1.t) rtol = 1.0e-4

        ### SOSRI
        sol = solve(prob, SOSRI(), dt = dt, save_noise = true, adaptive = false)
        W1 = reverse(sol.W)

        prob1 = SDEProblem(
            prob.f, prob.g, sol.u[end], reverse(prob.tspan), prob.p,
            noise = W1
        )
        sol1 = solve(prob1, SOSRI(), dt = dt, adaptive = false)

        @test sol.u ≈ reverse(sol1.u) rtol = 1.0e-4
        @test sol.t ≈ reverse(sol1.t) rtol = 1.0e-4

        # adaptive

        ### SOSRA2
        sol = solve(prob, SOSRA2(), dt = dt, save_noise = true, adaptive = true)
        W1 = reverse(sol.W)

        prob1 = SDEProblem(
            prob.f, prob.g, sol.u[end], reverse(prob.tspan), prob.p,
            noise = W1
        )
        sol1 = solve(prob1, SOSRA2(), dt = dt, adaptive = true)

        ts = prob.tspan[1]:0.1:prob.tspan[2]
        @test sol(ts) ≈ sol1(ts) rtol = 1.0e-3
        @test length(sol.t) != length(sol1.t)

        ### SOSRI
        sol = solve(prob, SOSRI(), dt = dt, save_noise = true, adaptive = true)
        W1 = reverse(sol.W)

        prob1 = SDEProblem(
            prob.f, prob.g, sol.u[end], reverse(prob.tspan), prob.p,
            noise = W1
        )
        sol1 = solve(prob1, SOSRI(), dt = dt, adaptive = true)

        ts = prob.tspan[1]:0.1:prob.tspan[2]
        @test sol(ts) ≈ sol1(ts) rtol = 1.0e-3
        @test length(sol.t) != length(sol1.t)
    end

    @testset "iip 4D" begin
        # iip 4D
        prob = prob_sde_additivesystem

        # non-adaptive

        ### SOSRA2
        sol = solve(prob, SOSRA2(), dt = dt, save_noise = true, adaptive = false)
        W1 = reverse(sol.W)

        prob1 = SDEProblem(
            prob.f, prob.g, sol.u[end], reverse(prob.tspan), prob.p,
            noise = W1
        )
        sol1 = solve(prob1, SOSRA2(), dt = dt, adaptive = false)

        @test sol.u ≈ reverse(sol1.u) rtol = 1.0e-3
        @test sol.t ≈ reverse(sol1.t) rtol = 1.0e-3

        ### SOSRI
        sol = solve(prob, SOSRI(), dt = dt, save_noise = true, adaptive = false)
        W1 = reverse(sol.W)

        prob1 = SDEProblem(
            prob.f, prob.g, sol.u[end], reverse(prob.tspan), prob.p,
            noise = W1
        )
        sol1 = solve(prob1, SOSRI(), dt = dt, adaptive = false)

        @test sol.u ≈ reverse(sol1.u) rtol = 1.0e-3
        @test sol.t ≈ reverse(sol1.t) rtol = 1.0e-3

        # adaptive

        abstol = 1.0e-5
        reltol = 1.0e-8

        ### SOSRA2
        sol = solve(
            prob, SOSRA2(), dt = dt, save_noise = true, adaptive = true,
            abstol = abstol, reltol = reltol
        )
        W1 = reverse(sol.W)

        prob1 = SDEProblem(
            prob.f, prob.g, sol.u[end], reverse(prob.tspan), prob.p,
            noise = W1
        )
        sol1 = solve(
            prob1, SOSRA2(), dt = dt, adaptive = true, abstol = abstol,
            reltol = reltol
        )

        ts = prob.tspan[1]:0.1:prob.tspan[2]
        @test sol(ts) ≈ sol1(ts) rtol = 1.0e-3
        @test length(sol.t) != length(sol1.t)

        ### SOSRI
        sol = solve(
            prob, SOSRI(), dt = dt, save_noise = true, adaptive = true,
            abstol = abstol, reltol = reltol
        )
        W1 = reverse(sol.W)

        prob1 = SDEProblem(
            prob.f, prob.g, sol.u[end], reverse(prob.tspan), prob.p,
            noise = W1
        )
        sol1 = solve(
            prob1, SOSRI(), dt = dt, adaptive = true, abstol = abstol,
            reltol = reltol
        )

        ts = prob.tspan[1]:0.1:prob.tspan[2]
        @test sol(ts) ≈ sol1(ts) rtol = 1.0e-3
        @test length(sol.t) != length(sol1.t)
    end
end
