"""Subspaces participating in the quotient: symmetric ones, restricted by `h` (empty = all)."""
_scale_selected(ctx::CanonCtx, h::Vector{Int}) =
    isempty(h) ? ctx.symmetric : intersect(ctx.symmetric, Set(h))

function quotient(g::MomentGraph; h::Vector{Int} = Int[])
    ctx = g.ctx
    selected = _scale_selected(ctx, h)
    sym_to_space = _build_sym_to_space(g)
    # Key nodes in the graph's treatment with selected subspaces now Scaled (not bare
    # `scaled_key`, which would alpha-rename already-Concrete subspaces and collapse them).
    treatments = copy(g.treatments)
    for sp in selected
        treatments[sp] = Scaled
    end
    # `_materialised_key` folds permutation symmetry but not conjugation; a conjugation-folded
    # graph must also collapse moments that coincide only after conjugation, else the count
    # depends on which conjugate side closure stored (#295). Key by the permutation rep (so the
    # drift matches the label), dedup by the conjugation-aware `canonical_rep`.
    fold_conj = conj_folded(g)
    seen_reps = Set{NodeKey}()
    nodes = OrderedCollections.OrderedDict{NodeKey, NodeData}()
    # A moment ⟨X⟩ recurs across many drifts; `_scale_leaf` is a pure function of the leaf
    # (given the fixed `selected`/`ctx`), so cache its result once per `quotient`. The cache
    # is local, not `ctx.cache`: the scaled form depends on `selected`, which varies with `h`.
    leaf_cache = Dict{Any, Any}()
    for (k, nd) in g.nodes
        ok = _materialised_key(k, ctx, treatments)
        haskey(nodes, ok) && continue   # permutation image already kept
        if fold_conj
            rep = canonical_rep(k, ctx; treatments)[1]
            rep in seen_reps && continue   # conjugate image already kept
            push!(seen_reps, rep)
        end
        drift = _scale_expr(nd.drift, ctx, selected, sym_to_space, leaf_cache)
        noise = nd.noise === nothing ? nothing :
            _scale_expr(nd.noise, ctx, selected, sym_to_space, leaf_cache)
        nodes[ok] = NodeData(drift, nd.op_drift, noise, nd.op_noise, nd.order, nd.aon)
    end
    return MomentGraph(nodes, g.sys, ctx, treatments)
end

"""
Symbol -> space_index map so an `IndexedVariable(:g, i)` coefficient can be traced to
its source subspace. Built from the ctx vocabulary AND the indices actually present on the
graph's moments: closure mints extra per-subspace indices (e.g. a second `i_2_2`) for
two-body moments that the frozen ctx vocab does not list, and an unmapped coefficient index
would otherwise be wrongly flattened (`_indexed_var_in_h` treats unknown as in-`h`).
"""
function _build_sym_to_space(g::MomentGraph)
    out = Dict{Symbol, Int}()
    for (sp, indices) in g.ctx.vocab, idx in indices
        out[SQA.index_name(idx)] = sp
    end
    for k in keys(g.nodes), idx in SQA.get_indices(k)
        out[SQA.index_name(idx)] = _aon(idx)
    end
    return out
end

"""
Rewrite an averaged RHS for scale in a single post-order pass: each moment unit ⟨X⟩ becomes
`prefactor * ⟨X_rep⟩`, where `X_rep` is the permutation-symmetry representative of `X` (via
`_scale_leaf`), and every `IndexedVariable(:g, i)` coefficient flattens to its scalar `g`
(all atoms share the coupling under the scale symmetry). `_scale_leaf`'s `average` can factor
an indexed-variable coefficient back out of a sum body, so the per-leaf result is flattened
before it is cached; coefficients sitting outside the moments are flattened during descent.
Memoise per leaf to skip re-folding a moment that recurs across drifts.
"""
function _scale_expr(x, ctx, selected, sym_to_space, cache)
    treatments = Dict{Int, SubspaceTreatment}()
    for sp in ctx.symmetric
        treatments[sp] = sp in selected ? Scaled : Concrete
    end
    reduced = rewrite(SymbolicUtils.unwrap(x)) do l
        _scale_node(l, ctx, selected, treatments, sym_to_space, cache)
    end
    return Symbolics.Num(reduced)
end

"""
Reduce one node under the permutation symmetry `scale` exploits: return the reduced
subtree, or `nothing` to let `rewrite` descend. A sum whose summand is a *product of several
moments* (a mean-field-factorised `Σ c·⟨A⟩⟨B⟩⋯`) needs care: rebuilding its operator would
fuse the moment product `⟨A⟩⟨B⟩` into a single higher moment `⟨A B⟩`, a different and
higher-order quantity. Such a sum is reduced by [`_scale_sum`](@ref); an ordinary moment (or
a sum over a single moment) is reduced by `_scale_leaf`, and a stray site-dependent coupling
outside a moment is replaced by its symmetric (site-independent) value.
"""
function _scale_node(l, ctx, selected, treatments, sym_to_space, cache)
    if SQA.is_indexed_sum(l) && length(eachleaf(_indexed_sum_body(l))) >= 2
        return get!(cache, l) do
            _scale_sum(l, ctx, selected, treatments, sym_to_space, cache)
        end
    end
    if _is_moment_unit(l)
        return get!(cache, l) do
            _flatten_indexed_vars_in_tree(
                _scale_leaf(l, ctx, selected, treatments), selected, sym_to_space,
            )
        end
    end
    (_is_indexed_var(l) && _indexed_var_in_h(l, selected, sym_to_space)) &&
        return _flatten_indexed_var(l)
    return nothing
