using OrdinaryDiffEqSDIRK, OrdinaryDiffEqBDF, OrdinaryDiffEqRosenbrock, OrdinaryDiffEqFIRK,
    OrdinaryDiffEqNonlinearSolve, Test, LinearAlgebra, Statistics
using OrdinaryDiffEqCore
using OrdinaryDiffEqNonlinearSolve: NLFunctional, NLAnderson, NLNewton
using SciMLOperators: MatrixOperator, update_coefficients!
using LinearAlgebra: Diagonal
using ADTypes: AutoForwardDiff

# create mass matrix problems
function make_mm_probs(mm_A, ::Val{iip}) where {iip}
    # iip
    function mm_f(du, u, p, t)
        update_coefficients!(mm_A, OrdinaryDiffEqCore.constvalue.(u), p, t)
        mm_b = vec(sum(mm_A; dims = 2))
        mul!(du, mm_A, u)
        du .+= t * mm_b
        return nothing
    end
    mm_g(du, u, p, t) = (@. du = u + t; nothing)

    # oop
    mm_f(u, p, t) = (
        update_coefficients!(mm_A, OrdinaryDiffEqCore.constvalue.(u), p, t);
        mm_A * (u .+ t)
    )
    mm_g(u, p, t) = u .+ t

    mm_analytic(u0, p, t) = @. 2 * u0 * exp(t) - t - 1

    u0 = ones(3)
    tspan = (0.0, 1.0)

    prob = ODEProblem(
        ODEFunction{iip, true}(
            mm_f, analytic = mm_analytic,
            mass_matrix = mm_A
        ), u0, tspan
    )
    prob2 = ODEProblem(ODEFunction{iip, true}(mm_g, analytic = mm_analytic), u0, tspan)

    return prob, prob2
end

function _norm_dsol(alg, prob, prob2, dt = 1 / 10)
    sol = solve(prob, alg, dt = dt, adaptive = false)
    sol2 = solve(prob2, alg, dt = dt, adaptive = false)
    return norm(sol .- sol2)
end

function update_func1!(A, u, p, t)
    A[1, 1] = cos(t)
    A[2, 1] = sin(t) * u[1]
    A[3, 1] = t^2
    A[1, 2] = cos(t) * sin(t)
    A[2, 2] = cos(t)^2 + u[3]
    A[3, 2] = sin(t) * u[2]
    A[1, 3] = sin(t)
    A[2, 3] = t^2
    return A[3, 3] = t * cos(t) + 1
end

function update_func1(A, u, p, t)
    A = copy(A)
    update_func1!(A, u, p, t)
    return A
end

function update_func2!(A, u, p, t)
    A[1, 1] = cos(t)
    A[2, 1] = sin(t)
    A[3, 1] = t^2
    A[1, 2] = cos(t) * sin(t)
    A[2, 2] = cos(t)^2
    A[3, 2] = sin(t)
    A[1, 3] = sin(t)
    A[2, 3] = t^2
    return A[3, 3] = t * cos(t) + 1
end

function update_func2(A, u, p, t)
    A = copy(A)
    update_func2!(A, u, p, t)
    return A
end

almost_I = Matrix{Float64}(1.01I, 3, 3)
mm_A = Float64[-2 1 4; 4 -2 1; 2 1 3]
dependent_M1 = MatrixOperator(
    ones(3, 3), update_func = update_func1,
    update_func! = update_func1!
)
dependent_M2 = MatrixOperator(
    ones(3, 3), update_func = update_func2,
    update_func! = update_func2!
)

