using OrdinaryDiffEqExponentialRK, Test, DiffEqDevTools, Random, LinearAlgebra, LinearSolve
using OrdinaryDiffEqVerner, OrdinaryDiffEqSDIRK
using OrdinaryDiffEqCore: alg_order
using SciMLOperators: MatrixOperator, ScalarOperator

@testset "Caching Out-of-place" begin
    println("Caching Out-of-place")
    μ = 1.01
    linnonlin_f2 = (u, p, t) -> μ * u
    linnonlin_f1 = ScalarOperator(μ)
    linnonlin_fun = SplitFunction(
        linnonlin_f1, linnonlin_f2;
        analytic = (u0, p, t) -> u0 .* exp.(2μ * t)
    )
    prob = SplitODEProblem(linnonlin_fun, 1 / 2, (0.0, 1.0))

    Random.seed!(100)
    dts = 1 ./ 2 .^ (7:-1:4) #14->7 good plot
    for Alg in [
            LawsonEuler,
            NorsettEuler,
            ETDRK2,
            ETDRK3,
            ETDRK4,
            HochOst4,
            ETD2,
            KenCarp3,
            CFNLIRK3,
        ]
        sim = test_convergence(dts, prob, Alg())
        @test sim.𝒪est[:l2] ≈ alg_order(Alg()) atol = 0.2
    end

    # Dense test
    sim = test_convergence(dts, prob, ETDRK4(), dense_errors = true)
    @test abs(sim.𝒪est[:l2] - 4) < 0.2
end

@testset "Caching Inplace" begin
    println("Caching Inplace")
    μ = 1.01
    u0 = rand(2)
    A = [2.0 -1.0; -1.0 2.0]
    linnonlin_f1 = MatrixOperator(A)
    linnonlin_f2 = (du, u, p, t) -> du .= μ .* u
    linnonlin_fun_iip = SplitFunction(
        linnonlin_f1, linnonlin_f2;
        analytic = (u0, p, t) -> exp((A + μ * I) * t) * u0
    )
    prob = SplitODEProblem(linnonlin_fun_iip, u0, (0.0, 1.0))

    dts = 1 ./ 2 .^ (8:-1:4)
    for Alg in [
            LawsonEuler(),
            NorsettEuler(),
            ETDRK2(),
            ETDRK3(),
            ETDRK4(),
            HochOst4(),
            ETD2(),
        ]
        sim = test_convergence(dts, prob, Alg)
        @test sim.𝒪est[:l2] ≈ alg_order(Alg) atol = 0.15
    end

    dts = 1 ./ 2 .^ (13:-1:9)
    Alg = KenCarp3(linsolve = LinearSolve.KrylovJL_GMRES())
    sim = test_convergence(dts, prob, Alg, reltol = 1.0e-16)
    @test sim.𝒪est[:l2] ≈ alg_order(Alg) atol = 0.5

    dts = 1 ./ 2 .^ (8:-1:4)
    sim = test_convergence(dts, prob, ETDRK4(), dense_errors = true)
    @test sim.𝒪est[:l2] ≈ 4 atol = 0.1
    @test sim.𝒪est[:L2] ≈ 4 atol = 0.1
end

@info "CFNLIRK3() is broken"

@testset "EPIRK Out-of-place" begin
    println("EPIRK Out-of-place")
    # Setup nonlinear problem
    A = [-2.0 1.0; 1.0 -2.0]
    f = (u, p, t) -> A * u - u .^ 3
    jac = (u, p, t) -> A - [3u[1]^2 0.0; 0.0 3u[2]^2]
    fun = ODEFunction(f; jac = jac)
    Random.seed!(0)
    u0 = rand(2)
    tspan = (0.0, 1.0)
    prob = ODEProblem(fun, u0, tspan)
    # Setup approximate solution
    test_setup = Dict(:alg => Vern9(), :reltol => 1.0e-16, :abstol => 1.0e-16)
    # Convergence simulation
    dts = 1 ./ 2 .^ (7:-1:4)
    Algs = [Exp4, EPIRK4s3A, EPIRK4s3B, EPIRK5s3, EXPRB53s3, EPIRK5P1, EPIRK5P2]
    for Alg in Algs
        sim = analyticless_test_convergence(
            dts, prob, Alg(adaptive_krylov = false),
            test_setup
        )
        if Alg == EPIRK5s3
            @test_broken sim.𝒪est[:l2] ≈ alg_order(Alg()) atol = 0.1
        else
            @test sim.𝒪est[:l2] ≈ alg_order(Alg()) atol = 0.1
        end
    end
end

