using ModelingToolkitBase, OrdinaryDiffEq, JumpProcesses, DynamicQuantities
using Symbolics
import SymbolicUtils as SU
using Test
MT = ModelingToolkitBase
using ModelingToolkitBase: t, D
@parameters τ [unit = u"s"] γ
@variables E(t) [unit = u"J"] P(t) [unit = u"W"]

const unitless = MT.get_unit(0.5)

# Basic access
@test MT.get_unit(t) == u"s"
@test MT.get_unit(E) == u"J"
@test MT.get_unit(τ) == u"s"
@test MT.get_unit(γ) == unitless
@test MT.get_unit(MT.SciMLBase.NullParameters()) == unitless

eqs = [
    D(E) ~ P - E / τ
    0 ~ P
]
@test MT.validate(eqs)
@named sys = System(eqs, t)

@test !MT.validate(D(D(E)) ~ P)
@test !MT.validate(0 ~ P + E * τ)

# Disabling unit validation/checks selectively
@test_throws MT.ArgumentError System(eqs, t, [E, P, t], [τ], name = :sys)
@test_nowarn System(eqs, t, [E, P, t], [τ], name = :sys, checks = MT.CheckUnits)
eqs = [
    D(E) ~ P - E / τ
    0 ~ P + E * τ
]
@test_throws MT.ValidationError System(eqs, t, name = :sys, checks = MT.CheckAll)
@test_throws MT.ValidationError System(eqs, t, name = :sys, checks = true)
System(eqs, t, name = :sys, checks = MT.CheckNone)
System(eqs, t, name = :sys, checks = false)
@test_throws MT.ValidationError System(
    eqs, t, name = :sys,
    checks = MT.CheckComponents | MT.CheckUnits
)
@named sys = System(eqs, t, checks = MT.CheckComponents)
@test_throws MT.ValidationError System(
    eqs, t, [E, P], [τ], name = :sys,
    checks = MT.CheckUnits
)

# connection validation
@connector function Pin(; name)
    sts = @variables(
        v(t) = 1.0, [unit = u"V"],
        i(t) = 1.0, [unit = u"A", connect = Flow]
    )
    System(Equation[], t, sts, []; name = name)
end
@connector function OtherPin(; name)
    sts = @variables(
        v(t) = 1.0, [unit = u"mV"],
        i(t) = 1.0, [unit = u"mA", connect = Flow]
    )
    System(Equation[], t, sts, []; name = name)
end
@connector function LongPin(; name)
    sts = @variables(
        v(t) = 1.0, [unit = u"V"],
        i(t) = 1.0, [unit = u"A", connect = Flow],
        x(t) = 1.0
    )
    System(Equation[], t, sts, []; name = name)
end
@named p1 = Pin()
@named p2 = Pin()
@named lp = LongPin()
good_eqs = [connect(p1, p2)]
@test MT.validate(good_eqs)
@named sys = System(good_eqs, t, [], [])
@named op = OtherPin()
bad_eqs = [connect(p1, op)]
@test !MT.validate(bad_eqs)
@test_throws MT.ValidationError @named sys = System(bad_eqs, t, [], [])
@named op2 = OtherPin()
good_eqs = [connect(op, op2)]
@test MT.validate(good_eqs)
@named sys = System(good_eqs, t, [], [])

# Array variables
@variables x(t)[1:3] [unit = u"m"]
@parameters v[1:3] = [1, 2, 3] [unit = u"m/s"]
eqs = [D(x) ~ v]
System(eqs, t, name = :sys)

# Nonlinear system
@parameters a [unit = u"kg"^-1]
@variables x [unit = u"kg"]
eqs = [
    0 ~ a * x,
]
@named nls = System(eqs, [x], [a])

# SDE test w/ noise vector
@parameters τ [unit = u"s"] Q [unit = u"W"]
@variables E(t) [unit = u"J"] P(t) [unit = u"W"]
eqs = [
    D(E) ~ P - E / τ
    P ~ Q
]