@testset "Mass Matrix Accuracy Tests" for mm in (almost_I, mm_A)
    # test each method for exactness
    for iip in (false, true)
        prob, prob2 = make_mm_probs(mm, Val(iip))

        println("SDIRKs")
        @test _norm_dsol(ImplicitEuler(), prob, prob2) ≈ 0 atol = 1.0e-12
        @test _norm_dsol(Trapezoid(), prob, prob2) ≈ 0 atol = 1.0e-12
        @test _norm_dsol(RadauIIA5(), prob, prob2) ≈ 0 atol = 1.0e-12
        @test _norm_dsol(ImplicitMidpoint(extrapolant = :constant), prob, prob2) ≈ 0 atol = 1.0e-10

        println("BDFs")
        @test _norm_dsol(ABDF2(), prob, prob2) ≈ 0 atol = 1.0e-12

        @test _norm_dsol(QBDF1(), prob, prob2) ≈ 0 atol = 1.0e-12
        @test _norm_dsol(QBDF2(), prob, prob2) ≈ 0 atol = 1.0e-12
        @test _norm_dsol(QBDF(), prob, prob2) ≈ 0 atol = 1.0e-12
        @test _norm_dsol(FBDF(), prob, prob2) ≈ 0 atol = 1.0e-12

        @test _norm_dsol(QNDF1(), prob, prob2) ≈ 0 atol = 1.0e-12
        @test _norm_dsol(QNDF2(), prob, prob2) ≈ 0 atol = 1.0e-12
        @test _norm_dsol(QNDF(), prob, prob2) ≈ 0 atol = 1.0e-12

        println("Rosenbrocks")
        if mm == almost_I
            @test _norm_dsol(Rosenbrock23(), prob, prob2) ≈ 0 atol = 1.0e-11
            @test _norm_dsol(Rosenbrock32(), prob, prob2) ≈ 0 atol = 1.0e-11
        end
        @test _norm_dsol(ROS3P(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(Rodas3(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(ROS2(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(ROS2PR(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(ROS2S(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(ROS3(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(ROS3PR(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(Scholz4_7(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(RosShamp4(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(Veldd4(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(Velds4(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(GRK4T(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(GRK4A(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(Ros4LStab(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(RosenbrockW6S4OS(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(ROS34PW1a(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(ROS34PW1b(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(ROS34PW2(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(ROS34PW3(), prob, prob2) ≈ 0 atol = 1.0e-10
        @test _norm_dsol(ROS34PRw(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(ROS3PRL(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(ROS3PRL2(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(ROK4a(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(Rodas4(), prob, prob2) ≈ 0 atol = 1.0e-9
        @test _norm_dsol(Rodas42(), prob, prob2) ≈ 0 atol = 1.0e-9
        @test _norm_dsol(Rodas4P(), prob, prob2) ≈ 0 atol = 1.0e-9
        @test _norm_dsol(Rodas5(), prob, prob2) ≈ 0 atol = 1.0e-7
        @test _norm_dsol(Rodas23W(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(Rodas3P(), prob, prob2) ≈ 0 atol = 1.0e-11
        @test _norm_dsol(Rodas5P(), prob, prob2) ≈ 0 atol = 1.0e-7
        @test _norm_dsol(Rodas5Pe(), prob, prob2) ≈ 0 atol = 1.0e-7
        @test _norm_dsol(Rodas5Pr(), prob, prob2) ≈ 0 atol = 1.0e-7
    end
end

@testset "Mass Matrix Functional Iteration Tests" for mm in (almost_I,) # Requires contraction map
    for iip in (false, true)
        @show iip
        prob, prob2 = make_mm_probs(mm, Val(iip))

        sol = solve(
            prob, ImplicitEuler(nlsolve = NLFunctional()), dt = 1 / 10,
            adaptive = false, reltol = 1.0e-7, abstol = 1.0e-10
        )
        sol2 = solve(
            prob2, ImplicitEuler(nlsolve = NLFunctional()), dt = 1 / 10,
            adaptive = false, reltol = 1.0e-7, abstol = 1.0e-10
        )
        @test norm(sol .- sol2) ≈ 0 atol = 1.0e-7

        sol = solve(
            prob,
            ImplicitMidpoint(
                extrapolant = :constant,
                nlsolve = NLFunctional()
            ), dt = 1 / 10, reltol = 1.0e-7,
            abstol = 1.0e-10
        )
        sol2 = solve(
            prob2,
            ImplicitMidpoint(extrapolant = :constant, nlsolve = NLFunctional()),
            dt = 1 / 10, reltol = 1.0e-7, abstol = 1.0e-10
        )
        @test norm(sol .- sol2) ≈ 0 atol = 1.0e-7

        sol = solve(
            prob, ImplicitEuler(nlsolve = NLAnderson()), dt = 1 / 10,
            adaptive = false
        )
        sol2 = solve(
            prob2, ImplicitEuler(nlsolve = NLAnderson()), dt = 1 / 10,
            adaptive = false
        )
        @test norm(sol .- sol2) ≈ 0 atol = 1.0e-7
        @test norm(sol.u[end] .- sol2.u[end]) ≈ 0 atol = 1.0e-7

        sol = solve(
            prob, ImplicitMidpoint(extrapolant = :constant, nlsolve = NLAnderson()),
            dt = 1 / 10, reltol = 1.0e-7, abstol = 1.0e-10
        )
        sol2 = solve(
            prob2,
            ImplicitMidpoint(extrapolant = :constant, nlsolve = NLAnderson()),
            dt = 1 / 10, reltol = 1.0e-7, abstol = 1.0e-10
        )
        @test norm(sol .- sol2) ≈ 0 atol = 1.0e-7
        @test norm(sol.u[end] .- sol2.u[end]) ≈ 0 atol = 1.0e-7
    end
end

# Singular mass matrices

@testset "Mass Matrix Tests with Singular Mass Matrices" begin
    function f!(du, u, p, t)
        du[1] = u[2]
        du[2] = u[2] - 1.0
        return
    end

    u0 = [0.0, 1.0]
    tspan = (0.0, 1.0)

    M = Diagonal([1.0, 0.0])

    m_ode_prob = ODEProblem(ODEFunction(f!; mass_matrix = M), u0, tspan)
    @test_nowarn sol = @inferred solve(m_ode_prob, Rosenbrock23(autodiff = AutoForwardDiff(chunksize = 2)))

    M = [
        0.637947 0.637947
        0.637947 0.637947
    ]

    function f2!(du, u, p, t)
        du[1] = u[2]
        du[2] = u[1]
        return
    end
    u0 = zeros(2)

    m_ode_prob = ODEProblem(ODEFunction(f2!; mass_matrix = M), u0, tspan)
    println("Rosenbrocks")
    sol1 = @test_nowarn solve(m_ode_prob, Rodas5P(), reltol = 1.0e-10, abstol = 1.0e-10)
    sol2 = @test_nowarn solve(m_ode_prob, RadauIIA5(), reltol = 1.0e-10, abstol = 1.0e-10)
    sol3 = @test_nowarn solve(m_ode_prob, Cash4(), reltol = 1.0e-10, abstol = 1.0e-10)
    println("SDIRKs")
    sol4 = @test_nowarn solve(m_ode_prob, ImplicitEuler(), reltol = 1.0e-10, abstol = 1.0e-10)
    sol5 = @test_nowarn solve(
        m_ode_prob, ImplicitMidpoint(), dt = 1 / 1000, reltol = 1.0e-10,
        abstol = 1.0e-10
    )
    sol6 = @test_nowarn solve(m_ode_prob, Trapezoid(), reltol = 1.0e-10, abstol = 1.0e-10)
    println("BDFs")
    sol7 = @test_nowarn solve(m_ode_prob, QNDF1(), reltol = 1.0e-10, abstol = 1.0e-10)
    sol8 = @test_nowarn solve(m_ode_prob, QBDF1(), reltol = 1.0e-10, abstol = 1.0e-10)
    sol9 = @test_nowarn solve(m_ode_prob, QNDF2(), reltol = 1.0e-10, abstol = 1.0e-10)
    sol10 = @test_nowarn solve(m_ode_prob, QBDF2(), reltol = 1.0e-10, abstol = 1.0e-10)
    sol11 = @test_nowarn solve(m_ode_prob, ABDF2(), reltol = 1.0e-10, abstol = 1.0e-10)
    sol12 = @test_nowarn solve(m_ode_prob, QBDF(), reltol = 1.0e-10, abstol = 1.0e-10)
    sol13 = @test_nowarn solve(m_ode_prob, QNDF(), reltol = 1.0e-10, abstol = 1.0e-10)

    @test sol1.u[end] ≈ sol2.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol3.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol4.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol5.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol6.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol7.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol8.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol9.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol10.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol11.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol12.u[end] atol = 1.0e-9
    @test sol1.u[end] ≈ sol13.u[end] atol = 1.0e-9
end

function _norm_dsol2(alg, prob, prob2; kwargs...)
    sol = solve(prob, alg; kwargs...)
    sol2 = solve(prob2, alg; kwargs...)
    return norm(sol.u[end] .- sol2.u[end])
end
@testset "Dependent Mass Matrix Tests" for mm in (dependent_M1, dependent_M2)
    # test each method for exactness
    for iip in (false, true)
        @show iip
        prob, prob2 = make_mm_probs(mm, Val(iip))
        eulersol = solve(prob, ImplicitEuler(nlsolve = NLNewton(κ = 1.0e-10)), reltol = 1.0e-10)
        @test _norm_dsol2(
            ImplicitEuler(nlsolve = NLNewton(κ = 1.0e-10)), prob, prob2,
            reltol = 1.0e-10
        ) ≈ 0 atol = 5.0e-4
        @test_skip _norm_dsol2(
            ImplicitMidpoint(nlsolve = NLNewton(κ = 1.0e-10)), prob, prob2,
            tstops = eulersol.t
        ) ≈ 0 atol = 1.0e-6
        @test_skip _norm_dsol(RadauIIA5(), prob, prob2) ≈ 0 atol = 1.0e-12

        @test_skip _norm_dsol(QNDF1(), prob, prob2) ≈ 0 atol = 1.0e-7
        @test_skip _norm_dsol(QBDF1(), prob, prob2) ≈ 0 atol = 1.0e-7
        @test_skip _norm_dsol(QNDF2(), prob, prob2) ≈ 0 atol = 1.0e-7
        @test_skip _norm_dsol(QBDF2(), prob, prob2) ≈ 0 atol = 1.0e-7
        @test_skip _norm_dsol(ABDF2(), prob, prob2) ≈ 0 atol = 1.0e-7
        @test_skip _norm_dsol(QBDF(), prob, prob2) ≈ 0 atol = 1.0e-7
        @test_skip _norm_dsol(QNDF(), prob, prob2) ≈ 0 atol = 1.0e-7
    end

    # test functional iteration
    @test_skip for iip in (false, true)
        @show iip
        prob, prob2 = make_mm_probs(mm, Val(iip))

        sol = solve(
            prob, ImplicitEuler(nlsolve = NLFunctional()), dt = 1 / 10,
            adaptive = false, reltol = 1.0e-7, abstol = 1.0e-10
        )
        sol2 = solve(
            prob2, ImplicitEuler(nlsolve = NLFunctional()), dt = 1 / 10,
            adaptive = false, reltol = 1.0e-7, abstol = 1.0e-10
        )
        @test_broken norm(sol .- sol2) ≈ 0 atol = 1.0e-7

        sol = solve(
            prob,
            ImplicitMidpoint(
                extrapolant = :constant,
                nlsolve = NLFunctional()
            ), dt = 1 / 10, reltol = 1.0e-7,
            abstol = 1.0e-10
        )
        sol2 = solve(
            prob2,
            ImplicitMidpoint(extrapolant = :constant, nlsolve = NLFunctional()),
            dt = 1 / 10, reltol = 1.0e-7, abstol = 1.0e-10
        )
        @test_broken norm(sol .- sol2) ≈ 0 atol = 1.0e-7

        sol = solve(
            prob, ImplicitEuler(nlsolve = NLAnderson()), dt = 1 / 10,
            adaptive = false
        )
        sol2 = solve(
            prob2, ImplicitEuler(nlsolve = NLAnderson()), dt = 1 / 10,
            adaptive = false
        )
        @test norm(sol .- sol2) ≈ 0 atol = 1.0e-7
        @test norm(sol.u[end] .- sol2.u[end]) ≈ 0 atol = 1.0e-7

        sol = solve(
            prob, ImplicitMidpoint(extrapolant = :constant, nlsolve = NLAnderson()),
            dt = 1 / 10, reltol = 1.0e-7, abstol = 1.0e-10
        )
        sol2 = solve(
            prob2,
            ImplicitMidpoint(extrapolant = :constant, nlsolve = NLAnderson()),
            dt = 1 / 10, reltol = 1.0e-7, abstol = 1.0e-10
        )
        @test norm(sol .- sol2) ≈ 0 atol = 1.0e-7
        @test norm(sol.u[end] .- sol2.u[end]) ≈ 0 atol = 1.0e-7
    end
end

# Check https://github.com/SciML/DifferentialEquations.jl/issues/757
n = 3
Λ_func = t -> exp(-t) * ones(n) |> Diagonal |> Matrix
τ = 0.2
function dynamics!(du, u, p, t)
    du .= u - Λ_func(t - τ)
    return nothing
end
function dynamics(u, p, t)
    return u - Λ_func(t - τ)
end

x0 = zeros(n, n)
M = zeros(n * n) |> Diagonal
M[1, 1] = true # zero mass matrix breaks rosenbrock
f = ODEFunction{true, SciMLBase.AutoSpecialize}(dynamics!, mass_matrix = M)
tspan = (0, 10.0)
adalg = AutoForwardDiff(chunksize = n)
prob = ODEProblem(f, x0, tspan)
foop = ODEFunction{false, SciMLBase.AutoSpecialize}(dynamics, mass_matrix = M)
proboop = ODEProblem(f, x0, tspan)
@test_broken sol = @inferred solve(prob, Rosenbrock23(autodiff = adalg))
# The Rodas4 + ShampineCollocationInit cases infer since get_differential_vars
# got a statically resolvable Diagonal branch; keep @inferred so an inference
# regression errors, but @test_broken would error ("non-Boolean") on success.
@test (@inferred solve(prob, Rodas4(autodiff = adalg), initializealg = ShampineCollocationInit()); true)
@test_broken sol = @inferred solve(proboop, Rodas5())
@test (@inferred solve(proboop, Rodas4(), initializealg = ShampineCollocationInit()); true)

# InitialFailure with chunksize matching ODE size (> alg vars), #3157
@testset "Mass Matrix: less alg vars than ODE AD chunksize" begin
    function dae(du, u, p, t)
        du[1] = u[2]
        du[2] = exp(u[1]) + exp(u[2]) # deliberately not satisfiable
        return nothing
    end
    dae_func = ODEFunction(dae, mass_matrix = Diagonal([1.0, 0.0]))
    prob = ODEProblem(dae_func, [0.0, 1.0], (0.0, 1.0))
    adalg = AutoForwardDiff(chunksize = 2)
    sol = solve(prob, Rodas5P(autodiff = adalg), initializealg = BrownFullBasicInit())
    @test sol.retcode == SciMLBase.ReturnCode.InitialFailure # previously, errored instead

    function dae_oop(u, p, t)
        return [u[2], exp(u[1]) + exp(u[2])]
    end
    dae_f_oop = ODEFunction(dae_oop, mass_matrix = Diagonal([1.0, 0.0]))
    prob = ODEProblem(dae_f_oop, [0.0, 1.0], (0.0, 1.0))
    adalg = AutoForwardDiff(chunksize = 2)
    sol = solve(prob, Rodas5P(autodiff = adalg), initializealg = BrownFullBasicInit())
    @test sol.retcode == SciMLBase.ReturnCode.InitialFailure # previously, errored instead

    # Have more parameters than there are algebraic unknowns: parameters should not be
    # considered in the chunksize for initialization.
    function dae_p(du, u, p, t)
        du[1] = u[2] * p[1]
        du[2] = exp(u[1] * p[2]) + exp(u[2] * p[3]) # deliberately not satisfiable
        return nothing
    end
    dae_p_f = ODEFunction(dae_p, mass_matrix = Diagonal([1.0, 0.0]))
    prob = ODEProblem(dae_p_f, [0.0, 1.0], (0.0, 1.0), (1.0, 1.0, 1.0))
    adalg = AutoForwardDiff(chunksize = 2)
    sol = solve(prob, Rodas5P(autodiff = adalg), initializealg = BrownFullBasicInit())
    @test sol.retcode == SciMLBase.ReturnCode.InitialFailure # previously, errored instead
end
