# Concrete graph-backed equation wrappers. Included after graph.jl so the `graph` field can
# be typed `MomentGraph{Sys}`. The wrapper carries the authoritative `MomentGraph` as its
# source of truth; the equation/state/operator vectors are a derived view regenerated by
# `resync!` from the graph.

"""
    MeanfieldEquations

Concrete equation set produced by [`meanfield`](@ref) when no measurement backaction is
requested. The authoritative cumulant hierarchy is `graph`; the vector fields are a derived
view regenerated from it.

# Fields
* `graph`: the source of truth (moments, drifts, treatments, `sys`, `ctx`).
* `equations`: the averaged differential equations (derived view).
* `operator_equations`: the same equations at the operator level.
* `states`: the averages on the left-hand sides.
* `operators`: the operators on the left-hand sides.
* `hamiltonian`: the system Hamiltonian.
* `jumps`, `jumps_dagger`: the collapse operators and their adjoints.
* `rates`: the decay rates corresponding to `jumps`.
* `iv`: the independent (time) variable.
* `order`: the cumulant-expansion order, or `nothing`.
* `direction`: [`Forward`](@ref) or [`Backward`](@ref) evolution.
"""
mutable struct MeanfieldEquations{
        O <: Union{Nothing, Vector{Int}},
        H <: QField,
        Op <: QField,
        Jt,
        Jdt,
        R,
        S <: SymbolicUtils.BasicSymbolic,
        D <: EvolutionDirection,
        Sys <: SystemSpec,
    } <: AbstractMeanfieldEquations
    graph::MomentGraph{Sys}                  # source of truth: moments, drifts, treatments, sys, ctx
    equations::Vector{Symbolics.Equation}    # derived view (regenerated from `graph`)
    operator_equations::Vector{Symbolics.Equation}
    states::Vector{S}
    operators::Vector{Op}
    hamiltonian::H                           # immutable system data, copied from `graph.sys`
    jumps::Vector{Jt}
    jumps_dagger::Vector{Jdt}
    rates::Vector{R}
    iv::Symbolics.Num
    order::O
    direction::D

    function MeanfieldEquations(graph::MomentGraph{Sys}) where {Sys}
        sys = graph.sys
        sys.efficiencies === nothing || error(
            "MeanfieldEquations(::MomentGraph) requires a deterministic graph " *
                "(efficiencies === nothing); use NoiseMeanfieldEquations for a noise graph.",
        )
        v = det_view_columns(graph)
        ham = sys.hamiltonian
        jumps = collect(sys.jumps)
        jumps_dagger = collect(sys.jumps_dagger)
        rates = collect(sys.rates)
        return new{
            typeof(sys.order), typeof(ham), eltype(v.operators),
            eltype(jumps), eltype(jumps_dagger), eltype(rates),
            eltype(v.states), typeof(sys.direction), Sys,
        }(
            graph, v.equations, v.operator_equations, v.states, v.operators,
            ham, jumps, jumps_dagger, rates, sys.iv, sys.order, sys.direction,
        )
    end
end

"""
    NoiseMeanfieldEquations

Concrete equation set produced by [`meanfield`](@ref) when `efficiencies` is supplied,
adding a measurement-backaction noise drift to the deterministic equations. The
`direction` tag ([`Forward`](@ref) or [`Backward`](@ref)) drives compile-time dispatch of
the noise assembly. Like [`MeanfieldEquations`](@ref) it is graph-backed: `graph` is the
source of truth and the vector fields are a derived view.

# Fields
* `graph`: the source of truth.
* `equations`: the deterministic averaged differential equations (derived view).
* `noise_equations`: the averaged noise (Brownian) drift per state.
* `operator_equations`, `operator_noise_equations`: the operator-level counterparts.
* `states`: the averages on the left-hand sides.
* `operators`: the operators on the left-hand sides.
* `hamiltonian`: the system Hamiltonian.
* `jumps`, `jumps_dagger`: the collapse operators and their adjoints.
* `rates`: the decay rates corresponding to `jumps`.
* `efficiencies`: the detector efficiencies per jump.
* `iv`: the independent (time) variable.
* `order`: the cumulant-expansion order, or `nothing`.
* `direction`: [`Forward`](@ref) or [`Backward`](@ref) evolution.
"""
mutable struct NoiseMeanfieldEquations{
        O <: Union{Nothing, Vector{Int}},
        H <: QField,
        Op <: QField,
        Jt,
        Jdt,
        R, E,
        S <: SymbolicUtils.BasicSymbolic,
        D <: EvolutionDirection,
        Sys <: SystemSpec,
    } <: AbstractMeanfieldEquations
    graph::MomentGraph{Sys}
    equations::Vector{Symbolics.Equation}
    noise_equations::Vector{Symbolics.Equation}
    operator_equations::Vector{Symbolics.Equation}
    operator_noise_equations::Vector{Symbolics.Equation}
    states::Vector{S}
    operators::Vector{Op}
    hamiltonian::H
    jumps::Vector{Jt}
    jumps_dagger::Vector{Jdt}
    rates::Vector{R}
    efficiencies::Vector{E}
    iv::Symbolics.Num
    order::O
    direction::D

    function NoiseMeanfieldEquations(graph::MomentGraph{Sys}) where {Sys}
        sys = graph.sys
        sys.efficiencies === nothing && error(
            "NoiseMeanfieldEquations(::MomentGraph) requires a noise graph " *
                "(efficiencies !== nothing).",
        )
        v = noise_view_columns(graph)
        ham = sys.hamiltonian
        jumps = collect(sys.jumps)
        jumps_dagger = collect(sys.jumps_dagger)
        rates = collect(sys.rates)
        efficiencies = collect(sys.efficiencies)
        return new{
            typeof(sys.order), typeof(ham), eltype(v.operators),
            eltype(jumps), eltype(jumps_dagger), eltype(rates), eltype(efficiencies),
            eltype(v.states), typeof(sys.direction), Sys,
        }(
            graph, v.equations, v.noise_equations, v.operator_equations,
            v.operator_noise_equations, v.states, v.operators,
            ham, jumps, jumps_dagger, rates, efficiencies, sys.iv, sys.order, sys.direction,
        )
    end
