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

@testset "indexed meanfield: collective atomic emission" begin
    hc = FockSpace(:cavity); ha = NLevelSpace(:atom, 2); h = hc ⊗ ha
    @qnumbers a::Destroy(h, 1)
    σ(x, y, k) = IndexedOperator(Transition(h, :σ, x, y, 2), k)
    @variables N Δ g κ γ
    i = Index(h, :i, N, ha)
    H = Δ * a' * a + Σ(g * (a * σ(2, 1, i) + a' * σ(1, 2, i)), i)
    eqs = meanfield([a], H, [a, σ(1, 2, i)]; rates = [κ, γ])
    @test eqs isa MeanfieldEquations
    # `meanfield` derives one equation per seed operator; one seed (`a`) → 1.
    @test length(eqs.equations) == 1
end

@testset "indexed average appears in RHS" begin
    hc = FockSpace(:cavity); ha = NLevelSpace(:atom, 2); h = hc ⊗ ha
    @qnumbers a::Destroy(h, 1)
    σ(x, y, k) = IndexedOperator(Transition(h, :σ, x, y, 2), k)
    @variables N Δ g κ γ
    i = Index(h, :i, N, ha)
    H = Δ * a' * a + Σ(g * (a * σ(2, 1, i) + a' * σ(1, 2, i)), i)
    eqs = meanfield([a], H, [a, σ(1, 2, i)]; rates = [κ, γ])
    @test occursin("σ", repr(eqs.equations[1].rhs))
end

