"""
ResidualEngine: residuals for governing laws, constitutive, and scaling/similarity audits.

- compute_residuals(state_pred, state_ref=None, *, log_to_python=True)
- build_loss: cascaded RMS over laws only (training).
- audit / visualize: same metrics (RMS, R_norm, admissibility) for all residual keys;
  visualize builds Plotly training/test dashboards (optional spatial law panel for x slices).

Constitutive and scaling/similarity audits are tied to Models.* and Groups.* functions via
**ref_delta**, **implied_delta** (constitutive only), and **π-constant** checks (scaling, Path A).
**implied_delta** and **ref_delta** are always **nondimensional** (see
``moju.monitor.closure_registry.apply_closure_discrepancy_normalize``). For **R_norm**, default
**scale_k** is **≈ 1** (plus ``_SCALE_EPS``) for **laws/** and for nondimensional closure keys;
other **constitutive/** / **scaling/** residuals and **data/** use state- or reference-derived scales. Optional ``audit(..., r_ref=...)`` overrides
**scale_k** per key. Per-key **admissibility_score** is ``1 / (1 + R_norm)`` when finite. Metrics are
consistency indicators, not certification.
"""

from __future__ import annotations

import datetime
import math
from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Union

import jax
import jax.numpy as jnp

from moju.piratio.groups import Groups
from moju.piratio.laws import Laws
from moju.monitor.closure_registry import (
    GROUP_FNS,
    MODEL_FNS,
    compute_implied_delta,
    compute_ref_delta,
)
from moju.monitor.derived_state_chain import all_ref_keys_from_chain, keys_produced_by_chain
from moju.monitor.pi_constant_recipes import (
    GROUP_PI_CONSTANT_RECIPES,
    apply_pi_constant_recipe,
)
from moju.monitor.law_implied_diagnostics import (
    merge_fragment_law_implied_audit_specs,
    merge_law_implied_audit_specs,
)
from moju.monitor.spatial_rnorm_panels import build_spatial_rnorm_panels_from_residuals
from moju.monitor.visualize_labels import pretty_category_name, pretty_residual_key

DEFAULT_VISUALIZE_TITLE_TRAINING = "Physics admissibility audit (model training)"
DEFAULT_VISUALIZE_TITLE_TEST = "State prediction audit (physics residuals)"


def _resolve_visualize_figure_title(mode: str, figure_title: Optional[str]) -> str:
    """Figure title for :func:`visualize`; non-empty ``figure_title`` overrides mode default."""
    if figure_title is not None and str(figure_title).strip():
        return str(figure_title).strip()
    if mode == "training":
        return DEFAULT_VISUALIZE_TITLE_TRAINING
    return DEFAULT_VISUALIZE_TITLE_TEST


def _state_key_suggests_law_fd_fill(key: str) -> bool:
    """True if missing key may be fillable via fill_law_fd (derived law input)."""
    if key.endswith("_laplacian") or key.endswith("_grad") or key.endswith("_tt"):
        return True
    if key.endswith("_t") and not key.endswith("_tt"):
        return True
    return key == "u_grad"


def _string_placeholder_like(v: Any, *, key: str, arg_name: str) -> bool:
    """Best-effort check for UI placeholders accidentally saved as strings."""
    if not isinstance(v, str):
        return False
    s = v.strip()
    if not s:
        return True
    low = s.lower()
    if low in {"none", "null", "nan", "na"}:
        return True
    return s == key or s == arg_name


def _kwargs_from_state(
    state: Dict[str, Any],
    constants: Dict[str, Any],
    state_map: Dict[str, str],
    *,
    law_context: Optional[str] = None,
) -> Dict[str, Any]:
    """Build kwargs for a law/group from state_map (arg_name -> state_key)."""
    out = {}
    for arg_name, key in state_map.items():
        val = state.get(key)
        if _string_placeholder_like(val, key=key, arg_name=str(arg_name)):
            val = None
        if val is None:
            val = constants.get(key)
            if _string_placeholder_like(val, key=key, arg_name=str(arg_name)):
                val = None
        if val is None:
            msg = f"Key {key!r} not found in state or constants (arg {arg_name})"
            if law_context is not None:
                msg += f" (law {law_context!r})"
                if _state_key_suggests_law_fd_fill(key):
                    msg += (
                        " — on Path B, enable auto_path_b_derivatives=True and fill_law_fd=True, "
                        "provide the primitive field (e.g. T for T_laplacian) and mesh coordinates "
                        "(x, and y/z/t as needed); see moju.monitor.law_fd_recipes.LAW_FD_RECIPES. "
                        "The engine does not rename NPZ keys: use key `T` (or change law state_map), "
                        "not only Studio alias hints."
                    )
            raise KeyError(msg)
        if isinstance(val, str):
            msg = (
                f"Key {key!r} resolved to string value {val!r} (arg {arg_name})"
                " but laws expect numeric array-like inputs."
            )
            if law_context is not None:
                msg += f" (law {law_context!r})"
            raise TypeError(msg)
        out[arg_name] = val
    return out


def _get_fn(spec: Dict[str, Any], builtin_class: Any) -> Any:
    if "fn" in spec:
        return spec["fn"]
    return getattr(builtin_class, spec["name"])


def _build_state(
    state_pred: Dict[str, Any],
    constants: Dict[str, Any],
    groups_spec: List[Dict],
) -> Dict[str, Any]:
    """Run group specs in order; write output_key into state."""
    state = dict(state_pred)
    merged = {**constants, **state}
    for spec in groups_spec:
        state_map = spec["state_map"]
        output_key = spec["output_key"]
        kwargs = _kwargs_from_state(merged, constants, state_map)
        fn = _get_fn(spec, Groups)
        state[output_key] = fn(**kwargs)
        merged[output_key] = state[output_key]
    return state


def _rms_scalar(x: jnp.ndarray) -> jnp.ndarray:
    a = jnp.asarray(x)
    if a.size == 0:
        return jnp.asarray(float("nan"))
    return jnp.sqrt(jnp.nanmean(jnp.square(a)))


def admissibility_level(score: float) -> str:
    if not math.isfinite(score):
        return "Unknown"
    if score >= 0.90:
        return "High Admissibility"
    if score >= 0.70:
        return "Moderate Admissibility"
    if score >= 0.40:
        return "Low Admissibility"
    return "Non-Admissible"


def _geom_mean_admissibility(values: Sequence[float]) -> float:
    """
    Geometric mean over finite scores only. If any finite value is <= 0, returns 0.0.
    If no finite values, returns NaN.
    """
    finite: List[float] = []
    for x in values:
        try:
            xf = float(x)
        except (TypeError, ValueError):
            continue
        if math.isfinite(xf):
            finite.append(xf)
    if not finite:
        return float("nan")
    if min(finite) <= 0.0:
        return 0.0
    return float(math.exp(sum(math.log(x) for x in finite) / len(finite)))


def _rms_per_key(
    residuals_flat: Dict[str, jnp.ndarray],
    *,
    to_python: bool = True,
) -> Dict[str, Any]:
    out: Dict[str, Any] = {}
    for key, arr in residuals_flat.items():
        r = _rms_scalar(arr)
        if to_python:
            out[key] = float(jax.device_get(r))
        else:
            out[key] = r
    return out


def _flatten_residual_dict(residuals: Dict[str, Any]) -> Dict[str, jnp.ndarray]:
    flat: Dict[str, jnp.ndarray] = {}
    for category, content in residuals.items():
        if not isinstance(content, dict):
            if hasattr(content, "shape"):
                flat[category] = jnp.asarray(content)
            else:
                flat[category] = jnp.asarray(content)
            continue
        for name, arr in content.items():
            flat[f"{category}/{name}"] = jnp.asarray(arr)
    return flat


_SCALE_EPS = 1e-12


def _default_unit_scale_k(*, to_python: bool) -> float:
    """scale_k = 1 + eps for nondimensional law / implied / ref residuals."""
    s = _SCALE_EPS + 1.0
    return float(jax.device_get(jnp.asarray(s))) if to_python else float(s)