end

"""
Reduce a mean-field-factorised sum `Σ_{i,j,…} c·⟨A_i⟩⟨B_j⟩⋯` under permutation symmetry.
Every atom/site is equivalent, so each single-site moment ⟨A_i⟩ equals the representative
⟨A⟩ and the sum collapses to `(number of distinct site assignments)·c·⟨A⟩⟨B⟩⋯`. That count
is the falling factorial `N(N−1)⋯` from `_scope_prefactor` (honouring any `i≠j`); each
moment folds to its representative and each site-dependent coupling takes its symmetric
value. Summation indices on subsystems that are *not* being scaled stay summed and are
re-wrapped in the residual `Σ` with their internal `≠` constraints.
"""
function _scale_sum(node, ctx, selected, treatments, sym_to_space, cache)
    indices = SQA.get_sum_indices(node)
    ne = SQA.get_sum_non_equal(node)
    body = _indexed_sum_body(node)
    keep = SQA.Index[b for b in indices if !(_aon(b) in selected)]
    pref = _scope_prefactor(indices, ne, selected)
    scaled_body = rewrite(body) do l
        _scale_node(l, ctx, selected, treatments, sym_to_space, cache)
    end
    result = pref === 1 ? scaled_body : SymbolicUtils.unwrap(pref) * scaled_body
    isempty(keep) && return result
    kept_ne = Tuple{SQA.Index, SQA.Index}[p for p in ne if p[1] in keep && p[2] in keep]
    return _make_indexed_sum(SymbolicUtils.unwrap(result), keep, kept_ne)
end

"""
Falling-factorial multiplicity `∏_b (index_range(b) − partners-already-fixed)` for the
`collapsing` indices, visited in the order of `ordered`. `ne` is any collection of index
pairs; a mutual `i≠j` between two collapsing indices is charged once (skipped for the later
of the two), while a constraint against an external or kept index is charged. Shared by the
single-moment (`_sum_scope_prefactor`) and multi-moment (`_scope_prefactor`) scale paths.
"""
function _falling_factorial(ordered, ne, collapsing::Set{SQA.Index})
    placed = Set{SQA.Index}()
    prefactor = 1
    for b in ordered
        b in collapsing || continue
        count_b = 0
        excluded = Set{SQA.Index}()
        for pair in ne
            if pair[1] == b
                other = pair[2]
            elseif pair[2] == b
                other = pair[1]
            else
                continue
            end
            # Charge a mutual `≠` once: skip a collapsed partner not yet placed.
            (other in collapsing && !(other in placed)) && continue
            other in excluded && continue
            count_b += 1
            push!(excluded, other)
        end
        prefactor = prefactor * (SQA.index_range(b) - count_b)
        push!(placed, b)
    end
    return prefactor
end

"""
How many distinct assignments the collapsing (scaled) summation indices have: the falling
factorial `N(N−1)⋯`, read straight from a sum node's `indices`/`ne`. This multiplicity is
what multiplies the representative moment when the symmetric sum collapses. Companion to
`_sum_scope_prefactor`, which derives the same count from an operator `QAdd`.
"""
_scope_prefactor(indices, ne, selected::Set{Int}) = _falling_factorial(
    indices, ne, Set{SQA.Index}(b for b in indices if _aon(b) in selected),
)

function _scale_leaf(avg, ctx, selected, treatments)
    op = undo_average(avg)
    op isa QAdd || return avg
    pref = _sum_scope_prefactor(op, selected)
    # Fold only the selected subspaces (Scaled); keep the rest verbatim (Concrete).
    # `_treatment_key` empties `.indices`, so re-attach the non-selected bound indices
    # afterwards to preserve their `Σ`.
    keep_idx = SQA.Index[b for b in op.indices if !(_aon(b) in selected)]
    folded = _treatment_key(op, ctx, treatments)
    # `_treatment_key` drops all non-equal index constraints, correct for the selected
    # subspace but wrong for a non-selected one whose off-diagonal constraint is
    # physical; re-attach the non-equal index pairs that involve any non-selected index.
    if folded isa QAdd && !isempty(keep_idx)
        kept_non_equal = Tuple{SQA.Index, SQA.Index}[]
        for (term, _) in op.arguments, p in term.ne
            (_aon(p[1]) in selected && _aon(p[2]) in selected) && continue
            p in kept_non_equal || push!(kept_non_equal, p)
        end
        reduced_op = isempty(kept_non_equal) ? SQA.QAdd(folded.arguments, keep_idx) :
            _reattach_non_equal(folded, kept_non_equal, keep_idx)
    else
        reduced_op = folded
    end
    reduced = average(reduced_op)
    pref === 1 && return reduced
    return SymbolicUtils.unwrap(pref) * reduced
