earthsci-ast 0.6.0

use super::*;
use indexmap::IndexMap;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// A `default` / override VALUE: either a scalar, or inline row-major nested
/// array data (esm-spec §6.3, §6.6.2 — the `NumericArrayLiteral` of the schema).
///
/// A SHAPED variable's `default`, and a shaped variable's entry in a test's
/// `parameter_overrides` / `initial_conditions`, may carry the whole field as a
/// nested JSON array whose nesting matches the declared `shape` after
/// metaparameter folding; a mismatch is a load-time error. A scalar in the same
/// position keeps its broadcast meaning — the one value applies to every element.
/// This is what lets a column test supply its θ, q_v, u, v, p profiles inline
/// rather than through a per-regime generated rewrite-rule library.
///
/// `#[serde(untagged)]`, so a scalar round-trips as a bare JSON number and array
/// data as the nested array it was authored as.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum InlineValue {
    /// A single number — a 0-D value, or the broadcast value of a shaped one.
    Scalar(f64),
    /// One axis of row-major nested array data.
    Array(Vec<InlineValue>),
}

/// A flattened inline array: its per-axis extents and its values in ROW-MAJOR
/// order (last index fastest), as [`InlineValue::to_dense`] produces them.
pub type DenseInlineArray = (Vec<usize>, Vec<f64>);

impl From<f64> for InlineValue {
    fn from(v: f64) -> Self {
        InlineValue::Scalar(v)
    }
}

impl InlineValue {
    /// The scalar this value carries, or `None` when it is array data.
    pub fn as_scalar(&self) -> Option<f64> {
        match self {
            InlineValue::Scalar(v) => Some(*v),
            InlineValue::Array(_) => None,
        }
    }

    /// Whether this value is inline ARRAY data rather than a scalar.
    pub fn is_array(&self) -> bool {
        matches!(self, InlineValue::Array(_))
    }

    /// Flatten row-major nested array data into `(shape, values)` — the dense
    /// buffer plus its per-axis extents, with the values in ROW-MAJOR order
    /// (last index fastest), exactly as the JSON nesting reads.
    ///
    /// # Errors
    ///
    /// A string naming the defect when an axis is ragged (sibling arrays of
    /// different lengths or mixed scalar/array siblings) — the schema admits
    /// the shape, so the check has to live here.
    pub fn to_dense(&self) -> Result<DenseInlineArray, String> {
        let mut shape = Vec::new();
        let mut probe = self;
        while let InlineValue::Array(items) = probe {
            shape.push(items.len());
            match items.first() {
                Some(first) => probe = first,
                None => break,
            }
        }
        let mut out = Vec::new();
        flatten_inline(self, 0, &shape, &mut out)?;
        Ok((shape, out))
    }
}

/// Row-major flatten of `value` at nesting `depth`, checking every axis against
/// the `shape` the first branch established.
fn flatten_inline(
    value: &InlineValue,
    depth: usize,
    shape: &[usize],
    out: &mut Vec<f64>,
) -> Result<(), String> {
    match value {
        InlineValue::Scalar(v) => {
            if depth != shape.len() {
                return Err(format!(
                    "inline array data is ragged: a number appears at nesting depth {depth} where an array of {} elements is expected",
                    shape[depth]
                ));
            }
            out.push(*v);
            Ok(())
        }
        InlineValue::Array(items) => {
            if depth >= shape.len() {
                return Err(
                    "inline array data is ragged: an array appears where a number is expected"
                        .to_string(),
                );
            }
            if items.len() != shape[depth] {
                return Err(format!(
                    "inline array data is ragged: axis {depth} has both {} and {} elements",
                    shape[depth],
                    items.len()
                ));
            }
            for item in items {
                flatten_inline(item, depth + 1, shape, out)?;
            }
            Ok(())
        }
    }
}

