using QuantumCumulants
using Test

@testset "mixed-order" begin

# Hilbertspace
hf = FockSpace(:cavity)
ha = NLevelSpace(:atom, 2)
h = hf⊗ha

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

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

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

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

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

# N-atom laser
# Parameters
N = 2 #number of atoms
κ, g, Γ23, Γ13, Γ12, Ω, Δc, Δ3 = cnumbers("κ g Γ_{23} Γ_{13} Γ_{12} Ω Δ_c Δ_3")

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

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

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

# Jumps
J = [a;[σ(1,2,i) for i=1:N];[σ(1,3,i) for i=1:N];[σ(2,3,i) for i=1:N]]

# Rates
rates = [κ;[Γ12 for i=1:N];[Γ13 for i=1:N];[Γ23 for i=1:N]]

# list of operators
ops = [a'a, σ(2,2,1), σ(3,3,1)]

he = meanfield(ops,H,J; rates=rates)
he_avg_ = cumulant_expansion(he,[2,1,1]) #second order average

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

he_avg_ = cumulant_expansion(he,[2,1,1];mix_choice=minimum)
he_avg = complete(he_avg_;order=[2,1,1],mix_choice=minimum)
@test isempty(find_missing(he_avg))
@test isempty(findall(x -> (aon = acts_on(x); 2 in aon && 3 in aon), he_avg.states))
@test isempty(findall(x -> (aon = acts_on(x); 1 in aon && 3 in aon), he_avg.states))
@test isempty(findall(x -> (aon = acts_on(x); 1 in aon && 2 in aon), he_avg.states))

end # testset