end

"""
Re-attach non-equal index pairs to a folded QAdd, restricting each to terms whose ops carry both
its indices, and set the kept bound indices as the new sum scope.
"""
function _reattach_non_equal(folded::QAdd, kept_non_equal, keep_idx)
    out = SQA.QTermDict()
    for (term, c) in folded.arguments
        present = Set{SQA.Index}()
        for o in term.ops
            SQA.has_index(o.index) && push!(present, o.index)
        end
        non_equal = Tuple{SQA.Index, SQA.Index}[p for p in kept_non_equal if p[1] in present && p[2] in present]
        out[SQA.QTerm(copy(term.ops), vcat(term.ne, non_equal))] = c
    end
    return SQA.QAdd(out, keep_idx)
end

"""
Sum-scope prefactor: each collapsed bound index `b` on a selected subspace contributes
`index_range(b)` minus the distinct partners it must differ from that are *already accounted for*
(external indices, or an earlier collapsed bound index). Charging a mutual `i≠j` constraint
once gives the falling-factorial count `N(N−1)…` rather than `(N−1)^k`. Indices not used by
any `term.ops` contribute 1.
"""
function _sum_scope_prefactor(op::QAdd, selected::Set{Int})
    isempty(op.indices) && return 1
    op_indices = Set{SQA.Index}()
    for (term, _) in op.arguments, o in term.ops
        SQA.has_index(o.index) && push!(op_indices, o.index)
    end
    collapsing = Set{SQA.Index}(
        b for b in op.indices if b in op_indices && _aon(b) in selected
    )
    ne = Tuple{SQA.Index, SQA.Index}[pair for (term, _) in op.arguments for pair in term.ne]
    return _falling_factorial(op.indices, ne, collapsing)
end

"""
True for an `IndexedVariable` call (`g(i)`): its operation is a `Sym` of
`FnType{…, Real}`.
"""
function _is_indexed_var(x)
    x isa SymbolicUtils.BasicSymbolic || return false
    SymbolicUtils.iscall(x) || return false
    op = SymbolicUtils.operation(x)
    op isa SymbolicUtils.BasicSymbolic || return false
    st = SymbolicUtils.symtype(op)
    return st <: SymbolicUtils.FnType && _fn_returns_real(st)
end

_fn_returns_real(::Type{<:SymbolicUtils.FnType{A, R}}) where {A, R} = R <: Real
_fn_returns_real(::Type) = false

_flatten_indexed_var(x::SymbolicUtils.BasicSymbolic) =
    SymbolicUtils.Sym{SymbolicUtils.SymReal}(nameof(SymbolicUtils.operation(x)); type = Real)

"""
True when the `IndexedVariable`'s source index sits on a selected subspace;
unrecognised argument shapes default to `true`, flattening every coupling.
"""
function _indexed_var_in_h(x::SymbolicUtils.BasicSymbolic, selected::Set{Int}, sym_to_space::Dict{Symbol, Int})
    args = SymbolicUtils.arguments(x)
    length(args) == 1 || return true
    a = args[1]
    a isa SymbolicUtils.BasicSymbolic || return true
    SymbolicUtils.iscall(a) && return true
    sp = get(sym_to_space, Base.nameof(a), 0)
    sp == 0 && return true
    return sp in selected
end

"""
Walk the coefficient tree flattening `IndexedVariable`s, without descending into moment
leaves (averages and indexed sums, the operator moments handled by `mapleaves`).
"""
function _flatten_indexed_vars_in_tree(x, selected::Set{Int}, sym_to_space::Dict{Symbol, Int})
    return rewrite(x; descend = _descend) do y
        (_is_indexed_var(y) && _indexed_var_in_h(y, selected, sym_to_space)) ?
            _flatten_indexed_var(y) : nothing
    end
end

"""
    scale(eqs::AbstractMeanfieldEquations; h=Int[])
    scale!(eqs::AbstractMeanfieldEquations; h=Int[])

Reduce a permutation-symmetric indexed system to one representative per symmetry class,
exploiting that every atom of an indexed family acts on the system identically. `scale!`
mutates `eqs` in place; `scale` returns a new system.

# Keyword arguments
* `h=Int[]`: the Hilbert subspaces to scale, given by their `acts_on` indices. Empty
  scales every symmetric (indexed) subspace and leaves the others untouched.
"""
scale(eqs::MeanfieldEquations; h::Vector{Int} = Int[]) =
    MeanfieldEquations(quotient(eqs.graph; h))
scale(eqs::NoiseMeanfieldEquations; h::Vector{Int} = Int[]) =
    NoiseMeanfieldEquations(quotient(eqs.graph; h))

function scale!(eqs::AbstractMeanfieldEquations; h::Vector{Int} = Int[])
    eqs.graph = quotient(eqs.graph; h)
    resync!(eqs)
    return eqs
end

@doc (@doc scale) scale!