/// A document-scoped index set declared in a model's `index_sets` registry
/// (RFC semiring-faq-unified-ir §5.2 / §8). Unifies ESM `domain.spatial` grid
/// dims and ESI categorical index sets under one shape. `kind` selects which
/// optional fields are meaningful (enforced by the schema's kind-conditional
/// `allOf`):
/// - `interval`    → `size`
/// - `categorical` → `members`
/// - `derived`     → `from_faq`
/// - `ragged`      → `of`, `offsets`, `values`
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct IndexSet {
    /// `"interval" | "categorical" | "derived" | "ragged"`.
    pub kind: String,

    /// Dense size for `kind: "interval"`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub size: Option<i64>,

    /// Enumerated members for `kind: "categorical"`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub members: Option<Vec<serde_json::Value>>,

    /// Source FAQ-node id for `kind: "derived"`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub from_faq: Option<String>,

    /// Members-fed-back-as-const-factor for `kind: "derived"` (projection-
    /// pushdown Phase 2b Hook 1, CONFORMANCE_SPEC §5.5.6).
    ///
    /// Names a model `parameter` const factor that the build fills with this
    /// set's materialized value-invention MEMBERS — the invented 1-based
    /// full-grid ids, `vi.members[from_faq]` — so an in-model gather
    /// `index(W, index(<member_factor>, c))` can pull the full-grid rows the
    /// compact derived axis selects. There is no `member(set, c)` IR op; this
    /// feedback IS the mechanism.
    ///
    /// `None` for every non-feedback set, so ordinary sets round-trip
    /// byte-identically. Mirrors the Julia reference's `IndexSet.member_factor`
    /// and the `member_factor` property of esm-schema.json. Without it the Rust
    /// binding silently DROPPED a normative field on parse, so a parse→emit
    /// round-trip of an overlap-gated document (the shared
    /// `overlap_gate_point_in_rect.esm` fixture among them) lost it.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub member_factor: Option<String>,

    /// Parent index sets for `kind: "ragged"`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub of: Option<Vec<String>>,

    /// CSR/length backing-factor name for `kind: "ragged"`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub offsets: Option<String>,

    /// Member backing-factor name for `kind: "ragged"`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub values: Option<String>,
}

/// Document-scoped, OPTIONAL solver hints (esm-spec §2.2).
///
/// Numerics the document knows about ITSELF, which each binding maps to its own
/// integrator. Every field is ADVISORY: a binding may ignore any or all of them
/// and still conform. Advisory governs the MECHANISM, never the OUTCOME — a
/// binding that ignores every field and still converges conforms; the
/// CONFORMANCE_SPEC §5.9 requirement to integrate successfully and agree within
/// the error band is untouched by this block and is not excused by it.
///
/// `abstol` / `reltol` are INTEGRATION tolerances and are a different quantity
/// from [`Tolerance`], which is what an assertion is COMPARED at (§6.6.4). They
/// resolve on independent chains and neither substitutes for the other.
///
/// Every field is `Option` because each is independently optional and absence
/// is NOT a default value: a document omitting `stiffness` has not declared its
/// stiffness, and a binding must not read absence as an assertion about it.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Solver {
    /// The author's declaration of the system's stiffness: `"low"`,
    /// `"moderate"` or `"high"`. A binding MAY select an implicit / BDF-family
    /// integrator on `"high"`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub stiffness: Option<String>,

    /// Absolute INTEGRATION tolerance the document asks for.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub abstol: Option<f64>,

    /// Relative INTEGRATION tolerance the document asks for.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reltol: Option<f64>,

    /// Advisory operator-splitting convention: `"none"`, `"lie"` or
    /// `"strang"`. Carries no prescribed substep structure.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub splitting: Option<String>,

    /// Reference to an external run-configuration document (esm-spec §2.2.6),
    /// held exactly as authored. EarthSciAST never opens it; the resolved
    /// location comes from [`crate::solver_external_location`] and is never
    /// written back. Arrives at esm 1.3.0.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub external: Option<String>,

    /// The interval the document is meant to be simulated over (esm-spec
    /// §2.2.7), held exactly as authored: a timestamp stays the string and a
    /// number stays the number written. Its value in binary64 seconds comes
    /// from [`crate::resolve_time_span`] and is never written back. Arrives at
    /// esm 1.3.0.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub time_span: Option<SolverTimeSpan>,
}

/// `solver.time_span` (esm-spec §2.2.7): the `{start, end}` shape of
/// `tests[].time_span` ([`TimeSpan`]), except that each bound may also be an
/// RFC 3339 UTC timestamp, so it is a separate type rather than `f64` bounds.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct SolverTimeSpan {
    /// Beginning of the interval.
    pub start: TimePoint,
    /// End of the interval; after `start` once both are converted to seconds.
    pub end: TimePoint,
}

/// One bound of a [`SolverTimeSpan`], kept in its authored form.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum TimePoint {
    /// Seconds of the independent variable, in the document's time units. A
    /// [`serde_json::Number`] so an integer bound emits as the integer written.
    Seconds(serde_json::Number),
    /// An RFC 3339 UTC timestamp (`YYYY-MM-DDThh:mm:ss(.f+)?Z`): seconds from
    /// the document's time origin to that instant, counted as POSIX time.
    Timestamp(String),
}