end

# ---- view construction / regeneration ----------------------------------------

"""Build the cold `operator_noise_equations` column from the noise builder, one per moment."""
function build_op_noise(g::MomentGraph)
    sys = g.sys
    out = Symbolics.Equation[]
    for k in keys(g.nodes)
        on, _ = _noise_builder(sys.direction)(
            [k], sys.jumps, sys.jumps_dagger, sys.rates, sys.efficiencies,
        )
        push!(out, on[1])
    end
    return out
end

"""Full view columns of a deterministic graph, for construction (`new` needs every field)."""
function det_view_columns(g::MomentGraph)
    ks = collect(keys(g.nodes))
    # `average(k)` is the LHS state of equation `k`; build it once and share it between the
    # `equations` column and the `states` column rather than averaging each key twice.
    states = SymbolicUtils.BasicSymbolic[average(k) for k in ks]
    equations = Symbolics.Equation[states[i] ~ g.nodes[ks[i]].drift for i in eachindex(ks)]
    operator_equations = Symbolics.Equation[k ~ g.nodes[k].op_drift for k in ks]
    operators = QAdd[k for k in ks]
    return (; equations, operator_equations, states, operators)
end

"""
Full view columns of a noise graph: the deterministic columns plus the noise drift and the
cold `operator_noise_equations`. The latter is rebuilt from the noise builder because
`NodeData.op_noise` is always `nothing`; nothing reads the column, so this matches the
behaviour of building it eagerly.
"""
function noise_view_columns(g::MomentGraph)
    ks = collect(keys(g.nodes))
    noise_equations = Symbolics.Equation[average(k) ~ g.nodes[k].noise for k in ks]
    return (;
        det_view_columns(g)...,
        noise_equations, operator_noise_equations = build_op_noise(g),
    )
end

refill!(dst, src) = (empty!(dst); append!(dst, src); dst)

"""
Regenerate the view in place from `eqs.graph` (the sole writer of the view besides the
constructor). The drift columns a rewrite can touch are always rebuilt; `moments_unchanged`
skips the moment-set columns (`operator_equations`, `states`, `operators`, and the costly
`operator_noise_equations`) so a drift-only rewrite stays cheap.
"""
function resync!(eqs::MeanfieldEquations; moments_unchanged::Bool = false)
    g = eqs.graph
    ks = collect(keys(g.nodes))
    refill!(eqs.equations, (average(k) ~ g.nodes[k].drift for k in ks))
    if !moments_unchanged
        refill!(eqs.operator_equations, (k ~ g.nodes[k].op_drift for k in ks))
        refill!(eqs.states, (average(k) for k in ks))
        refill!(eqs.operators, ks)
    end
    return eqs
end

function resync!(eqs::NoiseMeanfieldEquations; moments_unchanged::Bool = false)
    g = eqs.graph
    ks = collect(keys(g.nodes))
    refill!(eqs.equations, (average(k) ~ g.nodes[k].drift for k in ks))
    refill!(eqs.noise_equations, (average(k) ~ g.nodes[k].noise for k in ks))
    if !moments_unchanged
        refill!(eqs.operator_equations, (k ~ g.nodes[k].op_drift for k in ks))
        refill!(eqs.operator_noise_equations, build_op_noise(g))
        refill!(eqs.states, (average(k) for k in ks))
        refill!(eqs.operators, ks)
    end
    return eqs
end

# ---- noise -> deterministic conversion ---------------------------------------

"""
    MeanfieldEquations(eqs::NoiseMeanfieldEquations)

Strip the measurement-backaction columns (`noise_equations`, `operator_noise_equations`,
`efficiencies`) and return a plain [`MeanfieldEquations`](@ref) carrying only the
deterministic drift. The `direction` tag and the system's `mix_choice` are preserved.

Use this to compare a stochastic simulation against the no-measurement evolution: build the
noisy system once, then pass `MeanfieldEquations(eqs)` to `System` / `ODEProblem` for the
deterministic reference.
"""
function MeanfieldEquations(eqs::NoiseMeanfieldEquations)
    g = eqs.graph
    sys = g.sys
    # Strip the stochastic part: efficiencies -> nothing, drop each node's noise drift. The
    # deterministic drift, operator drift, treatments, ctx and `mix_choice` carry over from
    # the source graph verbatim (no view round-trip, so `op_drift` stays a faithful QAdd).
    det_sys = SystemSpec(
        sys.hamiltonian, sys.jumps, sys.jumps_dagger, sys.rates, nothing,
        sys.iv, sys.order, sys.mix_choice, sys.direction,
    )
    nodes = OrderedCollections.OrderedDict{NodeKey, NodeData}()
    for (k, nd) in g.nodes
        nodes[k] = NodeData(nd.drift, nd.op_drift, nothing, nothing, nd.order, nd.aon)
    end
    return MeanfieldEquations(MomentGraph(nodes, det_sys, g.ctx, copy(g.treatments)))
end
