using QuantumCumulants
using Symbolics: Symbolics, @variables
using SymbolicUtils
using ModelingToolkitBase: @named, mtkcompile, ODEProblem, unknowns
using OrdinaryDiffEqTsit5: Tsit5, solve, ReturnCode
using Test

# Higher-order cumulant closure of a damped JC system.

@testset "higher-order: 4th-order cumulant closure of JC laser" begin
    @variables Δ::Real g::Real γ::Real κ::Real ν::Real

    hf = FockSpace(:cavity)
    ha = NLevelSpace(:atom, 2)
    h = hf ⊗ ha
    a = Destroy(h, :a)
    s = Transition(h, :σ, 1, 2)

    H = Δ * a' * a + g * (a' * s + a * s')
    J = [a, s, s']
    rates = [κ, γ, ν]

    # `phase_invariant` is the shared U(1) filter from runtests.jl `init_code`.
    he4 = complete(
        meanfield(a' * a, H, J; rates = rates);
        order = 4, filter_func = phase_invariant
    )
    @test length(he4.equations) > 0
    @test isempty(find_missing(he4; filter_func = phase_invariant))

    @named sys4 = System(he4)
    sys4_c = mtkcompile(sys4)
    u0 = Dict(unknowns(sys4_c) .=> zeros(ComplexF64, length(unknowns(sys4_c))))
    ps = [Δ, g, γ, κ, ν]
    pn = [1.0, 1.5, 0.25, 1.0, 4.0]
    prob = ODEProblem(sys4_c, merge(u0, Dict(ps .=> pn)), (0.0, 20.0))
    sol = solve(prob, Tsit5(); abstol = 1.0e-8, reltol = 1.0e-8)
    @test sol.retcode == ReturnCode.Success
    n_ss = real(get_solution(sol, a' * a, he4)(sol.t[end]))
    @test isfinite(n_ss)
    @test n_ss >= 0
end