/// One entry of the document-scoped [`EsmFile::coordinates`] registry
/// (RFC streaming-output-sinks §8.3): marks an existing data array — or an
/// inline literal vector — as a physical coordinate and attaches CF metadata.
///
/// Exactly one of `source` and `values` is present. The coordinate's shape
/// comes from its source, so it is NOT attached to any single axis: that one
/// rule covers rectilinear (1-D monotonic → CF *dimension* coordinate),
/// unstructured (1-D over a shared dimension) and curvilinear (2-D `lat(y,x)`)
/// grids, the latter two emitting CF *auxiliary* coordinates.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Coordinate {
    /// Name of an existing data array (model variable, parameter, or loader
    /// field) supplying this coordinate's values. Mutually exclusive with
    /// `values`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub source: Option<String>,

    /// Inline literal 1-D coordinate vector for the simple rectilinear case,
    /// mirroring [`FunctionTableAxis::values`]. Mutually exclusive with
    /// `source`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub values: Option<Vec<f64>>,

    /// CF standard name (e.g. `"latitude"`), emitted verbatim.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub standard_name: Option<String>,

    /// CF/UDUNITS units string (e.g. `"degrees_north"`). Advisory at load
    /// time, matching [`FunctionTableAxis::units`].
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub units: Option<String>,

    /// CF axis role (`"X"` / `"Y"` / `"Z"` / `"T"`) for a 1-D monotonic
    /// dimension coordinate; absent for auxiliary coordinates, which have no
    /// single axis.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub axis: Option<String>,
}

/// ODE-based model component
///
/// # `name` / `description` are NOT schema keys
///
/// The schema's `Model` is `additionalProperties: false` and declares neither
/// `name` (a model's name is its key in the `models` map) nor `description`.
/// Both fields below are therefore schema-INVALID the moment they are set, and
/// `save()` will happily emit a document that fails its own schema. They are
/// harmless only while they stay `None`, which is what every load-parsed model
/// leaves them as — no deserializer ever fills them, because the wire form has
/// no such keys.
///
/// They are not dead, though: [`crate::update_model_metadata`] — a tier-
/// `extension`, Rust-only entry in `api-surface.json` — sets both, so its ONLY
/// observable effect is to make a valid document invalid. `graph.rs` also reads
/// `name` for a component-node label, where it is always `None` in practice.
/// They are part of the published API surface, so they stay; leave them
/// `None`.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct Model {
    /// Human-readable model name. NOT a schema key — see the type-level note.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,

    /// Academic citation or data source reference
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reference: Option<Reference>,

    /// State variables, parameters, and observed quantities (keyed by name)
    pub variables: IndexMap<String, ModelVariable>,

    /// Differential equations
    pub equations: Vec<Equation>,

    /// Discrete events
    #[serde(skip_serializing_if = "Option::is_none")]
    pub discrete_events: Option<Vec<DiscreteEvent>>,

    /// Continuous events
    #[serde(skip_serializing_if = "Option::is_none")]
    pub continuous_events: Option<Vec<ContinuousEvent>>,

    /// Named child models (subsystems), keyed by unique identifier
    #[serde(skip_serializing_if = "Option::is_none")]
    pub subsystems: Option<IndexMap<String, serde_json::Value>>,

    /// Brief description. NOT a schema key — see the type-level note.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// Model-level default numerical tolerance for inline tests.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tolerance: Option<Tolerance>,

    /// Inline validation tests that exercise this model in isolation.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub tests: Option<Vec<ModelTest>>,

    /// Inline illustrative analyses of how to run this model (schema
    /// `Model.analyses`); previously unmodelled here and so dropped on
    /// round trip.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub analyses: Option<Vec<ModelAnalysis>>,

    /// Equations that hold only at t=0 (initialization-only, not time-stepped).
    /// Introduced for aerosol equilibrium / plume-rise style models (gt-ebuq).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub initialization_equations: Option<Vec<Equation>>,

    /// Initial-guess seeds for nonlinear solvers during initialization, keyed
    /// by variable name. Values may be numeric literals or Expression graphs.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub guesses: Option<HashMap<String, serde_json::Value>>,

    /// MTK system-kind discriminator: "ode" (default), "nonlinear", "sde", "pde".
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub system_kind: Option<String>,
}