@testset "EPIRK Inplace" begin
    println("EPIRK Inplace")
    # Setup nonlinear problem
    A = [-2.0 1.0; 1.0 -2.0]
    f = (du, u, p, t) -> (mul!(du, A, u); du .-= u .^ 3)
    jac_update! = (J, u, p, t) -> (copyto!(J, A); J[1, 1] -= 3u[1]^2; J[2, 2] -= 3u[2]^2)
    jac_prototype = MatrixOperator(zeros(2, 2); update_func! = jac_update!)
    fun = ODEFunction(f; jac_prototype = jac_prototype)
    Random.seed!(0)
    u0 = rand(2)
    tspan = (0.0, 1.0)
    prob = ODEProblem(fun, u0, tspan)
    # Setup approximate solution
    test_setup = Dict(:alg => Vern9(), :reltol => 1.0e-16, :abstol => 1.0e-16)
    # Convergence simulation
    dts = 1 ./ 2 .^ (7:-1:4)
    Algs = [Exp4, EPIRK4s3A, EPIRK4s3B, EPIRK5s3, EXPRB53s3, EPIRK5P1, EPIRK5P2]
    for Alg in Algs
        sim = analyticless_test_convergence(
            dts, prob, Alg(adaptive_krylov = false),
            test_setup
        )
        if Alg == EPIRK5s3
            @test_broken sim.𝒪est[:l2] ≈ alg_order(Alg()) atol = 0.1
        else
            @test sim.𝒪est[:l2] ≈ alg_order(Alg()) atol = 0.1
        end
    end
end

@testset "ExpRK Autodiff" begin
    println("ExpRK Autodiff")
    # Setup nonlinear problem (without explicit jacobian)
    A = [-2.0 1.0; 1.0 -2.0]
    f = (u, p, t) -> A * u - u .^ 3
    f_ip = (du, u, p, t) -> (mul!(du, A, u); du .-= u .^ 3)
    fun = ODEFunction(f)
    fun_ip = ODEFunction(f_ip)
    Random.seed!(0)
    u0 = rand(2)
    tspan = (0.0, 1.0)
    prob = ODEProblem(fun, u0, tspan)
    prob_ip = ODEProblem(fun_ip, u0, tspan)
    # Setup approximate solution
    test_setup = Dict(:alg => Vern9(), :reltol => 1.0e-16, :abstol => 1.0e-16)
    # Convergence simulation
    dts = 1 ./ 2 .^ (7:-1:4)
    sim = analyticless_test_convergence(dts, prob, HochOst4(krylov = true), test_setup)
    @test sim.𝒪est[:l2] ≈ 4 atol = 0.1
    sim = analyticless_test_convergence(dts, prob_ip, HochOst4(krylov = true), test_setup)
    @test sim.𝒪est[:l2] ≈ 4 atol = 0.1
    sim = analyticless_test_convergence(
        dts, prob, EPIRK5P1(adaptive_krylov = false),
        test_setup
    )
    @test sim.𝒪est[:l2] ≈ 5 atol = 0.1
    sim = analyticless_test_convergence(
        dts, prob_ip, EPIRK5P1(adaptive_krylov = false),
        test_setup
    )
    @test sim.𝒪est[:l2] ≈ 5 atol = 0.1
end

@testset "Adaptive Exprb Out-of-place" begin
    println("Adaptive Exprb Out-of-place")
    # Setup nonlinear problem
    A = [-2.0 1.0; 1.0 -2.0]
    f = (u, p, t) -> A * u - u .^ 3
    jac = (u, p, t) -> A - [3u[1]^2 0.0; 0.0 3u[2]^2]
    fun = ODEFunction(f; jac = jac)
    Random.seed!(0)
    u0 = rand(2)
    tspan = (0.0, 1.0)
    prob = ODEProblem(fun, u0, tspan)
    # Setup approximate solution
    test_setup = Dict(:alg => Vern9(), :reltol => 1.0e-16, :abstol => 1.0e-16)
    # Convergence simulation
    dts = 1 ./ 2 .^ (7:-1:4)
    Algs = [Exprb32, Exprb43]
    for Alg in Algs
        sim = analyticless_test_convergence(dts, prob, Alg(), test_setup)
        @test sim.𝒪est[:l2] ≈ alg_order(Alg()) atol = 0.1
    end
end

@testset "Adaptive Exprb Inplace" begin
    println("Adaptive Exprb Inplace")
    # Setup nonlinear problem
    A = [-2.0 1.0; 1.0 -2.0]
    f = (du, u, p, t) -> (mul!(du, A, u); du .-= u .^ 3)
    jac_update! = (J, u, p, t) -> (copyto!(J, A); J[1, 1] -= 3u[1]^2; J[2, 2] -= 3u[2]^2)
    jac_prototype = MatrixOperator(zeros(2, 2); update_func! = jac_update!)
    fun = ODEFunction(f; jac_prototype = jac_prototype)
    Random.seed!(0)
    u0 = rand(2)
    tspan = (0.0, 1.0)
    prob = ODEProblem(fun, u0, tspan)
    # Setup approximate solution
    test_setup = Dict(:alg => Vern9(), :reltol => 1.0e-16, :abstol => 1.0e-16)
    # Convergence simulation
    dts = 1 ./ 2 .^ (7:-1:4)
    Algs = [Exprb32, Exprb43]
    for Alg in Algs
        sim = analyticless_test_convergence(dts, prob, Alg(), test_setup)
        @test sim.𝒪est[:l2] ≈ alg_order(Alg()) atol = 0.1
    end
end