noiseeqs = [
    0.1us"W",
    0.1us"W",
]
@named sys = SDESystem(eqs, noiseeqs, t, [P, E], [τ, Q])

noiseeqs = [
    0.1u"W",
    0.1u"W",
]
@test_throws MT.ValidationError @named sys = SDESystem(eqs, noiseeqs, t, [P, E], [τ, Q])

# With noise matrix
noiseeqs = [
    0.1us"W" 0.1us"W"
    0.1us"W" 0.1us"W"
]
@named sys = SDESystem(eqs, noiseeqs, t, [P, E], [τ, Q])

# Invalid noise matrix
noiseeqs = [
    0.1us"W" 0.1us"W"
    0.1us"W" 0.1us"s"
]
@test !MT.validate(eqs, noiseeqs)

# Non-trivial simplifications
@variables V(t) [unit = u"m"^3] L(t) [unit = u"m"]
@parameters v [unit = u"m/s"] r [unit = u"m"^3 / u"s"]
eqs = [
    D(L) ~ v,
    V ~ L^3,
]
@named sys = System(eqs, t)
sys_simple = mtkcompile(sys)

eqs = [
    D(V) ~ r,
    V ~ L^3,
]
@named sys = System(eqs, t)
sys_simple = mtkcompile(sys)

@variables V [unit = u"m"^3] L [unit = u"m"]
@parameters v [unit = u"m/s"] r [unit = u"m"^3 / u"s"]
eqs = [
    V ~ r * t,
    V ~ L^3,
]
@named sys = System(eqs, [V, L], [t, r])
sys_simple = mtkcompile(sys)

eqs = [
    L ~ v * t,
    V ~ L^3,
]
@named sys = System(eqs, [V, L], [t, r, v])
sys_simple = mtkcompile(sys)

#Jump System
@parameters β [unit = u"(mol^2*s)^-1"] γ [unit = u"(mol*s)^-1"] jumpmol [
    unit = u"mol",
]
@variables S(t) [unit = u"mol"] I(t) [unit = u"mol"] R(t) [unit = u"mol"]
rate₁ = β * S * I
affect₁ = [S ~ S - 1 * jumpmol, I ~ I + 1 * jumpmol]
rate₂ = γ * I
affect₂ = [I ~ I - 1 * jumpmol, R ~ R + 1 * jumpmol]
j₁ = ConstantRateJump(rate₁, affect₁)
j₂ = VariableRateJump(rate₂, affect₂)
js = JumpSystem([j₁, j₂], t, [S, I, R], [β, γ], name = :sys)

affect_wrong = [S ~ S - jumpmol, I ~ I + 1]
j_wrong = ConstantRateJump(rate₁, affect_wrong)
@test_throws MT.ValidationError JumpSystem([j_wrong, j₂], t, [S, I, R], [β, γ], name = :sys)

rate_wrong = γ^2 * I
j_wrong = ConstantRateJump(rate_wrong, affect₂)
@test_throws MT.ValidationError JumpSystem([j₁, j_wrong], t, [S, I, R], [β, γ], name = :sys)

# mass action jump tests for SIR model
maj1 = SymbolicMassActionJump(2 * β / 2, [S => 1, I => 1], [S => -1, I => 1])
maj2 = SymbolicMassActionJump(γ, [I => 1], [I => -1, R => 1])
@named js3 = JumpSystem([maj1, maj2], t, [S, I, R], [β, γ])

#Test unusual jump system
@parameters β γ
@variables S(t) I(t) R(t)

maj1 = SymbolicMassActionJump(2.0, [0 => 1], [S => 1])
maj2 = SymbolicMassActionJump(γ, [S => 1], [S => -1])
@named js4 = JumpSystem([maj1, maj2], ModelingToolkitBase.t_nounits, [S], [β, γ])