/// Variable within a model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelVariable {
    /// Variable type
    #[serde(rename = "type")]
    pub var_type: VariableType,

    /// Physical units
    #[serde(skip_serializing_if = "Option::is_none")]
    pub units: Option<String>,

    /// Default/initial value: a scalar, or — for a SHAPED variable — inline
    /// row-major nested array data (esm-spec §6.3). See [`InlineValue`].
    #[serde(skip_serializing_if = "Option::is_none")]
    pub default: Option<InlineValue>,

    /// The unit the `default` VALUE is expressed in, when it is not the
    /// variable's declared `units` (schema `ModelVariable.default_units`).
    ///
    /// The schema's contract is that `default` is given in the declared `units`,
    /// so a `default_units` naming a DIFFERENT unit is a defect: `units: "K"`
    /// with `default: 25.0, default_units: "degC"` means the stored number is 25
    /// but the variable actually reads 298.15. Rust did not model the field at
    /// all, so the mismatch was silently dropped on load.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub default_units: Option<String>,

    /// Floating-point precision of THIS variable and of the arithmetic over it,
    /// overriding the document-wide `domain.element_type` (esm-spec §11.3.1).
    ///
    /// This is NOT [`ModelVariable::var_type`]. That field is the variable's
    /// semantic role (`unknown` / `parameter`); this one is its element type.
    ///
    /// It exists for the split a single document-wide precision cannot express.
    /// A relational model may be binary32 in its floating-point *quantities* —
    /// reproducing a `real*4` reference — while its join keys are integers
    /// binary32 cannot represent: a ten-digit category code is ~2.26e9, 135× above
    /// binary32's exact-integer limit of 2^24, so under a document-wide
    /// `Float32` two distinct codes collapse onto one and a `join.on` over them
    /// merges unrelated rows. Declaring the key columns `Float64` keeps them
    /// exact while the quantities follow the document.
    ///
    /// `None` means the document's `domain.element_type`. See
    /// [`crate::precision_infer`] for how the declaration propagates through an
    /// expression and what it refuses.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub element_type: Option<String>,

    /// Brief description
    #[serde(skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// Arrayed-variable shape: ordered dimension names drawn from the
    /// enclosing model's domain.spatial. `None` means scalar.
    /// See discretization RFC §10.2.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub shape: Option<Vec<String>>,

    /// Staggered-grid location tag (e.g., "cell_center", "edge_normal",
    /// "vertex"). `None` means no explicit staggering.
    /// See discretization RFC §10.2.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub location: Option<String>,

    /// Parameter-only: draw the value from a distribution instead of fixing it
    /// at `default` (esm-spec §6.3, schema `ModelVariable.distribution`).
    /// Mutually exclusive with `default`. With no `update` the draw happens
    /// ONCE at setup (`sampled_parameters`); with `update.kind: "wiener"` it is
    /// redrawn every step with √dt scaling (`brownian_parameters`), which is
    /// what promotes the enclosing model to an SDE. This subsumes the 0.x
    /// `brownian` type's `noise_kind` / `correlation_group` tags: correlated
    /// noise is one vector-valued parameter carrying a `cov` matrix, which the
    /// tags never actually gave.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub distribution: Option<Distribution>,

    /// Parameter-only: when this parameter refreshes and what from
    /// (esm-spec §5.4). Absent means it never changes after setup. One
    /// mechanism replacing three 0.x constructs — the `brownian` type, the
    /// `discrete` type's `refresh` trigger, and the `discrete_parameters`
    /// event lists with their `functional_affect`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub update: Option<ParameterUpdateSpec>,
}

/// The bare variable: every optional field absent, `var_type` a `parameter`.
/// A manual impl because [`VariableType`] has no derivable default — `type`
/// is a required field on the wire — and of the two kinds only a parameter is
/// coherent with nothing else set (an unknown with no defining equation is a
/// structural error). Construction sites are still expected to state the type
/// they mean: `ModelVariable { var_type: .., ..Default::default() }`.
impl Default for ModelVariable {
    fn default() -> Self {
        ModelVariable {
            var_type: VariableType::Parameter,
            units: None,
            default: None,
            default_units: None,
            element_type: None,
            description: None,
            shape: None,
            location: None,
            distribution: None,
            update: None,
        }
    }
}

impl ModelVariable {
    /// This variable's SCALAR `default`, or `None` when it has none or carries
    /// inline array data instead (esm-spec §6.3).
    ///
    /// Every consumer that seeds one f64 slot reads through here, so a shaped
    /// variable whose whole column is authored inline never lands a bogus number
    /// in a scalar slot — the array channel picks it up instead
    /// ([`crate::simulate_array`]).
    pub fn default_scalar(&self) -> Option<f64> {
        self.default.as_ref().and_then(InlineValue::as_scalar)
    }

    /// This variable's `default` as inline ARRAY data, flattened row-major with
    /// its per-axis extents. `None` when the default is absent or scalar.
    ///
    /// # Errors
    ///
    /// A string naming the defect when the authored array is ragged.
    pub fn default_array(&self) -> Option<Result<DenseInlineArray, String>> {
        match self.default.as_ref() {
            Some(v @ InlineValue::Array(_)) => Some(v.to_dense()),
            _ => None,
        }
    }