def _suffix_is_nd_closure(rest: str) -> bool:
    """True if flat tail is implied_delta or ref_delta (nondimensional closure residual)."""
    parts = rest.split("/")
    return bool(parts) and parts[-1] in ("implied_delta", "ref_delta")


def _state_derived_scale_per_key(
    flat_keys: Iterable[str],
    merged: Dict[str, Any],
    _laws_spec: List[Dict[str, Any]],
    constitutive_audit: List[Dict[str, Any]],
    scaling_audit: List[Dict[str, Any]],
    state_ref_built: Optional[Dict[str, Any]] = None,
    *,
    to_python: bool = True,
) -> Dict[str, float]:
    """
    Per-key scale for ``R_norm = RMS(r_k) / scale_k`` stored on each log entry (``entry["scale"]``).

    Default **scale_k** is **≈ 1** for governing **laws/** and for nondimensional
    **implied_delta** / **ref_delta** under **constitutive/** and **scaling/**. Other audit keys and
    **data/** use RMS of relevant state (or reference) fields. Optional ``r_ref`` in
    :func:`audit` overrides per key after logging.
    """
    out: Dict[str, float] = {}
    ref = state_ref_built if state_ref_built is not None else merged

    for k in flat_keys:
        if "/" not in k:
            out[k] = _default_unit_scale_k(to_python=to_python)
            continue
        prefix, rest = k.split("/", 1)
        if prefix == "laws":
            out[k] = _default_unit_scale_k(to_python=to_python)
            continue
        if prefix == "constitutive":
            if _suffix_is_nd_closure(rest):
                out[k] = _default_unit_scale_k(to_python=to_python)
                continue
            name = rest.split("/")[0]
            spec = next(
                (
                    s
                    for s in constitutive_audit
                    if s.get("name") == name
                    or (s.get("residual_basename") or "").split("/")[0] == name
                ),
                None,
            )
            if spec is None:
                scale = _SCALE_EPS + _rms_scalar(jnp.asarray(1.0))
            else:
                state_map = spec.get("state_map") or {}
                output_key = spec.get("output_key")
                parts = []
                for sk in state_map.values():
                    if sk in merged:
                        parts.append(jnp.ravel(jnp.asarray(merged[sk])))
                if output_key and output_key in merged:
                    parts.append(jnp.ravel(jnp.asarray(merged[output_key])))
                ivk = spec.get("implied_value_key")
                if ivk and ivk in merged:
                    parts.append(jnp.ravel(jnp.asarray(merged[ivk])))
                if parts:
                    big = jnp.concatenate(parts)
                    scale = _SCALE_EPS + _rms_scalar(big)
                else:
                    scale = _SCALE_EPS + _rms_scalar(jnp.asarray(1.0))
            out[k] = float(jax.device_get(scale)) if to_python else float(scale)
            continue
        if prefix == "scaling":
            if _suffix_is_nd_closure(rest):
                out[k] = _default_unit_scale_k(to_python=to_python)
                continue
            name = rest.split("/")[0]
            spec = next(
                (
                    s
                    for s in scaling_audit
                    if s.get("name") == name
                    or (s.get("residual_basename") or "").split("/")[0] == name
                ),
                None,
            )
            if spec is None:
                scale = _SCALE_EPS + _rms_scalar(jnp.asarray(1.0))
            else:
                state_map = spec.get("state_map") or {}
                output_key = spec.get("output_key")
                parts = []
                for sk in state_map.values():
                    if sk in merged:
                        parts.append(jnp.ravel(jnp.asarray(merged[sk])))
                if output_key and output_key in merged:
                    parts.append(jnp.ravel(jnp.asarray(merged[output_key])))
                ivk_s = spec.get("implied_value_key")
                if ivk_s and ivk_s in merged:
                    parts.append(jnp.ravel(jnp.asarray(merged[ivk_s])))
                if parts:
                    big = jnp.concatenate(parts)
                    scale = _SCALE_EPS + _rms_scalar(big)
                else:
                    scale = _SCALE_EPS + _rms_scalar(jnp.asarray(1.0))
            out[k] = float(jax.device_get(scale)) if to_python else float(scale)
            continue
        if prefix == "data":
            state_key = rest
            if state_key in ref:
                v = jnp.asarray(ref[state_key])
                scale = _SCALE_EPS + _rms_scalar(jnp.ravel(v))
            else:
                scale = _SCALE_EPS + _rms_scalar(jnp.asarray(1.0))
            out[k] = float(jax.device_get(scale)) if to_python else float(scale)
            continue
        out[k] = _default_unit_scale_k(to_python=to_python)
    return out


def build_loss(
    residual_dict: Dict[str, Any],
    option: str = "cascaded",
    law_weights: Optional[Dict[str, float]] = None,
) -> jnp.ndarray:
    if option != "cascaded":
        raise ValueError(f"Only option='cascaded' is implemented, got {option!r}")
    laws = residual_dict.get("laws", {})
    if not laws:
        return jnp.array(0.0)
    names = list(laws.keys())
    n = len(names)
    weights = law_weights or {}
    w = jnp.array([weights.get(name, 1.0 / n) for name in names])
    rms_vals = jnp.array([_rms_scalar(jnp.asarray(laws[name])) for name in names])
    return jnp.sum(w * rms_vals)


def _compute_log_step_metrics(
    log: List[Dict[str, Any]],
    r_ref: Optional[Dict[str, float]] = None,
) -> List[Dict[str, Any]]:
    """
    Per-log-entry admissibility metrics (same rules as ``audit``), without mutating ``log``.

    Returns one dict per entry with keys: ``r_norm``, ``admissibility_score``,
    ``category_admissibility_score``, ``overall_admissibility_score``, and
    ``per_key_report`` (flat key -> {rms, r_norm, admissibility_score, admissibility_level}).
    Scale precedence per key: **r_ref** (if given) > entry ``scale`` > first-step RMS > 1.
    Category scores are **0** if any per-key admissibility in that category is non-finite.
    """
    if not log:
        return []
    first_rms = log[0].get("rms", {})
    out: List[Dict[str, Any]] = []
    category_buckets = ("laws", "constitutive", "scaling", "data")
    for entry in log:
        rms = entry.get("rms", {})
        entry_scale = entry.get("scale") or {}
        r_norm: Dict[str, float] = {}
        admissibility: Dict[str, float] = {}
        per_key_report: Dict[str, Any] = {}
        for k, v in rms.items():
            if r_ref is not None and k in r_ref and r_ref[k] is not None and r_ref[k] > 0:
                scale_k = r_ref[k]
            elif k in entry_scale and entry_scale[k] is not None and entry_scale[k] > 0:
                scale_k = entry_scale[k]
            elif k in first_rms and first_rms[k] is not None and first_rms[k] > 0:
                scale_k = first_rms[k]
            else:
                scale_k = 1.0
            if scale_k <= 0 or not math.isfinite(float(scale_k)):
                scale_k = 1.0
            try:
                v_f = float(v)
            except (TypeError, ValueError):
                v_f = float("nan")
            if not math.isfinite(v_f):
                r_norm[k] = float("nan")
                admissibility[k] = float("nan")
            else:
                r_norm[k] = v_f / scale_k
                if not math.isfinite(r_norm[k]):
                    admissibility[k] = float("nan")
                else:
                    admissibility[k] = 1.0 / (1.0 + r_norm[k])
            per_key_report[k] = {
                "rms": v,
                "r_norm": r_norm[k],
                "admissibility_score": admissibility[k],
                "admissibility_level": admissibility_level(admissibility[k]),
            }
        category_keys: Dict[str, List[str]] = {c: [] for c in category_buckets}
        for k in admissibility.keys():
            if "/" in k:
                prefix = k.split("/", 1)[0]
                if prefix in category_keys:
                    category_keys[prefix].append(k)
        category_scores: Dict[str, float] = {}
        for cat, keys in category_keys.items():
            if not keys:
                continue
            vals: List[float] = []
            category_failed = False
            for kk in keys:
                try:
                    vf = float(admissibility[kk])
                except (TypeError, ValueError):
                    category_failed = True
                    break
                if not math.isfinite(vf):
                    category_failed = True
                    break
                vals.append(vf)
            if category_failed:
                category_scores[cat] = 0.0
            else:
                category_scores[cat] = _geom_mean_admissibility(vals)
        cats_present = list(category_scores.keys())
        if not cats_present:
            overall = float("nan")
        else:
            overall = _geom_mean_admissibility([float(category_scores[c]) for c in cats_present])
        out.append(
            {
                "r_norm": r_norm,
                "admissibility_score": admissibility,
                "category_admissibility_score": category_scores,
                "overall_admissibility_score": overall,
                "per_key_report": per_key_report,
            }
        )
    return out


