using ModelingToolkit
using NonlinearSolve, SCCNonlinearSolve
using OrdinaryDiffEq
using OrdinaryDiffEqBDF
using SciMLBase, Symbolics
using StaticArrays
using LinearAlgebra, Test
using SymbolicUtils
import ModelingToolkitBase
using Symbolics: SymbolicT, VartypeT, unwrap
using ModelingToolkit: t_nounits as t, D_nounits as D

@testset "Trivial case" begin
    function f!(du, u, p)
        du[1] = cos(u[2]) - u[1]
        du[2] = sin(u[1] + u[2]) + u[2]
        du[3] = 2u[4] + u[3] + 1.0
        du[4] = u[5]^2 + u[4]
        du[5] = u[3]^2 + u[5]
        du[6] = u[1] + u[2] + u[3] + u[4] + u[5] + 2.0u[6] + 2.5u[7] + 1.5u[8]
        du[7] = u[1] + u[2] + u[3] + 2.0u[4] + u[5] + 4.0u[6] - 1.5u[7] + 1.5u[8]
        du[8] = u[1] + 2.0u[2] + 3.0u[3] + 5.0u[4] + 6.0u[5] + u[6] - u[7] - u[8]
    end
    @variables u[1:8] [irreducible = true]
    eqs = Any[0 for _ in 1:8]
    f!(eqs, u, nothing)
    eqs = 0 .~ eqs
    @named model = System(eqs)
    @test_throws ["simplified", "required"] SCCNonlinearProblem(model, [])
    _model = mtkcompile(model; split = false)
    @test_throws ["not compatible"] SCCNonlinearProblem(_model, [])
    model = mtkcompile(model)
    prob = NonlinearProblem(model, [u => zeros(8)])
    sccprob = SCCNonlinearProblem(model, collect(u[1:5]) .=> zeros(5))
    sol1 = solve(prob, NewtonRaphson())
    sol2 = solve(sccprob, NewtonRaphson())
    @test SciMLBase.successful_retcode(sol1)
    @test SciMLBase.successful_retcode(sol2)
    @test sol1[u] ≈ sol2[u]

    sccprob = SCCNonlinearProblem{false}(model, SA[(collect(u[1:5]) .=> zeros(5))...])
    for prob in sccprob.probs
        if prob isa LinearProblem
            @test prob.A isa SMatrix
            @test prob.b isa SVector
        else
            @test prob.u0 isa SVector
        end
        @test !SciMLBase.isinplace(prob)
    end

    # Test BLT sorted
    @test istril(StructuralTransformations.sorted_incidence_matrix(model), 2)
end