    /// Visit every Expression this variable carries, in a stable order.
    ///
    /// From esm 1.0.0 a variable has no `expression` field — an unknown's
    /// definition is an EQUATION. What is left on the variable is the
    /// parameter's `update`: each rule's `when` trigger, its `expression`
    /// value form, and the `unit_conversion` of a `from` binding. All three
    /// are Expression positions and so are subject to reference integrity
    /// (esm-spec §4.9.5), which is what this walk exists to serve.
    pub fn for_each_expression(&self, f: &mut impl FnMut(&Expr)) {
        let Some(spec) = &self.update else { return };
        for rule in spec.rules() {
            if let Some(when) = rule.when() {
                f(when);
            }
            let Some(value) = rule.value() else { continue };
            if let Some(expr) = &value.expression {
                f(expr);
            }
            if let Some(UnitConversion::Expression(expr)) =
                value.from.as_ref().and_then(|b| b.unit_conversion.as_ref())
            {
                f(expr);
            }
        }
    }

    /// As [`ModelVariable::for_each_expression`], but also handing the callback
    /// the JSON-pointer suffix of the SITE the expression came from, relative to
    /// the variable's `update`.
    ///
    /// A diagnostic must name the field it is about. Reporting every one of a
    /// parameter's update expressions at a bare `.../update` is ambiguous the
    /// moment the ordered ARRAY form is used (esm-spec §5.4) — and that form
    /// exists precisely because parameters written by two or more events are
    /// common — so `when`, `expression` and `from/unit_conversion` all have to
    /// be distinguishable, per rule. The suffixes match the other four
    /// bindings: `""` for the single-rule object form and `/i` for entry `i` of
    /// the array form, then the site (`/when`, `/expression`,
    /// `/from/unit_conversion`).
    pub fn for_each_expression_at(&self, f: &mut impl FnMut(&Expr, &str)) {
        let Some(spec) = &self.update else { return };
        let indexed = matches!(spec, ParameterUpdateSpec::Several(_));
        for (i, rule) in spec.rules().iter().enumerate() {
            let rule_path = if indexed {
                format!("/{i}")
            } else {
                String::new()
            };
            if let Some(when) = rule.when() {
                f(when, &format!("{rule_path}/when"));
            }
            let Some(value) = rule.value() else { continue };
            if let Some(expr) = &value.expression {
                f(expr, &format!("{rule_path}/expression"));
            }
            if let Some(UnitConversion::Expression(expr)) =
                value.from.as_ref().and_then(|b| b.unit_conversion.as_ref())
            {
                f(expr, &format!("{rule_path}/from/unit_conversion"));
            }
        }
    }

    /// Mutable counterpart of [`ModelVariable::for_each_expression`], for the
    /// namespacing / renaming / range-resolution passes.
    pub fn for_each_expression_mut(&mut self, f: &mut impl FnMut(&mut Expr)) {
        let Some(spec) = &mut self.update else { return };
        let rules: &mut [ParameterUpdate] = match spec {
            ParameterUpdateSpec::Single(rule) => std::slice::from_mut(rule),
            ParameterUpdateSpec::Several(rules) => rules,
        };
        for rule in rules {
            if let Some(when) = rule.when_mut() {
                f(when);
            }
            let Some(value) = rule.value_mut() else {
                continue;
            };
            if let Some(expr) = &mut value.expression {
                f(expr);
            }
            if let Some(UnitConversion::Expression(expr)) =
                value.from.as_mut().and_then(|b| b.unit_conversion.as_mut())
            {
                f(expr);
            }
        }
    }

    /// Error-propagating [`ModelVariable::for_each_expression_mut`]: visits
    /// the same expression set in the same order and returns the FIRST error
    /// `f` raises (later expressions are not passed to `f`, so they are left
    /// unmodified). Defined in terms of `for_each_expression_mut` itself, so
    /// the two can never disagree about which fields carry expressions.
    pub fn try_for_each_expression_mut<E>(
        &mut self,
        f: &mut impl FnMut(&mut Expr) -> Result<(), E>,
    ) -> Result<(), E> {
        let mut first_err: Option<E> = None;
        self.for_each_expression_mut(&mut |expr| {
            if first_err.is_none()
                && let Err(e) = f(expr)
            {
                first_err = Some(e);
            }
        });
        match first_err {
            Some(e) => Err(e),
            None => Ok(()),
        }
    }

    /// The `data_sources` keys this variable's update rules read from.
    pub fn update_sources(&self) -> Vec<&str> {
        self.update
            .iter()
            .flat_map(|spec| spec.rules())
            .filter_map(|rule| rule.data_source())
            .collect()
    }
}