def audit(
    log: List[Dict[str, Any]],
    r_ref: Optional[Dict[str, float]] = None,
    weights: Optional[Dict[str, float]] = None,
    *,
    export_dir: Optional[str] = None,
    save_residuals: bool = False,
    last_residual_dict: Optional[Dict[str, Any]] = None,
    model_name: Optional[str] = None,
    model_id: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Physics admissibility from logged RMS and scales.

    **R_norm** uses ``RMS(r_k) / scale_k`` with ``scale_k`` from each entry's ``scale`` when
    positive, else the first step's RMS fallback, else 1. **ResidualEngine** logs default
    ``scale_k ≈ 1`` for **laws/** and nondimensional **implied_delta** / **ref_delta** keys;
    other residuals and **data/** use state-derived scales. Optional **r_ref** (flat key → positive
    float) overrides ``scale_k`` for those keys. Per-key **admissibility_score** is
    ``1 / (1 + R_norm)`` when finite.

    Reporting uses three levels (no extra aggregation API): (1) per residual key in
    ``per_key``; (2) geometric mean within each category — ``per_category`` keys
    ``laws``, ``constitutive``, ``scaling``, ``data`` — only when **every** per-key
    admissibility in that category is finite; if any key is non-finite (NaN/Inf) or
    non-numeric, the category score is **0** (inadmissible for that bucket); empty
    categories omitted; (3) overall score — geometric mean of present category scores
    (any category at **0** drives overall to **0**).
    """
    if not log:
        return {"per_key": {}, "overall_admissibility_score": 0.0, "overall_admissibility_level": "Non-Admissible"}
    step_metrics = _compute_log_step_metrics(log, r_ref)
    last_report_per_key: Dict[str, Any] = {}
    for entry, m in zip(log, step_metrics):
        entry["r_norm"] = m["r_norm"]
        entry["admissibility_score"] = m["admissibility_score"]
        entry["category_admissibility_score"] = m["category_admissibility_score"]
        entry["overall_admissibility_score"] = m["overall_admissibility_score"]
        last_report_per_key = dict(m["per_key_report"])
    overall = log[-1].get("overall_admissibility_score", 0.0) if log else 0.0
    report = {
        "per_key": last_report_per_key,
        "per_category": log[-1].get("category_admissibility_score", {}) if log else {},
        "overall_admissibility_score": overall,
        "overall_admissibility_level": admissibility_level(overall),
    }

    if export_dir:
        import zipfile
        from pathlib import Path
        session_name = datetime.datetime.now().strftime("audit_%Y%m%d_%H%M")
        session_dir = Path(export_dir) / session_name
        session_dir.mkdir(parents=True, exist_ok=True)
        try:
            from moju.monitor.report import write_audit_pdf, write_residuals_json
            pdf_path = session_dir / "report.pdf"
            write_audit_pdf(report, str(pdf_path), model_name=model_name, model_id=model_id)
            if save_residuals and last_residual_dict is not None:
                json_path = session_dir / "residuals.json"
                write_residuals_json(last_residual_dict, str(json_path))
            zip_path = Path(export_dir) / f"{session_name}.zip"
            with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
                for f in session_dir.iterdir():
                    zf.write(f, f"{session_name}/{f.name}")
        except ImportError as e:
            raise ImportError(
                "PDF export requires reportlab. Install with: pip install moju[report] or pip install reportlab"
            ) from e

    return report


def _keys_by_category(plot_keys: Sequence[str]) -> Dict[str, List[str]]:
    order = ("laws", "constitutive", "scaling", "data")
    buckets: Dict[str, List[str]] = {c: [] for c in order}
    for k in plot_keys:
        if "/" not in k:
            continue
        p = k.split("/", 1)[0]
        if p in buckets:
            buckets[p].append(k)
    return buckets


def _parse_spatial_values_panel(panel: Mapping[str, Any], *, law_style: bool) -> Optional[Dict[str, Any]]:
    """
    Parse spatial law or constitutive panel with 1D, 2D, or 3D ``values`` arrays.

    **1D:** ``x`` (length ``nx``) and each ``values[k]`` shape ``(nx,)``. Output includes
    ``kind="1d"``, ``Z`` with shape ``(n_keys, nx)``.

    **2D:** ``x`` (``nx``), ``y`` (``ny``), each ``values[k]`` shape ``(ny, nx)`` with
    entry ``[j, i]`` at ``y[j], x[i]``. Output: ``kind="2d"``, ``Z`` ``(n_keys, ny, nx)``.

    **3D:** ``x``, ``y``, ``z`` 1D and each ``values[k]`` shape ``(nx, ny, nz)`` with
    ``[i, j, k]`` at ``x[i], y[j], z[k]``. Output: ``kind="3d"``, ``V`` ``(n_keys, nx, ny, nz)``.

    Optional ``log_step_index`` (int) is copied through for captions.
    Optional ``position_axis`` is kept for 1D (horizontal axis label hint).
    """
    values = panel.get("values")
    if values is None or not isinstance(values, Mapping):
        return None
    import numpy as np

    keys_sorted = sorted(values.keys())
    if not keys_sorted:
        return None
    arrs = [np.asarray(values[k], dtype=float) for k in keys_sorted]
    if not arrs:
        return None
    s0 = arrs[0].shape
    if not all(a.shape == s0 and a.ndim == arrs[0].ndim for a in arrs):
        return None
    nd = arrs[0].ndim

    log_step_index = panel.get("log_step_index")
    if log_step_index is not None and not isinstance(log_step_index, int):
        log_step_index = None

    def _labels() -> List[str]:
        lab: List[str] = []
        for k in keys_sorted:
            lk = str(k)
            if law_style:
                flat = lk if lk.startswith("laws/") else f"laws/{lk}"
                lab.append(pretty_residual_key(flat))
            else:
                lab.append(pretty_residual_key(lk))
        return lab

    labels = _labels()
    extra: Dict[str, Any] = {"row_labels": labels}
    if log_step_index is not None:
        extra["log_step_index"] = int(log_step_index)

    if nd == 1:
        if panel.get("y") is not None or panel.get("z") is not None:
            return None
        x = panel.get("x")
        if x is None:
            return None
        x_arr = np.asarray(x, dtype=float).ravel()
        if x_arr.size != s0[0]:
            return None
        out: Dict[str, Any] = {
            "kind": "1d",
            "x": x_arr,
            "Z": np.stack(arrs, axis=0),
            **extra,
        }
        pos_ax = panel.get("position_axis")
        if pos_ax is not None and str(pos_ax).strip():
            out["position_axis"] = str(pos_ax).strip()
        return out

    if nd == 2:
        if panel.get("z") is not None:
            return None
        x = panel.get("x")
        y = panel.get("y")
        if x is None or y is None:
            return None
        x_arr = np.asarray(x, dtype=float).ravel()
        y_arr = np.asarray(y, dtype=float).ravel()
        ny, nx = s0
        if ny != y_arr.size or nx != x_arr.size:
            return None
        return {
            "kind": "2d",
            "x": x_arr,
            "y": y_arr,
            "Z": np.stack(arrs, axis=0),
            **extra,
        }

    if nd == 3:
        x = panel.get("x")
        y = panel.get("y")
        z = panel.get("z")
        if x is None or y is None or z is None:
            return None
        x_arr = np.asarray(x, dtype=float).ravel()
        y_arr = np.asarray(y, dtype=float).ravel()
        z_arr = np.asarray(z, dtype=float).ravel()
        nx, ny, nz = s0
        if nx != x_arr.size or ny != y_arr.size or nz != z_arr.size:
            return None
        return {
            "kind": "3d",
            "x": x_arr,
            "y": y_arr,
            "z": z_arr,
            "V": np.stack(arrs, axis=0),
            **extra,
        }

    return None


def _parse_spatial_law_panel(spatial_law_panel: Optional[Any]) -> Optional[Dict[str, Any]]:
    """See :func:`_parse_spatial_values_panel` (``law_style=True``)."""
    if spatial_law_panel is None:
        return None
    if not isinstance(spatial_law_panel, Mapping):
        return None
    return _parse_spatial_values_panel(spatial_law_panel, law_style=True)


def _parse_spatial_rnorm_panel(spatial_rnorm_panel: Optional[Any]) -> Optional[Dict[str, Any]]:
    """See :func:`_parse_spatial_values_panel` (``law_style=False``)."""
    if spatial_rnorm_panel is None:
        return None
    if not isinstance(spatial_rnorm_panel, Mapping):
        return None
    return _parse_spatial_values_panel(spatial_rnorm_panel, law_style=False)


def _build_visualize_bundle(
    log: List[Dict[str, Any]],
    keys: Optional[List[str]],
    r_ref: Optional[Dict[str, float]],
    max_legend_keys: int,
    *,
    spatial_parsed: Optional[Dict[str, Any]],
    spatial_rnorm_parsed: Optional[Dict[str, Any]] = None,
    mode: str,
    spatial_normalize: bool = False,
) -> Optional[Dict[str, Any]]:
    """
    Shared arrays and metadata for :func:`visualize` (Plotly).

    Does not mutate ``log``. Requires ``numpy`` (already a moju dependency).
    """
    import numpy as np

    if not log:
        return None
    first_rms = log[0].get("rms", {})
    plot_keys = list(keys) if keys is not None else list(first_rms.keys())
    if not plot_keys:
        return None
    metrics = _compute_log_step_metrics(log, r_ref)
    n = len(log)
    indices = np.arange(n, dtype=float)
    cap = max(1, int(max_legend_keys))
    use_bar_chart = mode == "test" or (mode == "training" and n == 1)
    bar_keys = plot_keys[: min(48, len(plot_keys))]

    buckets = _keys_by_category(plot_keys)
    category_training: Dict[str, Dict[str, Any]] = {}
    cat_order = ("laws", "constitutive") if mode == "training" else ("laws", "constitutive", "scaling")
    for cat in cat_order:
        cat_keys = sorted(buckets.get(cat, []))[:cap]
        nk = len(cat_keys)
        mat = np.zeros((nk, n)) if nk else np.zeros((0, n))
        for j in range(n):
            for i, kk in enumerate(cat_keys):
                v = metrics[j]["r_norm"].get(kk)
                mat[i, j] = float(v) if v is not None else float("nan")
        category_training[cat] = {
            "keys": cat_keys,
            "displays": [pretty_residual_key(k) for k in cat_keys],
            "r_norm_mat": mat,
        }

    legend_keys = plot_keys[: min(cap, len(plot_keys))]
    r_norm_mat = np.zeros((len(legend_keys), n))
    for j in range(n):
        for i, kk in enumerate(legend_keys):
            v = metrics[j]["r_norm"].get(kk)
            r_norm_mat[i, j] = float(v) if v is not None else float("nan")

    legend_display = [pretty_residual_key(k) for k in legend_keys]
    bar_display = [pretty_residual_key(k) for k in bar_keys]
    bar_values = []
    last_rn = metrics[-1]["r_norm"]
    for kk in bar_keys:
        v = last_rn.get(kk)
        bar_values.append(float(v) if v is not None else float("nan"))
    bar_values_arr = np.asarray(bar_values, dtype=float)

    cat_colors = {
        "laws": "#4e79a7",
        "constitutive": "#f28e2b",
        "scaling": "#59a14f",
        "data": "#b07aa1",
    }
    last_cat = metrics[-1]["category_admissibility_score"]
    cats_fin = [
        (pretty_category_name(c), float(last_cat[c]))
        for c in cat_colors
        if c in last_cat and np.isfinite(last_cat[c])
    ]
    nr_title = "Normalized Residuals"
    if not use_bar_chart and len(plot_keys) > len(legend_keys):
        nr_title = f"Normalized Residuals (showing {len(legend_keys)} of {len(plot_keys)} keys)"

    category_titles = {
        "laws": "Normalized Governing Laws Residuals",
        "constitutive": "Normalized Constitutive Residuals",
        "scaling": "Normalized Scaling Residuals",
    }

    return {
        "log": log,
        "metrics": metrics,
        "n": n,
        "indices": indices,
        "plot_keys": plot_keys,
        "legend_keys": legend_keys,
        "legend_display": legend_display,
        "r_norm_mat": r_norm_mat,
        "category_training": category_training,
        "category_titles": category_titles,
        "use_bar_chart": use_bar_chart,
        "bar_keys": bar_keys,
        "bar_display": bar_display,
        "bar_values": bar_values_arr,
        "overall_adm": [float(metrics[i]["overall_admissibility_score"]) for i in range(n)],
        "cats_fin": cats_fin,
        "cat_colors": cat_colors,
        "spatial": spatial_parsed,
        "spatial_rnorm": spatial_rnorm_parsed,
        "mode": mode,
        "nr_title": nr_title,
        "np": np,
        "spatial_normalize": bool(spatial_normalize),
    }


def _maybe_build_spatial_panels(
    log_full: List[Dict[str, Any]],
    spatial_law_panel: Optional[Dict[str, Any]],
    spatial_rnorm_panel: Optional[Dict[str, Any]],
    residuals: Optional[Dict[str, Any]],
    state_pred: Optional[Mapping[str, Any]],
    r_ref: Optional[Dict[str, float]],
    spatial_coord_key: str,
    spatial_prefer_last_t: bool,
    *,
    spatial_normalize: bool = False,
) -> tuple[Optional[Dict[str, Any]], Optional[Dict[str, Any]]]:
    if residuals is None:
        return spatial_law_panel, spatial_rnorm_panel

    pred: Dict[str, Any] = {}
    if state_pred is not None:
        pred.update(dict(state_pred))
    if log_full:
        snap = log_full[-1].get("coord_snapshot")
        if isinstance(snap, Mapping):
            for k in ("x", "y", "z", "t"):
                if k in snap and snap[k] is not None and k not in pred:
                    pred[k] = snap[k]

    auto_law: Optional[Dict[str, Any]] = None
    auto_rnorm: Optional[Dict[str, Any]] = None
    log_entry = log_full[-1] if log_full else None
    first_rms = (log_full[0].get("rms") or {}) if log_full else {}
    log_step_index = len(log_full) - 1 if log_full else None
    rr = dict(r_ref) if r_ref else {}
    auto_law, auto_rnorm = build_spatial_rnorm_panels_from_residuals(
        residuals,
        pred,
        coord_key=str(spatial_coord_key),
        prefer_last_t=bool(spatial_prefer_last_t),
        log_entry=log_entry,
        first_rms=first_rms,
        r_ref=rr,
        log_step_index=log_step_index,
        normalize_spatial=bool(spatial_normalize),
    )

    law_out = spatial_law_panel if spatial_law_panel is not None else auto_law
    rnorm_out = spatial_rnorm_panel if spatial_rnorm_panel is not None else auto_rnorm
    return law_out, rnorm_out


def build_monitor_visualize_bundle(
    log: List[Dict[str, Any]],
    keys: Optional[List[str]] = None,
    r_ref: Optional[Dict[str, float]] = None,
    max_legend_keys: int = 16,
    *,
    spatial_law_panel: Optional[Dict[str, Any]] = None,
    spatial_rnorm_panel: Optional[Dict[str, Any]] = None,
    mode: str = "training",
    residuals: Optional[Dict[str, Any]] = None,
    state_pred: Optional[Mapping[str, Any]] = None,
    spatial_coord_key: str = "x",
    spatial_prefer_last_t: bool = True,
    spatial_normalize: bool = False,
    engine: Optional[Any] = None,
) -> Optional[Dict[str, Any]]:
    """
    Build the internal visualization bundle (same as :func:`visualize` uses for Plotly).

    Intended for Studio and other callers that want several small Plotly figures instead of
    one combined subplot grid.

    ``spatial_normalize``: if True, auto-built spatial panels use R_norm-style :math:`|r|/s_k`;
    default False uses absolute :math:`|r|` vs position (labels match).

    ``engine``: when ``residuals`` is omitted, use :attr:`ResidualEngine.last_residuals` if given.
    """
    eff_residuals = residuals
    if eff_residuals is None and engine is not None:
        eff_residuals = getattr(engine, "last_residuals", None)
    law_panel, rnorm_panel = _maybe_build_spatial_panels(
        log,
        spatial_law_panel,
        spatial_rnorm_panel,
        eff_residuals,
        state_pred,
        r_ref,
        spatial_coord_key,
        spatial_prefer_last_t,
        spatial_normalize=spatial_normalize,
    )
    work_log = list(log)
    if mode == "test" and len(work_log) > 1:
        work_log = work_log[-1:]
    spatial_parsed = _parse_spatial_law_panel(law_panel)
    spatial_rnorm_parsed = _parse_spatial_rnorm_panel(rnorm_panel)
    return _build_visualize_bundle(
        work_log,
        keys,
        r_ref,
        max_legend_keys,
        spatial_parsed=spatial_parsed,
        spatial_rnorm_parsed=spatial_rnorm_parsed,
        mode=mode,
        spatial_normalize=spatial_normalize,
    )


def visualize(
    log: List[Dict[str, Any]],
    keys: Optional[List[str]] = None,
    backend: str = "plotly",
    *,
    r_ref: Optional[Dict[str, float]] = None,
    max_legend_keys: int = 16,
    mode: str = "training",
    spatial_law_panel: Optional[Dict[str, Any]] = None,
    spatial_rnorm_panel: Optional[Dict[str, Any]] = None,
    residuals: Optional[Dict[str, Any]] = None,
    state_pred: Optional[Mapping[str, Any]] = None,
    spatial_coord_key: str = "x",
    spatial_prefer_last_t: bool = True,
    spatial_normalize: bool = False,
    engine: Optional[Any] = None,
    figure_title: Optional[str] = None,
    step_label: str = "Step",
    r_norm_scale: str = "log",
    spatial_heatmap_colorscale: Optional[str] = None,
) -> Any:
    """
    Monitor dashboard from ``ResidualEngine`` log entries (``rms``, ``scale``).

    Uses the same R_norm / admissibility rules as :func:`audit` via
    :func:`_compute_log_step_metrics` and **does not mutate** ``log``.

    **Modes**

    - ``training`` (multi-step) — **Top row:** overall admissibility vs step (with final
      value marker) and **horizontal category admissibility bars** (laws / constitutive,
      final step). **Second row:** two panels — :math:`R_{\\mathrm{norm}}` vs
      step for **governing laws** and **constitutive** (``data/`` and ``scaling/`` omitted);
      **y-axis** is ``log10(R_{\\mathrm{norm}} + \\varepsilon)`` by default, or linear if
      ``r_norm_scale="linear"``. **Third row:** **vs spatial coordinate** (last logged step)
      for laws and constitutive: by default :math:`|r|` (**absolute residual**); set
      ``spatial_normalize=True`` for :math:`R_{\\mathrm{norm}}`-style :math:`|r|/s_k`.
      Placeholders if no spatial data. Pass ``spatial_*_panel``, or ``residuals`` with coordinates from
      ``state_pred``, or ``residuals`` with the **last** log entry's ``coord_snapshot``
      (written by :meth:`ResidualEngine.compute_residuals` when the built state has ``x`` /
      ``y`` / ``z`` / ``t``).
    - ``training`` (single log entry) — horizontal bars for normalized residuals, category
      admissibility bars, spatial row (|residual| vs position by default).
    - ``test`` — Uses the **last** log entry for scalar metrics; **spatial row** shows
      |residual| vs position by default (same log/linear display scale as training for heatmaps), plus
      horizontal bar chart of normalized residuals and category admissibility.

    **Spatial panel**

    The log stores **scalar** RMS per key per step only. Pass panels built from the
    **last** ``compute_residuals`` / final log step (callers may set optional
    ``log_step_index`` on the dict for captions).

    **1D:** ``x`` (length ``nx``) and ``values`` with 1D arrays ``(nx,)``. Law keys may be
    bare names or ``laws/<name>``. Optional ``position_axis`` sets the horizontal axis
    label (default ``x``).

    **2D:** ``x``, ``y`` 1D and each ``values[k]`` with shape ``(len(y), len(x))``.

    **3D:** ``x``, ``y``, ``z`` 1D and each ``values[k]`` with shape
    ``(len(x), len(y), len(z))``.

    For constitutive-only slices use ``spatial_rnorm_panel`` (flat keys). **Plotly**
    renders 1D/2D/3D spatial panels.

    **Backends**

    - ``plotly`` (default) — interactive figure (requires ``pip install plotly`` or ``moju[viz]``).
    - ``none`` — returns ``None``.

    ``backend="matplotlib"`` is **not supported** (raises ``ValueError``); use ``plotly``.

    Parameters
    ----------
    log
        Entries from ``ResidualEngine.log`` (after ``compute_residuals``).
    keys
        Subset of flat residual keys to plot; default = all keys in the first entry.
    backend
        ``plotly`` (default) or ``none``.
    r_ref
        Optional per-key reference scale overrides (same as :func:`audit`).
    max_legend_keys
        Cap legend entries on per-key line plots for readability (training mode, multi-step).
    mode
        ``training`` or ``test``. Test mode slices to the last entry when ``len(log) > 1``.
    spatial_law_panel
        Optional ``dict`` with ``x`` and ``values`` (see above).
    spatial_rnorm_panel
        Optional ``dict`` with ``x`` and ``values`` for per-key R_norm spatial slices.
    residuals
        Optional nested residual dict (same shape as ``compute_residuals`` output). When
        both ``spatial_law_panel`` and ``spatial_rnorm_panel`` are omitted, panels are built
        for the **last** log step from ``state_pred`` coordinates and/or the last entry's
        ``coord_snapshot`` (lists of floats), merged so snapshot fills only missing axes.
        If coordinates are still missing but law/constitutive arrays share a consistent 1D
        length, a neutral ``linspace(0, 1, n)`` axis is inferred so ``residuals=`` alone can
        suffice (e.g. ``visualize(log, mode=\"test\", residuals=...)``).
    state_pred
        Optional state dict with coordinate arrays (``x``, optional ``y`` / ``z``, ``t``) for
        auto spatial panels. May be omitted or partial when ``log[-1]["coord_snapshot"]``
        supplies the missing axes.
    spatial_coord_key
        Primary 1D axis name when falling back from 3D/2D to a line slice (default ``"x"``).
    spatial_prefer_last_t
        If True and ``t`` is present, use the last time slice of residual fields.
    spatial_normalize
        If True, auto-built spatial panels use :math:`|r|/s_k` (R_norm-style). Default False
        uses absolute residual :math:`|r|` vs position.
    engine
        Optional :class:`ResidualEngine` whose :attr:`~ResidualEngine.last_residuals` is used
        when ``residuals`` is omitted, so ``visualize(engine.log, engine=engine)`` can build
        spatial heatmaps using log ``coord_snapshot`` (same convenience as Moju Studio).
    figure_title
        Optional override for the figure title. If omitted or blank, a mode-specific
        default is used (training vs test).
    step_label
        X-axis label for training step axis (e.g. ``Iteration`` or ``Epoch``).
    r_norm_scale
        ``log`` (default) plots ``log10(R_norm + ε)`` on the three category residual
        panels; ``linear`` plots raw ``R_norm``. Does not affect the overall admissibility
        axis.
    spatial_heatmap_colorscale
        Plotly colorscale name for optional spatial heatmaps (e.g. ``\"Jet\"``, ``\"Viridis\"``).
        Default ``None`` uses ``Jet`` in the Plotly backend.
    """
    if backend == "none":
        return None
    if backend == "matplotlib":
        raise ValueError(
            'visualize(..., backend="matplotlib") is no longer supported; '
            'use backend="plotly" (default) and pip install plotly or moju[viz].'
        )
    if backend != "plotly":
        raise ValueError(f"Unknown visualize backend {backend!r}; use 'plotly' or 'none'.")
    if mode not in ("training", "test"):
        raise ValueError("mode must be 'training' or 'test'")
    if r_norm_scale not in ("log", "linear"):
        raise ValueError("r_norm_scale must be 'log' or 'linear'")

    eff_residuals = residuals
    if eff_residuals is None and engine is not None:
        eff_residuals = getattr(engine, "last_residuals", None)

    law_panel, rnorm_panel = _maybe_build_spatial_panels(
        log,
        spatial_law_panel,
        spatial_rnorm_panel,
        eff_residuals,
        state_pred,
        r_ref,
        spatial_coord_key,
        spatial_prefer_last_t,
        spatial_normalize=spatial_normalize,
    )

    work_log = list(log)
    if mode == "test" and len(work_log) > 1:
        work_log = work_log[-1:]

    spatial_parsed = _parse_spatial_law_panel(law_panel)
    spatial_rnorm_parsed = _parse_spatial_rnorm_panel(rnorm_panel)
    bundle = _build_visualize_bundle(
        work_log,
        keys,
        r_ref,
        max_legend_keys,
        spatial_parsed=spatial_parsed,
        spatial_rnorm_parsed=spatial_rnorm_parsed,
        mode=mode,
        spatial_normalize=spatial_normalize,
    )
    if bundle is None:
        return None

    resolved_title = _resolve_visualize_figure_title(mode, figure_title)

    try:
        from moju.monitor.visualize_plotly import build_plotly_monitor_figure

        return build_plotly_monitor_figure(
            bundle,
            figure_title=resolved_title,
            step_label=step_label,
            r_norm_scale=r_norm_scale,
            spatial_heatmap_colorscale=spatial_heatmap_colorscale,
        )
    except ImportError:
        return None


def _coord_snapshot_from_merged(merged: Mapping[str, Any]) -> Dict[str, List[float]]:
    """
    Extract 1D ``x`` / ``y`` / ``z`` / ``t`` vectors from built state for JSON-safe log storage.

    Used for :attr:`ResidualEngine.log` ``coord_snapshot`` so :func:`visualize` can build
    position heatmaps when ``state_pred`` is omitted but the last log step recorded coords.
    """
    import numpy as np

    out: Dict[str, List[float]] = {}
    for key in ("x", "y", "z", "t"):
        if key not in merged:
            continue
        raw = merged.get(key)
        if raw is None:
            continue
        try:
            arr = np.asarray(jax.device_get(jnp.asarray(raw)), dtype=float).ravel()
        except (TypeError, ValueError):
            continue
        if arr.size == 0:
            continue
        out[key] = [float(v) for v in arr.tolist()]
    return out


class ResidualEngine:
    """
    Governing laws (Laws.*), optional group specs to enrich state, and model/group closures.

    Entry points:
      - Path A (recommended): provide (model, params, collocation) and a state_builder
        so moju can build state_pred (and derivatives) internally.
      - Path B (advanced): provide state_pred directly.

    Closure policy:
      - **ref_delta** runs when ``state_ref`` is provided, unless the spec sets ``include_ref_delta=False``.
      - **implied_delta** runs for **constitutive** specs when ``implied_value_key`` or ``implied_fn``
        is set; omitted if implied is missing. **Scaling** audits support **ref_delta** and **π-constant**
        (Path A) only—not **implied_delta**.
      - **ref_delta** / **implied_delta** residuals are nondimensional (see ``closure_registry``).
      - A spec with no applicable closure does nothing (optional omit log).
      - Law-linked implied rows (see ``moju.monitor.law_implied_diagnostics``) are prepended when
        ``law_implied_audits`` is true (``MonitorConfig`` default). Use optional ``residual_basename``
        for unique flat keys under each category.

    Audit spec shape (constitutive_audit / scaling_audit items):
      {
        "name": "sutherland_mu",               # Models.<name> or Groups.<name>
        "output_key": "mu",                    # state key for F output (ref_delta / implied_delta)
        "state_map": {"T": "T", "mu0": "mu0", "T0": "T0", "S": "S"},  # fn arg -> state key
      }

    Each :meth:`compute_residuals` log entry may include ``coord_snapshot``: a dict with optional
    keys ``x``, ``y``, ``z``, ``t`` (1D float lists) copied from built state when present, for
    spatial dashboards.

    :attr:`last_residuals` references the nested dict from the latest successful
    :meth:`compute_residuals` (``None`` after :meth:`clear_log`) so :func:`visualize` can run as
    ``visualize(engine.log, engine=engine)`` without a separate ``residuals=`` argument.
    """

    def __init__(
        self,
        config: Optional[Any] = None,
        constants: Optional[Dict[str, Any]] = None,
        laws: Optional[List[Dict[str, Any]]] = None,
        groups: Optional[List[Dict[str, Any]]] = None,
        *,
        constitutive_audit: Optional[List[Dict[str, Any]]] = None,
        scaling_audit: Optional[List[Dict[str, Any]]] = None,
        constitutive_custom: Optional[List[Dict[str, Any]]] = None,
        derived_state_chain: Optional[List[Dict[str, Any]]] = None,
        state_builder: Optional[
            Callable[[Any, Any, Dict[str, Any], Dict[str, Any]], Dict[str, Any]]
        ] = None,
        enable_omit_messages: bool = True,
        primary_fields: Optional[List[str]] = None,
        law_implied_audits: bool = True,
    ):
        law_implied_enabled = bool(law_implied_audits)
        # MonitorConfig convenience
        if config is not None:
            from moju.monitor.config import MonitorConfig, audit_spec_to_engine_dict

            if isinstance(config, MonitorConfig):
                constants = config.constants
                laws = config.laws
                groups = config.groups
                constitutive_audit = [audit_spec_to_engine_dict(s) for s in config.constitutive_audit]
                scaling_audit = [audit_spec_to_engine_dict(s) for s in config.scaling_audit]
                constitutive_custom = config.constitutive_custom
                derived_state_chain = list(config.derived_state_chain or [])
                primary_fields = list(config.primary_fields)
                law_implied_enabled = bool(config.law_implied_audits)
                if config.state_builder is not None and state_builder is None:
                    state_builder = config.state_builder
            else:
                raise TypeError("config must be a MonitorConfig")

        self.constants = dict(constants or {})
        self.laws_spec = list(laws or [])
        self.groups_spec = list(groups or [])
        self.constitutive_audit = list(constitutive_audit or [])
        self.scaling_audit = list(scaling_audit or [])
        li_c, li_s = merge_law_implied_audit_specs(self.laws_spec, enabled=law_implied_enabled)
        mc, rc = merge_fragment_law_implied_audit_specs(li_c, self.constitutive_audit)
        ms, rs = merge_fragment_law_implied_audit_specs(li_s, self.scaling_audit)
        self.constitutive_audit = mc + rc
        self.scaling_audit = ms + rs
        self.constitutive_custom = list(constitutive_custom or [])
        self.derived_state_chain = list(derived_state_chain or [])
        self.state_builder = state_builder
        self.enable_omit_messages = bool(enable_omit_messages)
        self.primary_fields = list(primary_fields or ["T", "u", "v", "w", "p", "rho"])

        # Config-time validation (low effort)
        def _validate_specs(specs: Sequence[Dict[str, Any]], registry: Dict[str, Any], category: str) -> None:
            for spec in specs:
                if "name" not in spec:
                    raise ValueError(f"{category} spec missing 'name'")
                name = spec["name"]
                reg = registry.get(name)
                if reg is None:
                    raise ValueError(f"{category} spec name {name!r} is not registered")
                if "output_key" not in spec:
                    raise ValueError(f"{category}:{name} missing 'output_key'")
                if "state_map" not in spec or not isinstance(spec["state_map"], dict):
                    raise ValueError(f"{category}:{name} missing 'state_map' dict")
                _, arg_names = reg
                missing_args = [an for an in arg_names if an not in spec["state_map"]]
                if missing_args:
                    raise ValueError(f"{category}:{name} state_map missing args: {missing_args}")
                ivk = spec.get("implied_value_key")
                ifn = spec.get("implied_fn")
                if ivk and ifn is not None:
                    raise ValueError(
                        f"{category}:{name} use only one of implied_value_key and implied_fn, not both"
                    )
        _validate_specs(self.constitutive_audit, MODEL_FNS, "constitutive")
        _validate_specs(self.scaling_audit, GROUP_FNS, "scaling")
        self._validate_pi_constant_specs()

        self._log: List[Dict[str, Any]] = []
        self._index = 0
        self._last_residuals: Optional[Dict[str, Any]] = None

    @property
    def log(self) -> List[Dict[str, Any]]:
        return self._log

    @property
    def last_residuals(self) -> Optional[Dict[str, Any]]:
        """Nested residual dict from the latest successful :meth:`compute_residuals` (same object)."""
        return self._last_residuals

    def clear_log(self) -> None:
        """Remove all logged steps and reset the step counter.

        Safe to call between training runs; does not alter engine configuration.
        """
        self._log.clear()
        self._index = 0
        self._last_residuals = None

    def _validate_pi_constant_specs(self) -> None:
        for spec in self.scaling_audit:
            if not spec.get("invariance_pi_constant"):
                continue
            name = spec["name"]
            if name not in GROUP_PI_CONSTANT_RECIPES:
                raise ValueError(
                    f"scaling:{name} invariance_pi_constant requires a built-in π-constant recipe; "
                    f"supported: {sorted(GROUP_PI_CONSTANT_RECIPES.keys())}"
                )
            recipe = GROUP_PI_CONSTANT_RECIPES[name]
            sm = spec.get("state_map") or {}
            for arg_name, _ in recipe:
                if arg_name not in sm:
                    raise ValueError(
                        f"scaling:{name} π-constant recipe needs state_map entry for arg {arg_name!r}"
                    )
            cmp_keys = list(spec.get("invariance_compare_keys") or [])
            if not cmp_keys:
                raise ValueError(
                    f"scaling:{name} invariance_pi_constant requires non-empty invariance_compare_keys"
                )
            c = float(spec.get("invariance_scale_c", 10.0))
            if c <= 1.0:
                raise ValueError(f"scaling:{name} invariance_scale_c must be > 1, got {c}")
            for arg_name, _ in recipe:
                sk = sm[arg_name]
                if sk not in self.constants:
                    raise ValueError(
                        f"scaling:{name} π-constant requires key {sk!r} in ResidualEngine.constants "
                        f"(arg {arg_name!r})"
                    )
            if self.state_builder is None:
                raise ValueError(
                    f"scaling:{name} invariance_pi_constant requires ResidualEngine(state_builder=...) (Path A only)"
                )

    def _state_builder(
        self,
        state_pred: Dict[str, Any],
        constants_override: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        c = self.constants if constants_override is None else constants_override
        return _build_state(state_pred, c, self.groups_spec)

    def compute_residuals(
        self,
        state_pred: Optional[Dict[str, Any]] = None,
        state_ref: Optional[Dict[str, Any]] = None,
        *,
        model: Any = None,
        params: Any = None,
        collocation: Optional[Dict[str, Any]] = None,
        log_to_python: bool = True,
        auto_path_b_derivatives: Any = False,
        fill_law_fd: bool = False,
    ) -> Dict[str, Any]:
        """
        Compute residuals.

        Path A: pass (model, params, collocation) and configure engine.state_builder.
        Path B: pass state_pred directly.

        If ``auto_path_b_derivatives`` is True, uses default ``PathBGridConfig``; if a
        ``PathBGridConfig`` instance, uses that layout. When ``fill_law_fd`` is also True, missing
        **registered** ``Laws.*`` inputs (e.g. ``phi_laplacian``, ``u_grad``) are filled from
        primitives on the same grid via finite differences (see ``law_fd_recipes``).
        Warnings are appended to the log ``inferred`` list when enabled.

        If ``fill_law_fd`` is True, ``auto_path_b_derivatives`` must also be enabled.
        """
        path_a = state_pred is None
        pi_specs = [s for s in self.scaling_audit if s.get("invariance_pi_constant")]
        if pi_specs and not path_a:
            raise ValueError(
                "π-constant scaling audit (invariance_pi_constant) requires Path A: "
                "call compute_residuals(..., model=..., params=..., collocation=...) "
                "without passing state_pred."
            )

        residuals: Dict[str, Any] = {"laws": {}}
        pb_warn: List[str] = []

        if state_pred is None:
            if self.state_builder is None:
                raise ValueError("Path A requires ResidualEngine(state_builder=...)")
            if model is None or params is None or collocation is None:
                raise ValueError("Path A requires model, params, and collocation")
            state_pred = self.state_builder(model, params, collocation, self.constants)

        omitted_msgs: List[str] = []
        inferred_msgs: List[str] = []

        def _maybe_log_omit(msg: str) -> None:
            if self.enable_omit_messages:
                omitted_msgs.append(msg)

        def _maybe_log_infer(msg: str) -> None:
            if self.enable_omit_messages:
                inferred_msgs.append(msg)

        state_for_groups = dict(state_pred)
        if self.derived_state_chain:
            from moju.monitor.derived_state_chain import apply_derived_state_chain

            state_for_groups, _dwarn = apply_derived_state_chain(
                state_for_groups,
                self.constants,
                self.derived_state_chain,
            )
            for w in _dwarn:
                _maybe_log_infer(f"derived_state: {w}")

        state_pred_built = self._state_builder(state_for_groups)
        merged = {**self.constants, **state_pred_built}

        def _state_ref_raw_after_derived(sr: Dict[str, Any]) -> Dict[str, Any]:
            s0 = dict(sr)
            if self.derived_state_chain:
                from moju.monitor.derived_state_chain import apply_derived_state_chain

                s0, wr = apply_derived_state_chain(s0, self.constants, self.derived_state_chain)
                for x in wr:
                    _maybe_log_infer(f"derived_state(ref): {x}")
            return s0

        def _merge_state_ref(sr: Dict[str, Any]) -> Dict[str, Any]:
            return {**self._state_builder(_state_ref_raw_after_derived(sr)), **self.constants}

        if fill_law_fd and not auto_path_b_derivatives:
            raise ValueError(
                "fill_law_fd=True requires auto_path_b_derivatives=True or a PathBGridConfig"
            )

        if auto_path_b_derivatives:
            from moju.monitor.path_b_derivatives import PathBGridConfig, fill_path_b_derivatives

            if auto_path_b_derivatives is True:
                grid = PathBGridConfig()
            elif isinstance(auto_path_b_derivatives, PathBGridConfig):
                grid = auto_path_b_derivatives
            else:
                raise TypeError(
                    "auto_path_b_derivatives must be False, True, or a PathBGridConfig instance"
                )
            state_pred_built, pb_warn = fill_path_b_derivatives(
                state_pred_built,
                constitutive_audit=self.constitutive_audit,
                scaling_audit=self.scaling_audit,
                laws_spec=self.laws_spec,
                constants=self.constants,
                grid=grid,
                copy=False,
                fill_law_recipes=bool(fill_law_fd),
            )
            merged = {**self.constants, **state_pred_built}
            for w in pb_warn:
                _maybe_log_infer(f"path_b_derivatives: {w}")

        for spec in self.laws_spec:
            name = spec["name"]
            state_map = spec["state_map"]
            try:
                kwargs = _kwargs_from_state(
                    merged, self.constants, state_map, law_context=str(name)
                )
            except KeyError as err:
                if pb_warn:
                    tail = "\n".join(f"  - {w}" for w in pb_warn[:32])
                    raise KeyError(
                        f"{err.args[0]}\n\nFinite-difference / law-FD fill reported:\n{tail}"
                    ) from err
                raise
            fn = _get_fn(spec, Laws)
            residuals["laws"][name] = fn(**kwargs)

        def _run_specs(
            specs: Sequence[Dict[str, Any]],
            *,
            registry: Dict[str, Any],
            category: str,
        ) -> Dict[str, Any]:
            out: Dict[str, Any] = {}
            for spec in specs:
                name = spec["name"]
                output_key = spec.get("output_key")
                state_map = spec.get("state_map") or {}
                has_implied = bool(spec.get("implied_value_key")) or spec.get("implied_fn") is not None
                if state_ref is None and not has_implied:
                    _maybe_log_omit(
                        f"{category}:{name} omitted: no ref_delta or implied_delta applicable"
                    )
                    continue

                reg = registry.get(name)
                if reg is None:
                    # unknown function name -> omit silently (config validation should catch)
                    continue
                fn, arg_names = reg
                base = spec.get("residual_basename") or name

                if (
                    state_ref is not None
                    and output_key is not None
                    and spec.get("include_ref_delta", True)
                ):
                    arr = compute_ref_delta(
                        fn=fn,
                        arg_names=arg_names,
                        output_key=output_key,
                        state_map=state_map,
                        state_pred=merged,
                        state_ref=_merge_state_ref(state_ref),
                        constants=self.constants,
                        ref_delta_ref_key=spec.get("ref_delta_ref_key"),
                    )
                    if arr is not None:
                        out[f"{base}/ref_delta"] = jnp.asarray(arr)

                if has_implied:
                    if category == "scaling":
                        # Focused clean-out: scaling audits support ref_delta (Path A/B) and pi_constant (Path A) only.
                        _maybe_log_omit(f"{category}:{name} omitted: implied_delta is not supported for scaling audits")
                        continue
                    arr = compute_implied_delta(
                        fn=fn,
                        arg_names=arg_names,
                        state_map=state_map,
                        state_pred=merged,
                        constants=self.constants,
                        implied_value_key=spec.get("implied_value_key"),
                        implied_fn=spec.get("implied_fn"),
                        output_key=output_key,
                        implied_delta_ref_key=spec.get("implied_delta_ref_key"),
                    )
                    if arr is not None:
                        out[f"{base}/implied_delta"] = jnp.asarray(arr)

            return out

        if self.constitutive_audit or self.constitutive_custom:
            c = _run_specs(self.constitutive_audit, registry=MODEL_FNS, category="constitutive")
            if self.constitutive_custom:
                for spec in self.constitutive_custom:
                    cname = spec["name"]
                    arr = spec["fn"](merged, self.constants)
                    if arr is not None:
                        c[f"custom/{cname}"] = jnp.asarray(arr)
            if c:
                residuals["constitutive"] = c

        pi_constant_scales: Dict[str, float] = {}
        if self.scaling_audit:
            s = _run_specs(self.scaling_audit, registry=GROUP_FNS, category="scaling")
            if path_a and self.state_builder is not None:
                for spec in self.scaling_audit:
                    if not spec.get("invariance_pi_constant"):
                        continue
                    name = spec["name"]
                    c = float(spec.get("invariance_scale_c", 10.0))
                    if c <= 1.0:
                        raise ValueError(f"scaling:{name} invariance_scale_c must be > 1, got {c}")
                    recipe = GROUP_PI_CONSTANT_RECIPES[name]
                    state_map = spec.get("state_map") or {}
                    constants_scaled = apply_pi_constant_recipe(
                        self.constants, recipe, state_map, c
                    )
                    state_pred_pi = self.state_builder(model, params, collocation, constants_scaled)
                    merged_scaled = {
                        **self._state_builder(state_pred_pi, constants_scaled),
                        **constants_scaled,
                    }
                    fn, arg_names = GROUP_FNS[name]
                    kb = _kwargs_from_state(merged, self.constants, state_map)
                    ks = _kwargs_from_state(merged_scaled, constants_scaled, state_map)
                    val_b = fn(**{an: kb[an] for an in arg_names})
                    val_s = fn(**{an: ks[an] for an in arg_names})
                    if not jnp.allclose(
                        jnp.asarray(val_b), jnp.asarray(val_s), rtol=1e-4, atol=1e-6
                    ):
                        raise ValueError(
                            f"scaling:{name} π-constant scaled inputs did not preserve group value "
                            f"(baseline {val_b!r} vs scaled {val_s!r})"
                        )
                    compare_keys = list(spec.get("invariance_compare_keys") or [])
                    parts_r: List[jnp.ndarray] = []
                    parts_scale: List[jnp.ndarray] = []
                    for ck in compare_keys:
                        if ck not in merged:
                            raise KeyError(
                                f"scaling:{name} invariance_compare_keys: {ck!r} missing from baseline merged state"
                            )
                        if ck not in merged_scaled:
                            raise KeyError(
                                f"scaling:{name} invariance_compare_keys: {ck!r} missing from scaled merged state"
                            )
                        vb = jnp.asarray(merged[ck])
                        vs = jnp.asarray(merged_scaled[ck])
                        parts_r.append(jnp.ravel(vs - vb))
                        parts_scale.append(jnp.ravel(jnp.abs(vs)))
                    r_pi = jnp.concatenate(parts_r) if parts_r else jnp.array([0.0])
                    stacked_abs = jnp.concatenate(parts_scale) if parts_scale else jnp.array([0.0])
                    flat_pi = f"{name}/pi_constant"
                    s[flat_pi] = r_pi
                    scale_key = f"scaling/{flat_pi}"
                    mean_abs = float(jax.device_get(jnp.nanmean(stacked_abs)))
                    pi_constant_scales[scale_key] = float(_SCALE_EPS + mean_abs)
            if s:
                residuals["scaling"] = s

        if state_ref is not None:
            state_ref_built = self._state_builder(_state_ref_raw_after_derived(state_ref))
            common = set(state_pred_built.keys()) & set(state_ref_built.keys())
            residuals["data"] = {
                k: jnp.asarray(state_ref_built[k]) - jnp.asarray(state_pred_built[k])
                for k in common
            }

        flat = _flatten_residual_dict(residuals)
        rms_per_key = _rms_per_key(flat, to_python=log_to_python)
        state_ref_built_for_scale = None
        if state_ref is not None:
            state_ref_built_for_scale = self._state_builder(_state_ref_raw_after_derived(state_ref))
        scale_per_key = _state_derived_scale_per_key(
            flat.keys(),
            merged,
            self.laws_spec,
            self.constitutive_audit,
            self.scaling_audit,
            state_ref_built_for_scale,
            to_python=log_to_python,
        )
        scale_per_key.update(pi_constant_scales)
        entry: Dict[str, Any] = {"index": self._index, "rms": rms_per_key, "scale": scale_per_key}
        if omitted_msgs:
            entry["omitted"] = omitted_msgs
        if inferred_msgs:
            entry["inferred"] = inferred_msgs
        if "coord_snapshot" not in entry:
            cs = _coord_snapshot_from_merged(merged)
            if cs:
                entry["coord_snapshot"] = cs
        self._log.append(entry)
        self._index += 1
        self._last_residuals = residuals
        return residuals

    def required_state_keys(
        self,
        *,
        include_groups: bool = True,
        include_laws: bool = True,
        include_audits: bool = True,
    ) -> Set[str]:
        keys: Set[str] = set()
        if include_laws:
            for spec in self.laws_spec:
                keys |= set((spec.get("state_map") or {}).values())
        if include_groups:
            for spec in self.groups_spec:
                keys |= set((spec.get("state_map") or {}).values())
                ok = spec.get("output_key")
                if ok:
                    keys.add(ok)
        if include_audits:
            for spec in self.constitutive_audit + self.scaling_audit:
                keys |= set((spec.get("state_map") or {}).values())
                ok = spec.get("output_key")
                if ok:
                    keys.add(ok)
                ivk = spec.get("implied_value_key")
                if ivk:
                    keys.add(ivk)
        keys |= all_ref_keys_from_chain(self.derived_state_chain)
        keys -= keys_produced_by_chain(self.derived_state_chain)
        return keys



def list_constitutive_models():
    from moju.monitor.closure_registry import list_models
    return list_models()


def list_scaling_closure_ids():
    from moju.monitor.closure_registry import list_groups
    return list_groups()
