const TruncOrder = Union{Nothing, Int, Vector{Int}}

function average_and_truncate(R::QAdd, order::TruncOrder, mix_choice, ctx::CanonCtx)
    acc = 0
    for (term, c) in R.arguments
        c = _coeff_num(c)
        # Drop unsimplified zero coefficients (e.g. `λ/2 - (1//2)λ`) that would seed phantom
        # moments that never close.
        # TODO: check if this can be removed after https://github.com/qojulia/SecondQuantizedAlgebra.jl/pull/167
        _iszero_coeff(c) && continue
        # Index-dependent coefficients (Γ(i,j)/Ω(i,j)) must ride inside the sum so the diagonal
        # split substitutes them (including pure c-number terms; see `_scoped_average_coeff`);
        # scalar coefficients stay outside.
        if !isempty(_coeff_scope_indices(c, R.indices))
            acc = acc + cumulant_expansion(
                _scoped_average_coeff(c, term.ops, term.ne, R.indices), order; mix_choice,
            )
        else
            acc = acc + _im_form(_truncate_coeff(c, order, mix_choice)) *
                _truncate_term(term.ops, term.ne, R.indices, order, mix_choice)
        end
    end
    return acc
end

"""Sum-scope indices the coefficient `c` depends on (empty for scalar coefficients)."""
function _coeff_scope_indices(c, scope)
    isempty(scope) && return SQA.Index[]
    vars = _coeff_vars(c)
    isempty(vars) && return SQA.Index[]
    out = SQA.Index[]
    for idx in scope
        isym = SymbolicUtils.unwrap(SQA.index_sym(idx))
        any(v -> isequal(v, isym), vars) && push!(out, idx)
    end
    return out
end

"""
Variables a coefficient depends on. `Num`/`Complex{Num}` are both `<: Number`, so test
the symbolic carrier explicitly rather than dispatching on `Number`.
"""
function _coeff_vars(c)
    cc = c isa Complex ? (real(c) + imag(c)) : c
    (cc isa Symbolics.Num || cc isa SymbolicUtils.BasicSymbolic) || return ()
    return Symbolics.get_variables(cc)
end

_coeff_num(c) = Complex(real(c), imag(c))

"""
Average a block with its coefficient inside the sum, scoping over indices used by the
operators or the coefficient so the diagonal split substitutes the coefficient.
"""
function _scoped_average_coeff(c, ops::AbstractVector{<:SQA.QSym}, non_equal, scope)
    if isempty(ops)
        # Pure c-number term whose coefficient still carries a summation index (the `σ^gg →
        # 1 − σ^ee` fold turned a scoped `Σ_k c·σ^gg_l` into a scope-less `Σ_k c·1`). Wrap the
        # coefficient back in its `Σ` so `evaluate` unrolls the index; the scalar body needs
        # SQA's diagonal split (`_accumulate_with_diag!`, SQA ≥ 0.9.3) to keep the off-diagonal
        # coefficient exact. See `_carry_non_equal` for the QC half of the same over-count (#198).
        block_scope = _coeff_scope_indices(c, scope)
        isempty(block_scope) && return _im_form(c)
        ne = _scope_non_equal(non_equal, block_scope, scope)
        return _make_indexed_sum(SymbolicUtils.unwrap(_im_form(c)), block_scope, ne)
    end
    cblock = c * reduce(*, ops)
    cblock isa QAdd || return average(cblock)
    # Carry the non-equal index constraints on `cblock` (always a QAdd); a single-op
    # block reduces to a QSym with no such slot, losing a `Σ_{j≠ext}` constraint.
    isempty(non_equal) || (cblock = _carry_non_equal(cblock, non_equal, scope))
    used = Set{SQA.Index}()
    for (t, _) in cblock.arguments, o in t.ops
        SQA.has_index(o.index) && push!(used, o.index)
    end
    for idx in _coeff_scope_indices(c, scope)
        push!(used, idx)
    end
    block_scope = SQA.Index[i for i in scope if i in used]
    isempty(block_scope) && return average(cblock)
    return average(SQA.Σ(cblock, block_scope[1], block_scope[2:end]...))
end

"""
True when the `Complex{Num}` coefficient is structurally zero. Only a sum can hide an
un-cancelled zero (`λ/2 - (1//2)λ`), so just those pay for `expand`;
products/symbols/numbers stay on the cheap `iszero` hot path.
TODO: check if this can be removed after
https://github.com/qojulia/SecondQuantizedAlgebra.jl/pull/167
"""
_iszero_coeff(c::Complex) = _iszero_part(real(c)) && _iszero_part(imag(c))
_iszero_coeff(c) = _iszero_part(c)

function _iszero_part(p)
    u = p isa Symbolics.Num ? SymbolicUtils.unwrap(p) : p
    is_sum = u isa SymbolicUtils.BasicSymbolic && SymbolicUtils.iscall(u) &&
        SymbolicUtils.operation(u) === (+)
    return is_sum ? iszero(Symbolics.expand(p)) : iszero(p)
end

"""
Convert `Complex` coefficient literals to the factored `re + im*Symbolics.IM` form
(SymbolicUtils won't unify a raw `complex(0, …)` with what `average` emits).
"""
_im_form(c) = c
_im_form(c::Complex) = real(c) + imag(c) * Symbolics.IM

"""
Truncate any averages carried by a coefficient to `order`. Forward coefficients are plain
parameters with no averages and pass straight through; the backward trace term enters as an
average-valued coefficient (`⟨J†J⟩ − ⟨JJ†⟩`) that must respect `order` like the operator block.
"""
function _truncate_coeff(c, order, mix_choice)
    if c isa Complex
        return _truncate_coeff(real(c), order, mix_choice) +
            _truncate_coeff(imag(c), order, mix_choice) * im
    end
    u = c isa Symbolics.Num ? SymbolicUtils.unwrap(c) : c
    (u isa SymbolicUtils.BasicSymbolic && _has_average(u)) || return c
    return cumulant_expansion(c, order; mix_choice)