/// A parameter's value drawn from a probability distribution rather than fixed
/// (esm-spec §6.3, schema `Distribution`). The closed set is `normal`,
/// `lognormal`, `uniform` — bindings implement exactly these and reject
/// anything else.
///
/// Univariate when the location parameter is a number, multivariate when it is
/// an array — in which case `cov` gives the full covariance matrix and the
/// parameter's `shape` must agree.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum Distribution {
    /// Gaussian. Exactly one of `std` (independent components) or `cov`.
    Normal {
        /// Mean — a number (scalar) or one entry per component (vector).
        mean: DistributionParam,
        /// Standard deviation of INDEPENDENT components. Excludes `cov`.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        std: Option<DistributionParam>,
        /// Full covariance matrix, row-major. Excludes `std`.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        cov: Option<CovarianceMatrix>,
    },
    /// Log-normal: `ln(value)` is normal with mean `mu` and spread
    /// `sigma` / `cov` (both on the LOG scale).
    Lognormal {
        /// Mean of the underlying normal (log scale).
        mu: DistributionParam,
        /// Standard deviation of the underlying normal. Excludes `cov`.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        sigma: Option<DistributionParam>,
        /// Full covariance matrix on the log scale. Excludes `sigma`.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        cov: Option<CovarianceMatrix>,
    },
    /// Uniform on `[low, high]`. Components are independent by construction, so
    /// there is no covariance form.
    Uniform {
        /// Lower bound(s).
        low: DistributionParam,
        /// Upper bound(s).
        high: DistributionParam,
    },
}

impl Distribution {
    /// The distribution's `kind` string, as it appears on the wire.
    pub fn kind(&self) -> &'static str {
        match self {
            Distribution::Normal { .. } => "normal",
            Distribution::Lognormal { .. } => "lognormal",
            Distribution::Uniform { .. } => "uniform",
        }
    }

    /// The LOCATION parameter (`mean` / `mu` / `low`), whose form decides
    /// whether the distribution is univariate or multivariate.
    pub fn location(&self) -> &DistributionParam {
        match self {
            Distribution::Normal { mean, .. } => mean,
            Distribution::Lognormal { mu, .. } => mu,
            Distribution::Uniform { low, .. } => low,
        }
    }

    /// True when the location parameter is an array — the schema's
    /// univariate/multivariate discriminator.
    pub fn is_multivariate(&self) -> bool {
        self.location().is_vector()
    }

    /// The full covariance matrix, when this distribution carries one.
    pub fn cov(&self) -> Option<&CovarianceMatrix> {
        match self {
            Distribution::Normal { cov, .. } | Distribution::Lognormal { cov, .. } => cov.as_ref(),
            Distribution::Uniform { .. } => None,
        }
    }
}

/// A distribution's location/scale parameter: a number for a scalar parameter,
/// an array for a vector-valued one.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum DistributionParam {
    /// Scalar (univariate) form.
    Scalar(f64),
    /// Per-component (multivariate) form.
    Vector(Vec<f64>),
}

impl DistributionParam {
    /// True when this is the multivariate (array) form.
    pub fn is_vector(&self) -> bool {
        matches!(self, DistributionParam::Vector(_))
    }
}

/// Symmetric positive-semidefinite covariance matrix, row-major: entry
/// `[i][j]` is the covariance of components `i` and `j`.
pub type CovarianceMatrix = Vec<Vec<f64>>;

/// A parameter's update behaviour (esm-spec §5.4, schema
/// `ParameterUpdateSpec`): EITHER a single rule, OR an ordered array of two or
/// more applied in declaration order.
///
/// A single rule MUST be the object form — a one-element array is invalid — so
/// every update set has exactly one spelling and the round trip is stable.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
// The single form carries a whole `ParameterUpdate` inline while the array form
// carries a `Vec`, so the variants differ in size. Boxing the single one would
// pay an allocation on the COMMON spelling — most parameters carry one rule —
// to shrink a type that lives once per variable and is never in a hot loop, and
// it would cost the wire-facing `ParameterUpdateSpec::Single(rule)` match
// ergonomics on the way. Same call as `DiscreteEventTrigger` above.
#[allow(clippy::large_enum_variant)]
pub enum ParameterUpdateSpec {
    /// The object form: exactly one rule.
    Single(ParameterUpdate),
    /// The array form: two or more rules, applied in declaration order.
    Several(Vec<ParameterUpdate>),
}

impl ParameterUpdateSpec {
    /// The rules, in declaration order. The single form yields a one-element
    /// slice, so every consumer reads one shape.
    pub fn rules(&self) -> &[ParameterUpdate] {
        match self {
            ParameterUpdateSpec::Single(rule) => std::slice::from_ref(rule),
            ParameterUpdateSpec::Several(rules) => rules,
        }
    }

    /// True iff ANY rule is `wiener` — the `brownian_parameters` test
    /// (esm-spec §6.3.1). The schema forbids `wiener` inside the array form,
    /// so in practice only the single form can answer true.
    pub fn is_wiener(&self) -> bool {
        self.rules()
            .iter()
            .any(|r| matches!(r, ParameterUpdate::Wiener))
    }