@testset "indexed meanfield: two-atom collective JC drift" begin
    @variables Δc::Real η::Real Δa::Real κ::Real Γ::Real
    N = 2
    hc = FockSpace(:cavity)
    ha = NLevelSpace(Symbol(:atom), 2)
    h = hc ⊗ ha

    i_ind = Index(h, :i, N, ha)
    k_ind = Index(h, :k, N, ha)

    g(k) = IndexedVariable(:g, k)

    @qnumbers a::Destroy(h)
    σ(i, j, k) = IndexedOperator(Transition(h, :σ, i, j), k)

    Hc = Δc * a' * a + η * (a' + a)
    Ha = Δa * Σ(σ(2, 2, i_ind), i_ind)
    Hi = Σ(g(i_ind) * (a' * σ(1, 2, i_ind) + a * σ(2, 1, i_ind)), i_ind)
    H = Hc + Ha + Hi

    J = [a, σ(1, 2, i_ind)]
    rates = [κ, Γ]

    ops = [a, σ(2, 2, k_ind), σ(1, 2, k_ind)]
    eqs = meanfield(ops, H, J; rates = rates, order = 2)
    # `meanfield` derives one equation per seed; three seeds (`a`, `σ22`,
    # `σ12`) → 3.
    @test length(eqs.equations) == 3
    @test eqs isa MeanfieldEquations
    ceqs = complete(eqs)
    @test ceqs isa MeanfieldEquations
    @test length(ceqs.equations) >= 3

    # `complete` is idempotent on a closed system.
    ceqs_again = complete(ceqs)
    @test length(ceqs_again.equations) == length(ceqs.equations)

    # Every index in the completed system traces back to the user's declared
    # vocabulary, including slot reps `Symbol(base, "_", k)` derived from a
    # declared index.
    declared = Set([SecondQuantizedAlgebra.index_name(i_ind), SecondQuantizedAlgebra.index_name(k_ind)])
    function _is_user_derived(name::Symbol)
        name in declared && return true
        s = String(name)
        for d in declared
            ds = String(d) * "_"
            startswith(s, ds) || continue
            tail = s[(length(ds) + 1):end]
            isempty(tail) && continue
            all(isdigit, tail) && return true
        end
        return false
    end
    for op in ceqs.operators
        for idx in SecondQuantizedAlgebra.get_indices(op)
            @test _is_user_derived(SecondQuantizedAlgebra.index_name(idx))
        end
    end
end

@testset "indexed jumps: per-jump dissipator is summed over jump index" begin
    # Indexed jumps induce independent-decay semantics (Σ_i D[J_i] applied
    # separately for each atom), so the i ≠ j contribution vanishes by
    # commutation and no cross-site σ_i*σ_j products survive on the RHS.
    hc = FockSpace(:cavity); ha = NLevelSpace(:atom, 2); h = hc ⊗ ha
    @qnumbers a::Destroy(h)
    σ(α, β, k) = IndexedOperator(Transition(h, :σ, α, β), k)
    @variables N Δ R Γ ν
    g(k) = IndexedVariable(:g, k)
    i = Index(h, :i, N, ha); j = Index(h, :j, N, ha)

    H = -Δ * a' * a + Σ(g(i) * (a' * σ(1, 2, i) + a * σ(2, 1, i)), i)
    J = [σ(1, 2, i), σ(2, 1, i), σ(2, 2, i)]
    rates = [Γ, R, ν]

    eqs = meanfield([σ(2, 2, j)], H, J; rates = rates, order = 2)
    rhs = eqs.equations[1].rhs

    # The dissipator collapses to
    #   R + (-R - Γ) ⟨σ_j^{22}⟩ + i g_j ⟨a' σ_j^{12}⟩ - i g_j ⟨a σ_j^{21}⟩.
    expected = R + (-R - Γ) * average(σ(2, 2, j)) +
        Symbolics.IM * g(j) * average(a' * σ(1, 2, j)) -
        Symbolics.IM * g(j) * average(a * σ(2, 1, j))
    diff = rhs - expected
    @test SymbolicUtils._iszero(SymbolicUtils.simplify(diff; expand = true))
end

@testset "indexed jumps: σ_jj equation has no σ²/σ³ leak after scale" begin
    hc = FockSpace(:cavity); ha = NLevelSpace(:atom, 2); h = hc ⊗ ha
    @qnumbers a::Destroy(h)
    σ(α, β, k) = IndexedOperator(Transition(h, :σ, α, β), k)
    @variables N Δ κ Γ R ν g
    i = Index(h, :i, N, ha); j = Index(h, :j, N, ha)
    g_idx(k) = IndexedVariable(:g, k)

    H = -Δ * a' * a + Σ(g_idx(i) * (a' * σ(1, 2, i) + a * σ(2, 1, i)), i)
    J = [a, σ(1, 2, i), σ(2, 1, i), σ(2, 2, i)]
    rates = [κ, Γ, R, ν]

    # `phase_invariant` is the shared U(1) filter from runtests.jl `init_code`.
    eqs = meanfield([a' * a, σ(2, 2, j)], H, J; rates = rates, order = 2)
    eqs_c = complete(eqs; filter_func = phase_invariant)
    eqs_sc = scale(eqs_c)

    # Pick the σ_{j_1}^{22} equation. `scale` renames free indices to suffixed
    # slots; `j` is the only free LHS atom index, so canonical slot 1 is `j(1)`.
    j1 = j(1)
    σ22_op = average(σ(2, 2, j1))
    σ22_idx = findfirst(eq -> isequal(eq.lhs, σ22_op), eqs_sc.equations)
    @test σ22_idx !== nothing
    σ22_rhs = eqs_sc.equations[σ22_idx].rhs

    # `IndexedVariable(:g, i)` flattens to the scalar `g` under scale.
    expected = R + (-R - Γ) * average(σ(2, 2, j1)) +
        Symbolics.IM * g * average(a' * σ(1, 2, j1)) -
        Symbolics.IM * g * average(a * σ(2, 1, j1))
    diff = σ22_rhs - expected
    @test SymbolicUtils._iszero(SymbolicUtils.simplify(diff; expand = true))
end


@testset "indexed meanfield: alpha-equivalent leaves are one state" begin
    # With `H` summed over `i_ind` and `ops` using `k_ind`, the leaf
    # ⟨σ_{i_ind}^{12}⟩ and the user-declared state ⟨σ_{k_ind}^{12}⟩ denote the
    # same observable; `find_missing` must not introduce a duplicate state.
    @variables Δc::Real η::Real Δa::Real κ::Real Γ::Real
    N = 2
    hc = FockSpace(:cavity); ha = NLevelSpace(:atom, 2); h = hc ⊗ ha
    i_ind = Index(h, :i, N, ha); k_ind = Index(h, :k, N, ha)
    @qnumbers a::Destroy(h)
    σ(i, j, k) = IndexedOperator(Transition(h, :σ, i, j), k)
    g(k) = IndexedVariable(:g, k)

    H = Δc * a' * a + η * (a' + a) + Δa * Σ(σ(2, 2, i_ind), i_ind) +
        Σ(g(i_ind) * (a' * σ(1, 2, i_ind) + a * σ(2, 1, i_ind)), i_ind)
    ops = [a, σ(2, 2, k_ind), σ(1, 2, k_ind)]
    J = [a, σ(1, 2, i_ind)]
    eqs = meanfield(ops, H, J; rates = [κ, Γ], order = 2)

    # `find_missing` must not report the i-indexed single-σ states as missing;
    # they are alpha-equivalent to the declared k-indexed ones.
    missing_after_first = QuantumCumulants.find_missing(eqs)
    for m in missing_after_first
        s = repr(m)
        @test !occursin(r"⟨σ_i₁₂⟩", s)
        @test !occursin(r"⟨σ_i₂₂⟩", s)
        @test !occursin(r"⟨σ_i₂₁⟩", s)
    end
end

@testset "collective indexed dissipation: N=1 reduces to single-atom decay" begin
    # A collective `DoubleIndexedVariable` rate Γ(i,j) on a single-index jump must
    # give ordinary single-atom decay at N=1 — population at Γ(1,1), coherence at
    # Γ(1,1)/2 — for any level count. The explicit diagonal Γ(i,i) D[σ_i] supplies
    # it; an off-diagonal-only emission gave identically 0 decay (the silent bug).
    num_coeff(eq) = begin
        sub = Dict{Any, Any}(eq.lhs => 1.0)
        for v in Symbolics.get_variables(eq.rhs)
            sub[v] = 1.0
        end
        Symbolics.value(Symbolics.substitute(eq.rhs, sub))
    end
    for nlev in (2, 3)
        h = NLevelSpace(:atom, nlev)
        @variables N
        i = Index(h, :i, N, h); j = Index(h, :j, N, h)
        σ(x, y, k) = IndexedOperator(Transition(h, :σ, x, y), k)
        eqs = meanfield(
            [σ(2, 2, i), σ(1, 2, i)], 0 * Σ(σ(2, 2, i), i), [σ(1, 2, i)];
            rates = [DoubleIndexedVariable(:Γ, i, j)], order = 1,
        )
        complete!(eqs)
        ev = evaluate(eqs; limits = (N => 1))
        pop = only(e for e in ev.equations if endswith(string(e.lhs), "₂₂⟩"))
        coh = only(e for e in ev.equations if endswith(string(e.lhs), "₁₂⟩"))
        @test isapprox(num_coeff(pop), -1.0)   # d⟨σ²²⟩/dt = -Γ(1,1)⟨σ²²⟩
        @test isapprox(num_coeff(coh), -0.5)   # d⟨σ¹²⟩/dt = -½Γ(1,1)⟨σ¹²⟩
    end
end

@testset "collective indexed dissipation: off-diagonal recycling keeps j≠LHS" begin
    # The off-diagonal recycling Σ_{j≠i} must keep its constraint through cumulant
    # factorisation in the indexed-sum scope; otherwise it re-admits the diagonal j=i
    # and double-counts the self-decay (the cavity ~0.2% gap).
    h = NLevelSpace(:atom, 2)
    @variables N
    i = Index(h, :i, N, h); j = Index(h, :j, N, h)
    σ(x, y, k) = IndexedOperator(Transition(h, :σ, x, y), k)
    eqs = meanfield(
        [σ(1, 2, i)], 0 * Σ(σ(2, 2, i), i), [σ(1, 2, i)];
        rates = [DoubleIndexedVariable(:Γ, i, j)], order = 1,
    )
    complete!(eqs)
    rhs = only(e for e in eqs.equations if endswith(string(e.lhs), "₁₂⟩")).rhs
    found = Ref(false)
    walk(x) = begin
        x isa SymbolicUtils.BasicSymbolic || return
        if SecondQuantizedAlgebra.is_indexed_sum(x)
            isempty(SecondQuantizedAlgebra.get_sum_non_equal(x)) || (found[] = true)
            return
        end
        SymbolicUtils.iscall(x) && foreach(walk, SymbolicUtils.arguments(x))
    end
    walk(SymbolicUtils.unwrap(rhs))
    @test found[]
end

@testset "collective indexed dissipation: N=2 equals explicit collective Lindblad" begin
    # Lock the FULL mean-field of Σ_{i,j} Γ(i,j) D[σ_i¹², σ_j¹²] at N=2 against the
    # hand-derived exact collective Lindblad (master-independent). Directly tests the
    # diagonal+off-diagonal decomposition with j≠k preserved:
    #   d⟨σ_k¹²⟩ = -½Γ_kk⟨σ_k¹²⟩ + Σ_{j≠k} Γ_kj (⟨σ_k²²⟩ - ½) ⟨σ_j¹²⟩
    #   d⟨σ_k²²⟩ = -Γ_kk⟨σ_k²²⟩ - ½Σ_{i≠k} Γ_ik ⟨σ_i²¹⟩⟨σ_k¹²⟩ - ½Σ_{j≠k} Γ_kj ⟨σ_k²¹⟩⟨σ_j¹²⟩
    # Re-admitting the diagonal j=k (dropped NE) or losing the explicit self-decay
    # would change these values; the asymmetric Γ catches index transposition too.
    h = NLevelSpace(:atom, 2)
    @variables N
    i = Index(h, :i, N, h); j = Index(h, :j, N, h)
    σ(x, y, k) = IndexedOperator(Transition(h, :σ, x, y), k)
    eqs = meanfield(
        [σ(2, 2, i), σ(1, 2, i)], 0 * Σ(σ(2, 2, i), i), [σ(1, 2, i)];
        rates = [DoubleIndexedVariable(:Γ, i, j)], order = 1,
    )
    complete!(eqs)
    ev = evaluate(eqs; limits = (N => 2))

    av = Set{Any}()
    walk(x) = begin
        x isa SymbolicUtils.BasicSymbolic || return
        if SecondQuantizedAlgebra.is_average(x)
            push!(av, x)
        elseif SymbolicUtils.iscall(x)
            foreach(walk, SymbolicUtils.arguments(x))
        end
    end
    for e in ev.equations
        walk(SymbolicUtils.unwrap(e.rhs))
    end
    byname(sfx) = only(a for a in av if endswith(string(a), sfx))
    s12 = [byname("_$(k)₁₂⟩") for k in 1:2]
    s21 = [byname("_$(k)₂₁⟩") for k in 1:2]
    s22 = [byname("_$(k)₂₂⟩") for k in 1:2]
    pars = unique(reduce(vcat, [collect(Symbolics.get_variables(e.rhs)) for e in ev.equations]))
    Γ = [only(p for p in pars if string(p) == "Γ[$a, $b]") for a in 1:2, b in 1:2]

    vals = Dict{Any, Any}()
    for (n, x) in enumerate(vcat(s12, s21, s22))
        vals[x] = 0.1n
    end
    for a in 1:2, b in 1:2
        vals[Γ[a, b]] = float(2a + b)   # asymmetric so a transposed index is caught
    end
    v(x) = vals[x]
    rhsval(e) = Symbolics.value(Symbolics.substitute(e.rhs, vals))
    eqof(sfx) = only(e for e in ev.equations if endswith(string(e.lhs), sfx))

    for k in 1:2
        o = 3 - k
        exp_coh = -0.5 * v(Γ[k, k]) * v(s12[k]) +
            v(Γ[k, o]) * (v(s22[k]) - 0.5) * v(s12[o])
        exp_pop = -v(Γ[k, k]) * v(s22[k]) -
            0.5 * v(Γ[o, k]) * v(s21[o]) * v(s12[k]) -
            0.5 * v(Γ[k, o]) * v(s21[k]) * v(s12[o])
        @test isapprox(rhsval(eqof("_$(k)₁₂⟩")), exp_coh)
        @test isapprox(rhsval(eqof("_$(k)₂₂⟩")), exp_pop)
    end
end