@testset "Initialization checks" begin
    @component function PendulumUnits(; name, g = nothing, L = nothing)
        pars = @parameters begin
            g = g, [unit = u"m/s^2"]
            L = L, [unit = u"m"]
        end

        systems = @named begin
        end

        vars = @variables begin
            x(t), [unit = u"m"]
            y(t), [state_priority = 10, unit = u"m"]
            λ(t), [unit = u"s^-2"]
        end

        equations = Equation[
            D(D(x)) ~ λ * x
            D(D(y)) ~ λ * y - g
            x^2 + y^2 ~ L^2
        ]

        return System(equations, t, vars, pars; name, systems)
    end
    @mtkcompile pend = PendulumUnits()
    u0 = [pend.x => 1.0, pend.y => 0.0]
    if !@isdefined(ModelingToolkit)
        u0 = [u0; D(pend.x) => 0.0; D(pend.y) => 0.0]
    end
    p = [pend.g => 1.0, pend.L => 1.0]
    guess = [pend.λ => 0.0]
    @test prob = ODEProblem(
        pend, [u0; p], (0.0, 1.0); guesses = guess, check_units = false
    ) isa Any
end

@parameters p [unit = u"L/s"] d [unit = u"s^(-1)"]
@parameters tt [unit = u"s"]
@variables X(tt) [unit = u"L"]
DD = Differential(tt)
eqs = [DD(X) ~ p - d * X + d * X]
@test ModelingToolkitBase.validate(eqs)

@constants begin
    to_m = 1, [unit = u"m"]
end
@variables begin
    L(t), [unit = u"m"]
    L_out(t), [unit = u"1"]
end
@test to_m in SU.search_variables(Symbolics.unwrap(L_out * -to_m))

# test units for registered functions
let
    mm(X, v, K) = v * X / (X + K)
    mm2(X, v, K) = v * X / (X + K)
    Symbolics.@register_symbolic mm2(X, v, K)
    @parameters t [unit = u"s"] K [unit = u"mol/m^3"] v [unit = u"(m^6)/(s*mol^2)"]
    @variables X(t) [unit = u"mol/m^3"]
    mmunits = MT.get_unit(mm(X, v, K))
    mm2units = MT.get_unit(mm2(X, v, K))
    @test mmunits == MT.oneunit(mmunits)
    @test mm2units == MT.oneunit(mm2units)
    @test mmunits == mm2units
end

# test for array variable units https://github.com/SciML/ModelingToolkit.jl/issues/3009
let
    @variables x_vec(t)[1:3] [unit = u"1"] x_mat(t)[1:3, 1:3] [unit = u"1"]
    @test MT.get_unit(x_vec) == u"1"
    @test MT.get_unit(x_mat) == u"1"
end

# Issue #4211: collect_vars! must discover parameters in defaults with DynamicQuantities loaded
# On Julia 1.10, loading DynamicQuantities could cause collect_vars! to fail to discover
# parameters used in variable defaults due to a method invalidation bug.
@testset "Issue #4211: collect_vars! discovers parameters in defaults" begin
    using DataStructures: OrderedSet

    @parameters X0_test
    @variables X_test(t) = X0_test

    us = OrderedSet{Symbolics.SymbolicT}()
    ps = OrderedSet{Symbolics.SymbolicT}()
    MT.collect_vars!(us, ps, Symbolics.unwrap(X_test), Symbolics.unwrap(t), Symbolics.Operator; depth = 0)

    # X0_test should be discovered in X_test's default value
    @test Symbolics.unwrap(X0_test) in ps

    # Test with expression in default
    @parameters a_test b_test
    @variables Y_test(t) = a_test + 2 * b_test

    empty!(us)
    empty!(ps)
    MT.collect_vars!(us, ps, Symbolics.unwrap(Y_test), Symbolics.unwrap(t), Symbolics.Operator; depth = 0)

    @test Symbolics.unwrap(a_test) in ps
    @test Symbolics.unwrap(b_test) in ps
end