    /// True iff any rule requires the variable to declare a `shape` —
    /// `schedule`, `data`, or `remesh` (esm-spec §5.4 "Cadence"): such a
    /// parameter is a buffer whose extent is fixed at setup.
    pub fn requires_shape(&self) -> bool {
        self.rules().iter().any(|r| {
            matches!(
                r,
                ParameterUpdate::Schedule { .. }
                    | ParameterUpdate::Data { .. }
                    | ParameterUpdate::Remesh { .. }
            )
        })
    }
}

/// One update rule: WHEN a parameter refreshes, and WHAT from
/// (esm-spec §5.4). Every kind except `wiener` carries exactly one value form
/// — `expression`, `from`, or `handler` — held in the flattened
/// [`UpdateValue`]; `wiener` carries none, because the parameter's
/// `distribution` IS its value.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum ParameterUpdate {
    /// A driving Wiener (Brownian) process: the parameter's `distribution` is
    /// resampled every step with √dt increment scaling. Any such parameter
    /// promotes the enclosing model from an ODE system to an SDE system.
    Wiener,
    /// Time-driven refresh at preset `times` and/or on a periodic `interval`
    /// (at least one of the two is present) — the MTK `PresetTimeCallback` /
    /// `tstops` analogue.
    Schedule {
        /// Explicit simulation times (the tstops) at which to refresh.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        times: Option<Vec<f64>>,
        /// Periodic refresh interval in simulation time units.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        interval: Option<f64>,
        /// Offset from t=0 for the first periodic refresh.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        initial_offset: Option<f64>,
        /// The value form.
        #[serde(flatten)]
        value: UpdateValue,
    },
    /// Refresh at the end of any timestep at which the boolean `when` is true.
    Condition {
        /// Boolean expression tested at the end of each timestep.
        when: Expr,
        /// The value form.
        #[serde(flatten)]
        value: UpdateValue,
    },
    /// Refresh when `when` crosses zero, located by root-finding.
    Crossing {
        /// Expression whose zero crossing triggers the refresh.
        when: Expr,
        /// Which crossings count: `up`, `down`, or `any` (the default).
        #[serde(default, skip_serializing_if = "Option::is_none")]
        direction: Option<String>,
        /// The value form.
        #[serde(flatten)]
        value: UpdateValue,
    },
    /// Refresh when the named `data_sources` entry advances a record.
    Data {
        /// Key of the `data_sources` entry driving the refresh. MUST resolve
        /// (`data_source_undefined`).
        source: String,
        /// The value form.
        #[serde(flatten)]
        value: UpdateValue,
    },
    /// Refresh on a mesh-topology change (AMR refinement, moving mesh).
    Remesh {
        /// Optional name of the remesh hook. Absent ⇒ any remesh event.
        #[serde(default, skip_serializing_if = "Option::is_none")]
        hook: Option<String>,
        /// The value form.
        #[serde(flatten)]
        value: UpdateValue,
    },
}

impl ParameterUpdate {
    /// The rule's `kind` string, as it appears on the wire.
    pub fn kind(&self) -> &'static str {
        match self {
            ParameterUpdate::Wiener => "wiener",
            ParameterUpdate::Schedule { .. } => "schedule",
            ParameterUpdate::Condition { .. } => "condition",
            ParameterUpdate::Crossing { .. } => "crossing",
            ParameterUpdate::Data { .. } => "data",
            ParameterUpdate::Remesh { .. } => "remesh",
        }
    }

    /// The rule's value form, or `None` for `wiener` (which has none).
    pub fn value(&self) -> Option<&UpdateValue> {
        match self {
            ParameterUpdate::Wiener => None,
            ParameterUpdate::Schedule { value, .. }
            | ParameterUpdate::Condition { value, .. }
            | ParameterUpdate::Crossing { value, .. }
            | ParameterUpdate::Data { value, .. }
            | ParameterUpdate::Remesh { value, .. } => Some(value),
        }
    }

    /// Mutable access to the rule's value form, or `None` for `wiener`.
    pub fn value_mut(&mut self) -> Option<&mut UpdateValue> {
        match self {
            ParameterUpdate::Wiener => None,
            ParameterUpdate::Schedule { value, .. }
            | ParameterUpdate::Condition { value, .. }
            | ParameterUpdate::Crossing { value, .. }
            | ParameterUpdate::Data { value, .. }
            | ParameterUpdate::Remesh { value, .. } => Some(value),
        }
    }

    /// The boolean/zero-crossing trigger expression (`condition` / `crossing`).
    pub fn when(&self) -> Option<&Expr> {
        match self {
            ParameterUpdate::Condition { when, .. } | ParameterUpdate::Crossing { when, .. } => {
                Some(when)
            }
            _ => None,
        }
    }

    /// Mutable access to the trigger expression (`condition` / `crossing`).
    pub fn when_mut(&mut self) -> Option<&mut Expr> {
        match self {
            ParameterUpdate::Condition { when, .. } | ParameterUpdate::Crossing { when, .. } => {
                Some(when)
            }
            _ => None,
        }
    }

    /// The `data_sources` key this rule refreshes from, for the `data` kind.
    pub fn data_source(&self) -> Option<&str> {
        match self {
            ParameterUpdate::Data { source, .. } => Some(source.as_str()),
            _ => None,
        }
    }
}