@testset "With parameters" begin
    function f!(du, u, (p1, p2), t)
        x = (*)(p1[4], u[1])
        y = (*)(p1[4], (+)(0.1016, (*)(-1, u[1])))
        z1 = ifelse(
            (<)(p2[1], 0),
            (*)((*)(457896.07999999996, p1[2]), sqrt((*)(1.1686468413521012e-5, p1[3]))),
            0
        )
        z2 = ifelse(
            (>)(p2[1], 0),
            (*)((*)((*)(0.58, p1[2]), sqrt((*)(1 // 86100, p1[3]))), u[4]),
            0
        )
        z3 = ifelse(
            (>)(p2[1], 0),
            (*)((*)(457896.07999999996, p1[2]), sqrt((*)(1.1686468413521012e-5, p1[3]))),
            0
        )
        z4 = ifelse(
            (<)(p2[1], 0),
            (*)((*)((*)(0.58, p1[2]), sqrt((*)(1 // 86100, p1[3]))), u[5]),
            0
        )
        du[1] = p2[1]
        du[2] = (+)(z1, (*)(-1, z2))
        du[3] = (+)(z3, (*)(-1, z4))
        du[4] = (+)((*)(-1, u[2]), (*)((*)(1 // 86100, y), u[4]))
        du[5] = (+)((*)(-1, u[3]), (*)((*)(1 // 86100, x), u[5]))
    end
    p = (
        [
            0.04864391799335977, 7.853981633974484e-5, 1.4034843205574914,
            0.018241469247509915, 300237.05, 9.226186337232914,
        ],
        [0.0508],
    )
    u0 = [0.0, 0.0, 0.0, 789476.0, 101325.0]
    tspan = (0.0, 1.0)
    mass_matrix = [
        1.0 0.0 0.0 0.0 0.0; 0.0 1.0 0.0 0.0 0.0; 0.0 0.0 1.0 0.0 0.0;
        0.0 0.0 0.0 0.0 0.0; 0.0 0.0 0.0 0.0 0.0
    ]
    dt = 1.0e-3
    function nlf(u1, (u0, p))
        resid = Any[0 for _ in u0]
        f!(resid, u1, p, 0.0)
        return mass_matrix * (u1 - u0) - dt * resid
    end

    prob = NonlinearProblem(nlf, u0, (u0, p))
    @test_throws Exception solve(prob, SimpleNewtonRaphson(), abstol = 1.0e-9)
    sol = solve(prob, TrustRegion(); abstol = 1.0e-9)

    @variables u[1:5] [irreducible = true]
    @parameters p1[1:6] p2
    eqs = 0 .~ collect(nlf(u, (u0, (p1, p2))))
    @mtkcompile sys = System(eqs, [u], [p1, p2])
    sccprob = SCCNonlinearProblem(sys, [u => u0, p1 => p[1], p2 => p[2][]])
    sccsol = solve(sccprob, SimpleNewtonRaphson(); abstol = 1.0e-9)
    sccresid = prob.f(sccsol[u], (u0, p))
    @test SciMLBase.successful_retcode(sccsol)
    @test norm(sccresid) < norm(sol.resid)

    # Test BLT sorted
    @test istril(StructuralTransformations.sorted_incidence_matrix(sys), 1)
end

@testset "Transistor amplifier" begin
    C = [k * 1.0e-6 for k in 1:5]
    Ub = 6
    UF = 0.026
    α = 0.99
    β = 1.0e-6
    R0 = 1000
    R = 9000
    Ue(t) = 0.1 * sin(200 * π * t)

    function transamp(out, du, u, p, t)
        g(x) = 1.0e-6 * (exp(x / 0.026) - 1)
        y1, y2, y3, y4, y5, y6, y7, y8 = u
        out[1] = -Ue(t) / R0 + y1 / R0 + C[1] * du[1] - C[1] * du[2]
        out[2] = -Ub / R + y2 * 2 / R - (α - 1) * g(y2 - y3) - C[1] * du[1] + C[1] * du[2]
        out[3] = -g(y2 - y3) + y3 / R + C[2] * du[3]
        out[4] = -Ub / R + y4 / R + α * g(y2 - y3) + C[3] * du[4] - C[3] * du[5]
        out[5] = -Ub / R + y5 * 2 / R - (α - 1) * g(y5 - y6) - C[3] * du[4] + C[3] * du[5]
        out[6] = -g(y5 - y6) + y6 / R + C[4] * du[6]
        out[7] = -Ub / R + y7 / R + α * g(y5 - y6) + C[5] * du[7] - C[5] * du[8]
        out[8] = y8 / R - C[5] * du[7] + C[5] * du[8]
    end

    u0 = [0, Ub / 2, Ub / 2, Ub, Ub / 2, Ub / 2, Ub, 0]
    du0 = [
        51.338775,
        51.338775,
        -Ub / (2 * (C[2] * R)),
        -24.9757667,
        -24.9757667,
        -Ub / (2 * (C[4] * R)),
        -10.00564453,
        -10.00564453,
    ]
    daeprob = DAEProblem(transamp, du0, u0, (0.0, 0.1))
    daesol = solve(daeprob, DImplicitEuler())

    t0 = daesol.t[5]
    t1 = daesol.t[6]
    u0 = daesol.u[5]
    u1 = daesol.u[6]
    dt = t1 - t0

    @variables y(t)[1:8]
    eqs = Any[0 for _ in 1:8]
    transamp(eqs, collect(D(y)), y, nothing, t)
    eqs = 0 .~ eqs
    subrules = Dict(Symbolics.unwrap(D(y[i])) => ((y[i] - u0[i]) / dt) for i in 1:8)
    eqs = substitute.(eqs, (subrules,))
    @mtkcompile sys = System(eqs)
    prob = NonlinearProblem(sys, [y => u0, t => t0])
    sol = solve(prob, NewtonRaphson(); abstol = 1.0e-12)

    sccprob = SCCNonlinearProblem(sys, [y => u0, t => t0])
    sccsol = solve(sccprob, NewtonRaphson(); abstol = 1.0e-12)

    @test sol.u ≈ sccsol.u atol = 1.0e-8

    # Test BLT sorted
    @test istril(StructuralTransformations.sorted_incidence_matrix(sys), 1)
end

@testset "Expression caching" begin
    @variables x[1:4] = rand(4)
    val = Ref(0)
    function func(x, y)
        val[] += 1
        x + y
    end
    @register_symbolic func(x, y)
    @mtkcompile sys = System(
        [
            0 ~ x[1]^3 + x[2]^3 - 5
            0 ~ sin(x[1] - x[2]) - 0.5
            0 ~ func(x[1], x[2]) * exp(x[3]) - x[4]^3 - 5
            0 ~ func(x[1], x[2]) * exp(x[4]) - x[3]^3 - 4
        ]
    )
    sccprob = SCCNonlinearProblem(sys, [])
    @test val[] == 0
    sccsol = solve(sccprob, NewtonRaphson())
    @test SciMLBase.successful_retcode(sccsol)
    @test val[] == 1
end

import ModelingToolkitStandardLibrary.Blocks as B
import ModelingToolkitStandardLibrary.Mechanical.Translational as T
import ModelingToolkitStandardLibrary.Hydraulic.IsothermalCompressible as IC

@testset "Caching of subexpressions of different types" begin
    liquid_pressure(rho, rho_0, bulk) = (rho / rho_0 - 1) * bulk
    gas_pressure(rho, rho_0, p_gas, rho_gas) = rho * ((0 - p_gas) / (rho_0 - rho_gas))
    full_pressure(
        rho,
        rho_0,
        bulk,
        p_gas,
        rho_gas
    ) = ifelse(
        rho >= rho_0, liquid_pressure(rho, rho_0, bulk),
        gas_pressure(rho, rho_0, p_gas, rho_gas)
    )

    @component function Volume(;
            #parameters
            area,
            direction = +1,
            x_int,
            name
        )
        pars = @parameters begin
            area = area
            x_int = x_int
            rho_0 = 1000
            bulk = 1.0e9
            p_gas = -1000
            rho_gas = 1
        end

        vars = @variables begin
            x(t) = x_int
            dx(t), [guess = 0]
            p(t), [guess = 0]
            f(t), [guess = 0]
            rho(t), [guess = 0]
            m(t), [guess = 0]
            dm(t), [guess = 0]
        end

        systems = @named begin
            port = IC.HydraulicPort()
            flange = T.MechanicalPort()
        end

        eqs = [
            # connectors
            port.p ~ p
            port.dm ~ dm
            flange.v * direction ~ dx
            flange.f * direction ~ -f

            # differentials
            D(x) ~ dx
            D(m) ~ dm

            # physics
            p ~ full_pressure(rho, rho_0, bulk, p_gas, rho_gas)
            f ~ p * area
            m ~ rho * x * area
        ]

        return System(eqs, t, vars, pars; name, systems)
    end

    systems = @named begin
        fluid = IC.HydraulicFluid(; bulk_modulus = 1.0e9)

        src1 = IC.Pressure()
        src2 = IC.Pressure()

        vol1 = Volume(; area = 0.01, direction = +1, x_int = 0.1)
        vol2 = Volume(; area = 0.01, direction = +1, x_int = 0.1)

        mass = T.Mass(; m = 10)

        sin1 = B.Sine(; frequency = 0.5, amplitude = +0.5e5, offset = 10.0e5)
        sin2 = B.Sine(; frequency = 0.5, amplitude = -0.5e5, offset = 10.0e5)
    end

    eqs = [
        connect(fluid, src1.port)
        connect(fluid, src2.port)
        connect(src1.port, vol1.port)
        connect(src2.port, vol2.port)
        connect(vol1.flange, mass.flange, vol2.flange)
        connect(src1.p, sin1.output)
        connect(src2.p, sin2.output)
        src2.port.dm ~ 0
    ]

    initialization_eqs = [
        mass.s ~ 0.0
        mass.v ~ 0.0
    ]

    @mtkcompile sys = System(eqs, t, [], []; systems, initialization_eqs)
    prob = ODEProblem(sys, [], (0, 5))
    sol = solve(prob)
    @test SciMLBase.successful_retcode(sol)
end

@testset "Array variables split across SCCs" begin
    @variables x[1:3]
    @parameters (f::Function)(..)
    @mtkcompile sys = System(
        [
            0 ~ x[1]^2 - 9, x[2] ~ 2x[1], 0 ~ x[3]^2 - x[1]^2 + f(x),
        ]
    )
    prob = SCCNonlinearProblem(sys, [x => ones(3), f => sum])
    sol = solve(prob, NewtonRaphson())
    @test SciMLBase.successful_retcode(sol)
end

@testset "SCCNonlinearProblem retains parameter order" begin
    @variables x y z
    @parameters σ β ρ
    @mtkcompile fullsys = System(
        [0 ~ x^3 * β + y^3 * ρ - σ, 0 ~ x^2 + 2x * y + y^2, 0 ~ z^2 - 4z + 4],
        [x, y, z], [σ, β, ρ]
    )

    u0 = [
        x => 1.0,
        y => 0.0,
        z => 0.0,
    ]

    p = [
        σ => 28.0,
        ρ => 10.0,
        β => 8 / 3,
    ]

    sccprob = SCCNonlinearProblem(fullsys, [u0; p])
    @test isequal(parameters(fullsys), parameters(sccprob.f.sys))
end

@testset "Vector parameters in function arguments" begin
    @variables x y
    @parameters p[1:2] (f::Function)(..)

    @mtkcompile sys = System([x^2 - p[1]^2 ~ 0, y^2 ~ f(p)])
    prob = SCCNonlinearProblem(sys, [x => 1.0, y => 1.0, p => ones(2), f => sum])
    @test_nowarn solve(prob, NewtonRaphson())
end

@testset "Subproblems retain bound parameters" begin
    # This can cause bad codegen, since `update_A`/`update_b` for LinearProblem SCCs
    # won't have access to the necessary bound parameters. I'm unable to MWE that behavior,
    # but checking the `bound_parameters` of the subsystems is a good proxy.
    @parameters p[1:3, 1:3] q[1:3, 1:3] = p r[1:3] s[1:3] = r
    @variables x[1:3] y[1:3]
    @mtkcompile sys = System([p * x ~ r, q * y ~ s])
    @test issetequal(bound_parameters(sys), [q, s])
    prob = SCCNonlinearProblem(sys, [p => ones(3, 3), r => ones(3)]; combine_sccs = false)
    @test issetequal(bound_parameters(prob.probs[1].f.sys), [q, s])
    @test issetequal(bound_parameters(prob.probs[2].f.sys), [q, s])
end

@testset "SCCs are sorted when combined" begin
    # NOTE: This test is best-effort. It depends on the `active::Set{Int}` in the
    # `SCCDecomposition` constructor hashing `2` before `1`. In other words,
    # ```julia
    # active = Set{Int}([1, 2])
    # @assert collect(active) == [2, 1]
    # ```
    # But I can't guarantee that.
    @variables x [irreducible = true] y [irreducible = true] z [irreducible = true] a b
    @mtkcompile sys = System([2x ~ 4 + b, 3y ~ 7 + a, x^2 + z^2 ~ 100, a ~ 3, b ~ sin(a)])
    # Ensure the system simplifies the way we expect
    dvs = unknowns(sys)
    @test issetequal(dvs, [y, x, z])
    sched = ModelingToolkit.get_schedule(sys)
    # NOTE: Intentionally not `var_sccs`
    @test isequal(dvs[reduce(vcat, sched.var_sccs)], [y, x, z])
    prob = SCCNonlinearProblem(sys, [z => 1])
    sol = solve(prob)
    @test SciMLBase.successful_retcode(sol)
end

@testset "`subexpressions_not_involving_vars!`" begin
    @variables x[1:3]
    @parameters (f::Function)(..)
    ir = IRStructure{VartypeT}()
    expr = f([x[1], 2x[2], x[2]^2 + x[3]]) + f(x) + f(x[2] + 1)
    state = Dict{SymbolicT, SymbolicT}()
    banned_vars = Set{SymbolicT}([x[3], x])
    ModelingToolkitBase.subexpressions_not_involving_vars!(ir, [unwrap(expr)], banned_vars, state)
    @test issetequal(keys(state), [x[1], 2x[2], x[2]^2, f(x[2] + 1)])
end
