using QuantumCumulants
using Symbolics: Symbolics, @variables, expand
using SymbolicUtils
using Test

# Robust zero check across QAdd, Number, and BasicSymbolic. Uses expand() because
# simplify does not always collapse polynomially-equivalent symbolic-average products.
_iz(x) = (
    x isa Number ? iszero(x) :
        (
            x isa SymbolicUtils.BasicSymbolic ? SymbolicUtils._iszero(expand(x)) :
            iszero(x)
        )
)

@testset "get_order" begin
    hc = FockSpace(:cavity)
    @qnumbers a::Destroy(hc)
    @test get_order(a) == 1
    @test get_order(a' * a) == 2
    @test get_order(a' * a' * a) == 3
    @test get_order(2.0) == 0
    @test get_order(average(a' * a)) == 2

    ha = NLevelSpace(:atom, (:g, :e))
    h = hc ⊗ ha
    a2 = Destroy(h, :a)
    σ(i, j) = Transition(h, :σ, i, j)
    @test get_order(a2 * σ(:g, :e) + Destroy(hc, :b)) == 2
end

@testset "cumulant_expansion basics" begin
    hf = FockSpace(:cavity)
    ha = NLevelSpace(:atom, (:g, :e))
    h = hf ⊗ ha

    a = Destroy(h, :a)

    @test _iz(cumulant_expansion(average(a^2), 1) - average(a)^2)
    @test _iz(
        cumulant_expansion(average(a' * a^2), 2) -
            (
            average(a') * average(a^2) +
                -2 * average(a') * average(a)^2 +
                2 * average(a) * average(a' * a)
        )
    )

    n = average(a' * a)
    @test isequal(cumulant_expansion(n, 2), n)
    @test _iz(simplify(cumulant_expansion(n, 1) - average(a') * average(a)))
end

@testset "cumulant_expansion scalar order sizes over operators outside the Hamiltonian" begin
    hc = FockSpace(:cavity)
    hf = FockSpace(:filter)
    h = hc ⊗ hf
    a = Destroy(h, :a, 1)
    b = Destroy(h, :b, 2)

    @variables ω::Real κ::Real
    # H and the jump touch only subspace 1; observable b touches subspace 2, so a
    # scalar `order` must size the order vector over the full system inputs,
    # including observables, not just H and the jumps.
    eqs = meanfield([b' * b], ω * a' * a, [a]; rates = [κ])

    expanded = cumulant_expansion(eqs, 2)
    @test expanded.order == [2, 2]
end

@testset "cumulant_expansion: second order" begin
    hc = FockSpace(:cavity)
    @qnumbers a::Destroy(hc)
    avg3 = average(a' * a' * a)
    expanded = cumulant_expansion(avg3, 2)
    @test get_order(expanded) <= 2
    @test !isequal(expanded, avg3)
end

@testset "cumulant_expansion(eqs, order)" begin
    hc = FockSpace(:cavity)
    @qnumbers a::Destroy(hc)
    @variables ω κ
    H = ω * a' * a + κ * a' * a' * a
    eqs = meanfield([a], H, []; rates = [])
    eqs2 = cumulant_expansion(eqs, 2)
    @test eqs2.order == [2]
    @test get_order(eqs2.equations[1].rhs) <= 2
end

@testset "cumulant_expansion(eqs, order) preserves direction" begin
    hc = FockSpace(:cavity)
    @qnumbers a::Destroy(hc)
    @variables ω κ
    H = ω * a' * a + κ * a' * a' * a
    bw = meanfield([a], H, [a]; rates = [κ], direction = Backward())
    @test bw.direction isa Backward
    @test cumulant_expansion(bw, 2).direction isa Backward
end

@testset "cumulant" begin
    hs = FockSpace[]
    for i in 1:4
        push!(hs, FockSpace(Symbol(:fock, i)))
    end
    h = ⊗(hs...)

    a = Destroy(h, :a, 1)
    b = Destroy(h, :b, 2)
    c = Destroy(h, :c, 3)
    d = Destroy(h, :d, 4)

    @test _iz(cumulant(a * b) - (average(a * b) + -1 * average(a) * average(b)))
    @test isequal(cumulant(a * b, 1), average(a * b))
    @test _iz(
        cumulant(a * b * c) - (
            average(a * b * c) + 2 * average(a) * average(b) * average(c) -
                average(a) * average(b * c) - average(b) * average(a * c) -
                average(c) * average(a * b)
        ),
    )
    @test _iz(
        average(a * b * c * d) - cumulant(a * b * c * d) -
            cumulant_expansion(average(a * b * c * d), 3),
    )
end

@testset "mixed-order cumulant" begin
    hf = FockSpace(:cavity)
    ha = NLevelSpace(:atom, 2)
    h = hf ⊗ ha

    a = Destroy(h, :a, 1)
    σ(i, j) = Transition(h, :σ, i, j)

    @test _iz(cumulant_expansion(average(a' * a), [2, 1]) - average(a' * a))
    @test _iz(
        cumulant_expansion(average(a' * σ(1, 2)), [2, 1]; mix_choice = minimum) -
            average(a') * average(σ(1, 2)),
    )

    H = a' * a + σ(2, 2) + a' * σ(1, 2) + a * σ(2, 1)
    he = meanfield([a' * a, σ(2, 2)], H)

    he_min = meanfield([a' * a, σ(2, 2)], H; mix_choice = minimum)
    he_avg1 = cumulant_expansion(he, 2)
    he_avg2 = cumulant_expansion(he, [2, 1])
    he_avg3 = cumulant_expansion(he_min, [2, 1])

    @test isequal(he_avg1.equations, he_avg2.equations)
    @test !isequal(he_avg1.equations, he_avg3.equations)
end

@testset "mixed-order N-atom laser closure" begin
    N = 2
    @variables κ g Γ23 Γ13 Γ12 Ω Δc Δ3

    hf = FockSpace(:cavity)
    ha = ⊗([NLevelSpace(Symbol(:atom, i), 3) for i in 1:N]...)
    h = hf ⊗ ha

    a = Destroy(h, :a)
    σ(i, j, k) = Transition(h, Symbol("σ_{$k}"), i, j, k + 1)

    H = -Δc * a'a +
        sum(g * (a' * σ(1, 2, i) + a * σ(2, 1, i)) for i in 1:N) +
        sum(Ω * (σ(3, 1, i) + σ(1, 3, i)) for i in 1:N) -
        sum(Δ3 * σ(3, 3, i) for i in 1:N)

    J = [
        a; [σ(1, 2, i) for i in 1:N]; [σ(1, 3, i) for i in 1:N];
        [σ(2, 3, i) for i in 1:N]
    ]
    rates = [κ; [Γ12 for i in 1:N]; [Γ13 for i in 1:N]; [Γ23 for i in 1:N]]

    ops = [a'a, σ(2, 2, 1), σ(3, 3, 1)]
    he = meanfield(ops, H, J; rates = rates)
    he_avg = cumulant_expansion(he, [2, 1, 1])

    he_c = complete(he_avg; order = [2, 1, 1])
    @test isempty(find_missing(he_c))
    @test isempty(findall(x -> (aon = acts_on(x); 2 in aon && 3 in aon), he_c.states))
end

@testset "cumulant_expansion preserves sum scope on Wick terms" begin
    # `cumulant_expansion(average(Σ_i …), 2)` keeps Σ_i on every factorised Wick
    # term that references i, rebuilt through the canonical `average(SQA.Σ(...))`
    # form so the factored sum is isequal to the natural one.
    @variables N
    hc = FockSpace(:cavity); ha = NLevelSpace(:atom, 2); h = hc ⊗ ha
    a = Destroy(h, :a, 1)
    σ(α, β, idx) = IndexedOperator(Transition(h, :σ, α, β, 2), idx)
    i = Index(h, :i, N, ha)

    R = Σ(a' * a * σ(1, 2, i), i)            # 3-op product under a sum
    got = cumulant_expansion(average(R), 2)

    # Correct order-2 factorisation: the sum scope stays on the block that
    # contains σ_i, i.e. Σ_i is present on every Wick term that references i.
    want = average(a) * average(Σ(a' * σ(1, 2, i), i)) +
        average(Σ(σ(1, 2, i), i)) * average(a' * a) +
        average(a') * average(Σ(a * σ(1, 2, i), i)) +
        -2 * average(a) * average(a') * average(Σ(σ(1, 2, i), i))

    @test _iz(got - want)
end

@testset "mix_choice is system state (Stage 1)" begin
    @variables Δ::Real g::Real κ::Real γ::Real N::Real
    hc = FockSpace(:cavity); ha = NLevelSpace(:atom, 2); h = hc ⊗ ha
    a = Destroy(h, :a, 1)
    σ(x, y, k) = IndexedOperator(Transition(h, :σ, x, y, 2), k)
    i = Index(h, :i, N, ha)
    H = Δ * a' * a + g * (a' * Σ(σ(1, 2, i), i) + a * Σ(σ(2, 1, i), i))
    J = [a, σ(1, 2, i)]
    rates = [κ, γ]

    eqs_min = meanfield([a' * a], H, J; rates = rates, order = [2, 1], mix_choice = minimum)
    @test eqs_min.graph.sys.mix_choice === minimum
    # scale/evaluate now preserve the system's mix_choice (the latent bug is fixed).
    eqs_c = complete(eqs_min)
    @test eqs_c.graph.sys.mix_choice === minimum
    @test scale(eqs_c).graph.sys.mix_choice === minimum
    @test evaluate(eqs_c; limits = (N => 3)).graph.sys.mix_choice === minimum

    # Noise-to-det conversion preserves mix_choice.
    eqs_noise = meanfield(
        [a' * a], H, J; rates = rates, efficiencies = [0.5, 0.5],
        order = [2, 1], mix_choice = minimum,
    )
    @test MeanfieldEquations(eqs_noise).graph.sys.mix_choice === minimum
end