/// The value form of an update rule: exactly one of `expression`, `from`, or
/// `handler` (esm-spec §5.4, "What: exactly one value form"). Modelled as three
/// options rather than an enum so the flattened round trip stays verbatim; the
/// "exactly one" constraint is a schema-layer check.
#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
pub struct UpdateValue {
    /// The new value, computed symbolically (§4).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub expression: Option<Expr>,

    /// The new value, read from the `data_sources` entry named by
    /// `update.source` (§8).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub from: Option<DataSourceBinding>,

    /// The new value, computed by a registered handler (§5.5).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub handler: Option<FunctionalUpdate>,
}

/// A registered handler computing a parameter's new value when its update
/// fires (esm-spec §5.5). The 0.x event `functional_affect`, relocated onto the
/// parameter it writes — which is why it needs no `modified_params` list: it
/// writes exactly one thing, the parameter carrying it.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct FunctionalUpdate {
    /// Registered identifier for the handler implementation.
    pub handler_id: String,

    /// Unknowns read by the handler. Absent means none.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub read_vars: Option<Vec<String>>,

    /// Parameters read by the handler. Absent means none.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub read_params: Option<Vec<String>>,

    /// Handler-specific configuration, passed through verbatim.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub config: Option<serde_json::Value>,
}

/// Binds a parameter to ONE variable of a `data_sources` entry (esm-spec §8.5).
/// The 0.x `DataLoaderVariable` minus `units`: the units are the parameter's
/// own, declared once on the parameter instead of twice.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct DataSourceBinding {
    /// Name of the variable in the source file supplying the values.
    pub file_variable: String,

    /// Multiplicative factor or Expression AST reaching the parameter's
    /// declared `units`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub unit_conversion: Option<UnitConversion>,

    /// Text-label decoding table, passed through verbatim.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub codes: Option<serde_json::Value>,

    /// Per-parameter override of the source-level `select`.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub select: Option<serde_json::Value>,

    /// Brief description.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub description: Option<String>,

    /// Academic citation or data source reference.
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub reference: Option<Reference>,
}

/// Type of model variable (esm-spec §6.3). There are **two** declared types
/// and no third: everything else a solver needs — which unknowns are ODE
/// states, which parameters are Brownian — is DERIVED by
/// [`crate::classification`], never declared.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum VariableType {
    /// A quantity the solver solves for. Its behaviour is stated by the
    /// model's `equations` and NOWHERE else — there is no `expression` field
    /// on a variable. Subsumes the 0.x `state` and `observed` types; which of
    /// the two an unknown is follows from the form of the equation LHS that
    /// defines it ([`crate::classification::ode_states`] /
    /// [`crate::classification::observed_unknowns`] /
    /// [`crate::classification::algebraic_unknowns`]).
    Unknown,
    /// A quantity supplied to the solver. Its value is `default` or a
    /// `distribution`, optionally refreshed by an `update` (§5.4). Subsumes
    /// the 0.x `parameter`, `brownian` and `discrete` types
    /// ([`crate::classification::brownian_parameters`] /
    /// [`crate::classification::discrete_parameters`] /
    /// [`crate::classification::sampled_parameters`] /
    /// [`crate::classification::constant_parameters`]).
    Parameter,
}

/// Differential equation
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Equation {
    /// Left-hand side expression
    pub lhs: Expr,

    /// Right-hand side expression
    pub rhs: Expr,

    /// Author's annotation on this equation (schema `Equation._comment`).
    /// The schema allows it explicitly alongside `additionalProperties:
    /// false`, so it is normative document content and must survive
    /// `parse → emit` — the field exists purely for that pass-through.
    #[serde(rename = "_comment", default, skip_serializing_if = "Option::is_none")]
    pub comment: Option<String>,
}

impl Default for Equation {
    fn default() -> Self {
        Equation {
            lhs: Expr::Integer(0),
            rhs: Expr::Integer(0),
            comment: None,
        }
    }
}