end
average_and_truncate(R::SQA.QField, order::TruncOrder, mix_choice, ::CanonCtx) =
    order === nothing ? average(R) : cumulant_expansion(average(R), order; mix_choice)

function _truncate_term(ops::AbstractVector{<:SQA.QSym}, non_equal, scope, order, mix_choice)
    # Average WITH the sum scope first so SQA's diagonal split collapses same-index
    # operator pairs, THEN cumulant-truncate; truncating the raw product first splits
    # them into different blocks and the collapse never fires.
    return cumulant_expansion(_scoped_average(ops, non_equal, scope), order; mix_choice)
end

"""
Average a single block, attaching only the scope indices its ops use, routed through
`SQA.Σ` so the off-diagonal/diagonal split runs.
"""
function _scoped_average(ops::AbstractVector{<:SQA.QSym}, non_equal, scope)
    isempty(ops) && return 1   # empty operator product is the identity, ⟨I⟩ = 1
    block = reduce(*, ops)
    block isa QAdd && !isempty(non_equal) && (block = _carry_non_equal(block, non_equal, scope))
    used = Set{SQA.Index}()
    for o in ops
        SQA.has_index(o.index) && push!(used, o.index)
    end
    block_scope = SQA.Index[i for i in scope if i in used]
    isempty(block_scope) && return average(block)
    return average(SQA.Σ(block, block_scope[1], block_scope[2:end]...))
end

"""
Reattach the term's non-equal index pairs so the diagonal split sees them: keep any pair that
references an operator index on the block. This includes a `(scope_idx, block_idx)` pair (the
scope index is about to be summed and its partner sits on the block): dropping it lets the
subsequent `SQA.Σ` re-split a diagonal that upstream already peeled off, double-counting it.
A pair touching no block index is irrelevant to this block's sum and is dropped.
"""
function _carry_non_equal(block::QAdd, non_equal, scope)
    present = Set{SQA.Index}()
    for (term, _) in block.arguments, o in term.ops
        SQA.has_index(o.index) && push!(present, o.index)
    end
    kept = Tuple{SQA.Index, SQA.Index}[
        p for p in non_equal if p[1] in present || p[2] in present
    ]
    isempty(kept) && return block
    out = SQA.QTermDict()
    for (term, c) in block.arguments
        out[SQA.QTerm(copy(term.ops), vcat(term.ne, kept))] = c
    end
    return SQA.QAdd(out, block.indices)
end

"""
Moment-level ground-projector reduction: rewrite each averaged moment containing a
ground-state projector `σ^gg` via the completeness identity `σ^gg = 1 − Σ_{k≠g} σ^kk`,
canonicalising onto the minimal `N−1` population basis. A no-op without a ground projector.
"""
_reduce_ground_in_drift(x::Symbolics.Num) =
    Symbolics.Num(_reduce_ground_in_drift(SymbolicUtils.unwrap(x)))
function _reduce_ground_in_drift(x)
    return mapleaves(x) do leaf
        op = undo_average(leaf)
        op isa QAdd || return leaf
        expanded = SQA.expand_completeness(op)
        isequal(expanded, op) ? leaf : SymbolicUtils.unwrap(average(expanded))
    end
end

struct NodeData
    drift::Symbolics.Num                      # faithful averaged-and-truncated RHS
    op_drift::QAdd                            # operator RHS (latex / inspection / re-truncation)
    noise::Union{Nothing, Symbolics.Num}      # averaged + truncated noise drift (optional)
    op_noise::Union{Nothing, Symbolics.Num}   # operator-level noise form: deferred (always `nothing`)
    order::Int                                # cached
    aon::Vector{Int}                          # cached acts_on
end

"""
Derive the averaged-and-truncated drift (and noise) for one operator: build the
operator RHS via `_operator_rhs`, apply the `σ^gg` fold, then cross to moments
via `average_and_truncate`.
"""
function derive(op::QAdd, sys, ctx::CanonCtx)
    op_drift = _operator_rhs(
        sys.direction, op, im * sys.hamiltonian,
        sys.jumps, sys.jumps_dagger, sys.rates
    )
    op_drift = SQA.expand_completeness(op_drift)
    # Assert the LHS operator's free atom-space indices distinct (a multi-atom moment
    # ⟨σ_i σ_j⟩ is a distinct-slot cumulant, i≠j); else diagonal-split terms fail to
    # collapse and leak spurious higher-order cumulants.
    op_drift = _assume_distinct_atom_indices(op_drift, _distinct_atom_indices([op]))
    drift = Symbolics.Num(average_and_truncate(op_drift, sys.order, sys.mix_choice, ctx))
    drift = _reduce_ground_in_drift(drift)

    if sys.efficiencies === nothing
        op_noise = nothing
        noise = nothing
    else
        # `_noise_builder` returns (operator, averaged) noise eqs; the operator-level
        # drift is deferred (mixed operator/average), so truncate the averaged one.
        _, noise_eqs = _noise_builder(sys.direction)(
            [op], sys.jumps,
            sys.jumps_dagger, sys.rates, sys.efficiencies
        )
        op_noise = nothing
        noise_rhs = noise_eqs[1].rhs
        noise = Symbolics.Num(
            sys.order === nothing ? noise_rhs :
                cumulant_expansion(noise_rhs, sys.order; mix_choice = sys.mix_choice),
        )
        noise = _reduce_ground_in_drift(noise)
    end

    return NodeData(drift, op_drift, noise, op_noise, get_order(op), SQA.acts_on(op))
end
