# CelerisLab/tests/run_sah04_st_matrix.py
"""Sah04 St validation matrix from tests/Sah04_St_validation_matrix.md.

Runs the nine paper-anchored (Re, beta tier) cases with fixed D=30 channel
geometry, optional SRT/TRT/MRT sweep, and evaluates St against targets.

Each case reports **two** Strouhal estimates from the same lift window:

- **raw**: dominant rFFT peak in the paper-frequency band (no target prior).
- **guided**: same band with a Gaussian weight toward ``f0`` from the paper
  ``St_target`` (reduces harmonic confusion in narrow channels).

The matrix **5% / 10% hard-case rules** apply to **guided** ``St`` by default;
``St_raw`` is for diagnosing whether the guide helps or hides a spectrum issue.

Usage::
    conda run -n pycuda_3_10 python tests/run_sah04_st_matrix.py --collision MRT
    conda run -n pycuda_3_10 python tests/run_sah04_st_matrix.py --collision all --json-out sah04_matrix.json
    conda run -n pycuda_3_10 python tests/run_sah04_st_matrix.py --smoke --case 3 --collision SRT

Long diagnostic (lift + spectrum + **final-step vorticity** under ``tests/output/``)::

    conda run -n pycuda_3_10 python tests/run_sah04_st_matrix.py --collision MRT --case all \\
      --steps 200000 --burn 80000 \\
      --dump-npz-dir tests/output/sah04_long/npz \\
      --final-vorticity-dir tests/output/sah04_long/vorticity \\
      --json-out tests/output/sah04_long/matrix_mrt.json

Requires **matplotlib** for ``--final-vorticity-dir`` (not a core package dependency).

Design::
    Hard cases use a 5% relative St gate on **guided** St; no hard case worse
    than 10%. Soft cases (2, 8) only check ordering vs neighbors in printed summary.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import tempfile
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np
import pycuda.driver as cuda

_PKG_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
_DEFAULT_LBM = os.path.join(_PKG_ROOT, "src", "CelerisLab", "configs", "config_lbm.json")

# D=30 fixed; Lx_fluid = 80D per Sah04 confined setup
_D = 30
_NX = 80 * _D + 2
_CX = 40.0 * _D + 0.5
_R_CYL = 0.5 * _D

# Blockage tiers: fluid height H = ny - 2; cylinder center (1200.5, center_y)
_TIERS: Dict[str, Dict[str, float]] = {
    "low": {"ny": 62.0, "center_y": 30.5, "beta_nom": 0.5},
    "mid": {"ny": 40.0, "center_y": 19.5, "beta_nom": 0.8},
    "high": {"ny": 35.0, "center_y": 17.0, "beta_nom": 0.9},
}


@dataclass(frozen=True)
class MatrixCase:
    case_id: int
    tier: str
    re: float
    target_st: float
    hard: bool
    steps: int
    burn: int


# Table from Sah04_St_validation_matrix.md
MATRIX: Tuple[MatrixCase, ...] = (
    MatrixCase(1, "low", 124.09, 0.3393, True, 80_000, 30_000),
    MatrixCase(2, "low", 160.0, 0.3450, False, 60_000, 20_000),
    MatrixCase(3, "low", 200.0, 0.3513, True, 60_000, 20_000),
    MatrixCase(4, "mid", 110.24, 0.5363, True, 80_000, 30_000),
    MatrixCase(5, "mid", 160.0, 0.5537, True, 60_000, 20_000),
    MatrixCase(6, "mid", 200.0, 0.5510, True, 60_000, 20_000),
    MatrixCase(7, "high", 162.82, 0.5202, True, 80_000, 30_000),
    MatrixCase(8, "high", 180.0, 0.5254, False, 60_000, 20_000),
    MatrixCase(9, "high", 200.0, 0.5314, True, 60_000, 20_000),
)


def _load_json(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _write_json(path: str, d: dict) -> None:
    with open(path, "w", encoding="utf-8") as f:
        json.dump(d, f, indent=2)


def vorticity_z_from_velocity(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    """Z-component vorticity ``ωz = ∂uy/∂x − ∂ux/∂y`` on a 2D ``(ny, nx)`` slice.

    Lattice spacing is taken as 1 in ``np.gradient`` (LBM cell units).
    """
    ux = np.asarray(ux, dtype=np.float64)
    uy = np.asarray(uy, dtype=np.float64)
    duy_dx = np.gradient(uy, axis=1)
    dux_dy = np.gradient(ux, axis=0)
    return duy_dx - dux_dy


def save_final_vorticity_png(
    path: str,
    ux: np.ndarray,
    uy: np.ndarray,
    *,
    title: str,
) -> None:
    """Write ``ωz`` heatmap from last-step ``ux``/``uy`` to ``path`` (PNG)."""
    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError as e:
        raise RuntimeError(
            "save_final_vorticity_png requires matplotlib (e.g. pip install matplotlib)."
        ) from e

    omega = vorticity_z_from_velocity(ux, uy)
    abs_o = np.abs(omega[np.isfinite(omega)])
    if abs_o.size:
        vmax = float(np.percentile(abs_o, 99.5))
        if vmax <= 0.0:
            vmax = float(np.max(abs_o)) or 1.0
    else:
        vmax = 1.0

    ny, nx = omega.shape
    fw = min(18.0, max(8.0, nx / 100.0))
    fh = min(12.0, max(3.0, ny / 40.0))
    fig, ax = plt.subplots(figsize=(fw, fh))
    im = ax.imshow(
        omega,
        origin="lower",
        aspect="equal",
        cmap="RdBu_r",
        vmin=-vmax,
        vmax=vmax,
        extent=(0, nx - 1, 0, ny - 1),
    )
    ax.set_xlabel("x (lattice)")
    ax.set_ylabel("y (lattice)")
    ax.set_title(title)
    fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label="omega_z")
    fig.tight_layout()
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def rfft_power_spectrum(
    lift: np.ndarray,
    *,
    sample_dt: float,
) -> Tuple[np.ndarray, np.ndarray]:
    """Mean-subtracted Hanning-windowed lift → ``(freqs_hz, power)`` for positive rFFT bins."""
    x = np.asarray(lift, dtype=np.float64)
    x = x - np.mean(x)
    n = x.size
    if n < 64:
        return np.zeros(0, dtype=np.float64), np.zeros(0, dtype=np.float64)
    win = np.hanning(n)
    xw = x * win
    spec = np.abs(np.fft.rfft(xw)) ** 2
    freqs = np.fft.rfftfreq(n, d=float(sample_dt))
    return freqs.astype(np.float64), spec.astype(np.float64)


def _parabolic_peak_freq(freqs: np.ndarray, spec: np.ndarray, idx: int) -> float:
    """Refine discrete FFT peak index ``idx`` with log-power parabolic fit."""
    idx = int(np.clip(idx, 0, spec.size - 1))
    if idx <= 0 or idx + 1 >= spec.size:
        return float(freqs[idx])
    i0, i1, i2 = idx - 1, idx, idx + 1
    y0, y1, y2 = (
        np.log(spec[i0] + 1e-30),
        np.log(spec[i1] + 1e-30),
        np.log(spec[i2] + 1e-30),
    )
    denom = y0 - 2.0 * y1 + y2
    if abs(denom) < 1e-20:
        return float(freqs[i1])
    delta = 0.5 * (y0 - y2) / denom
    delta = float(np.clip(delta, -1.0, 1.0))
    df = float(freqs[i2] - freqs[i1])
    return float(freqs[i1]) + delta * df


def dual_strouhal_from_lift(
    lift: np.ndarray,
    *,
    diameter: float,
    u_max: float,
    sample_dt: float,
    f_hz_min: float,
    f_hz_max: float,
) -> Dict[str, float]:
    """Return raw and guided ``St`` and underlying ``f_peak`` (cycles per LBM step)."""
    nan = {"St_raw": float("nan"), "f_raw": float("nan"), "St_guided": float("nan"), "f_guided": float("nan")}
    freqs, spec = rfft_power_spectrum(lift, sample_dt=sample_dt)
    if freqs.size == 0:
        return nan
    mask = (freqs >= float(f_hz_min)) & (freqs <= float(f_hz_max))
    if not np.any(mask):
        return nan
    m = mask.astype(float)
    idx_raw = int(np.argmax(spec * m))
    f0 = 0.5 * (float(f_hz_min) + float(f_hz_max))
    sigma = max(1e-12, 0.18 * f0)
    weight = np.exp(-((freqs - f0) / sigma) ** 2)
    idx_g = int(np.argmax(spec * m * weight))
    f_raw = _parabolic_peak_freq(freqs, spec, idx_raw)
    f_g = _parabolic_peak_freq(freqs, spec, idx_g)
    return {
        "St_raw": float(f_raw * diameter / u_max),
        "f_raw": float(f_raw),
        "St_guided": float(f_g * diameter / u_max),
        "f_guided": float(f_g),
    }


def dominant_strouhal_from_lift(
    lift: np.ndarray,
    *,
    diameter: float,
    u_max: float,
    sample_dt: float = 1.0,
    f_hz_min: float,
    f_hz_max: float,
) -> Tuple[float, float]:
    """Backward-compatible: returns guided ``(St, f_peak)``."""
    d = dual_strouhal_from_lift(
        lift,
        diameter=diameter,
        u_max=u_max,
        sample_dt=sample_dt,
        f_hz_min=f_hz_min,
        f_hz_max=f_hz_max,
    )
    return d["St_guided"], d["f_guided"]


def shedding_freq_band_hz(
    target_st: float, u_max: float, d_lattice: float, *, half_width: float = 0.42
) -> Tuple[float, float]:
    """Cycles per LBM step around ``f0 = St*Umax/D`` so sparse rFFT bins still get coverage."""
    f0 = float(target_st) * float(u_max) / float(d_lattice)
    lo = max(1e-8, f0 * (1.0 - half_width))
    hi = f0 * (1.0 + half_width)
    return lo, hi


def run_one_simulation(
    *,
    collision: str,
    outlet: str,
    nx: int,
    ny: int,
    center: Tuple[float, float],
    re: float,
    d_lattice: float,
    r_cyl: float,
    u_max: float,
    steps: int,
    burn: int,
    record_every: int,
    f_hz_min: float,
    f_hz_max: float,
    dump_npz_path: Optional[str] = None,
    final_vorticity_png_path: Optional[str] = None,
    flow_figure_title: str = "Sah04 vorticity (final LBM step)",
) -> Dict[str, Any]:
    """Build configs, run steps, return St, rough Cd, curved stats."""
    u0_mean = u_max / 1.5
    nu = u_max * d_lattice / re

    lbm_path = _DEFAULT_LBM
    if not os.path.isfile(lbm_path):
        raise FileNotFoundError(lbm_path)

    cfg = _load_json(lbm_path)
    cfg["grid"]["nx"] = nx
    cfg["grid"]["ny"] = ny
    cfg["grid"]["nz"] = 1
    cfg["physics"]["viscosity"] = float(nu)
    cfg["physics"]["velocity"] = float(u0_mean)
    cfg["physics"]["rho"] = 1.0
    cfg["method"]["collision"] = collision
    cfg["method"]["streaming"] = "double_buffer"
    cfg["method"]["les"]["enabled"] = False
    cfg["method"]["outlet"]["mode"] = outlet

    body_doc = {
        "objects": [
            {
                "type": "cylinder",
                "center": list(center),
                "radius": float(r_cyl),
            }
        ]
    }

    tmpd = tempfile.mkdtemp(prefix="celeris_sah04_mtx_")
    lbm_tmp = os.path.join(tmpd, "config_lbm.json")
    body_tmp = os.path.join(tmpd, "config_body.json")
    _write_json(lbm_tmp, cfg)
    _write_json(body_tmp, body_doc)

    from CelerisLab import Simulation  # noqa: WPS433

    sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp)
    sim.initialize()

    stream = cuda.Stream()
    lift_hist: List[float] = []
    fx_hist: List[float] = []
    step_hist: List[int] = []
    rec_every = max(1, int(record_every))
    rho_every = max(2000, rec_every)
    rho_snap_step: List[int] = []
    rho_snap_min: List[float] = []
    rho_snap_max: List[float] = []

    n_curved = int(sim.field.n_curved)
    fb = int(sim.bodies.fallback_link_count())
    lq = int(sim.bodies.low_q_link_count())

    for step in range(1, int(steps) + 1):
        sim.bodies.zero_force_segment_async(stream)
        sim.stepper.step(
            1,
            action_gpu=sim.bodies.action_gpu,
            obs_gpu=sim.bodies.obs_gpu,
            stream=stream,
        )
        if step % rec_every == 0 or step == int(steps):
            stream.synchronize()
            sim.bodies.download_obs_full_async(stream)
            stream.synchronize()
            fvec = sim.bodies.read_force(0)
            lift_hist.append(float(fvec[1]))
            fx_hist.append(float(fvec[0]))
            step_hist.append(int(step))
            if not np.isfinite(lift_hist[-1]):
                sim.close()
                raise RuntimeError(f"NaN/Inf lift at step {step}")
        if step % rho_every == 0 or step == int(steps):
            stream.synchronize()
            macro = sim.get_macroscopic()
            rho_snap_step.append(step)
            rho_snap_min.append(float(np.min(macro["rho"])))
            rho_snap_max.append(float(np.max(macro["rho"])))

    if final_vorticity_png_path:
        stream.synchronize()
        macro_last = sim.get_macroscopic()
        _vdir = os.path.dirname(os.path.abspath(final_vorticity_png_path))
        if _vdir:
            os.makedirs(_vdir, exist_ok=True)
        save_final_vorticity_png(
            final_vorticity_png_path,
            macro_last["ux"],
            macro_last["uy"],
            title=flow_figure_title,
        )

    sim.close()

    lift_arr = np.array(lift_hist, dtype=np.float64)
    fx_arr = np.array(fx_hist, dtype=np.float64)
    step_arr = np.array(step_hist, dtype=np.int64)
    burn_samp = min(int(burn) // rec_every, max(0, lift_arr.size - 16))
    tail = lift_arr[burn_samp:]
    dual = dual_strouhal_from_lift(
        tail,
        diameter=d_lattice,
        u_max=u_max,
        sample_dt=float(rec_every),
        f_hz_min=f_hz_min,
        f_hz_max=f_hz_max,
    )
    st_guided = dual["St_guided"]
    f_guided = dual["f_guided"]
    mean_cd = float(np.mean(fx_arr[burn_samp:])) * 2.0 / (u_max ** 2 * d_lattice) if fx_arr.size else float("nan")
    if rho_snap_min:
        i0 = next((i for i, s in enumerate(rho_snap_step) if s >= burn), 0)
        rho_rng = (min(rho_snap_min[i0:]), max(rho_snap_max[i0:]))
    else:
        rho_rng = (float("nan"), float("nan"))

    if dump_npz_path:
        freqs_pb, spec_pb = rfft_power_spectrum(tail, sample_dt=float(rec_every))
        f0_hz = 0.5 * (float(f_hz_min) + float(f_hz_max))
        sigma = max(1e-12, 0.18 * f0_hz)
        w_guided = (
            np.exp(-((freqs_pb - f0_hz) / sigma) ** 2) if freqs_pb.size else np.zeros(0)
        )
        band = (
            ((freqs_pb >= float(f_hz_min)) & (freqs_pb <= float(f_hz_max))).astype(np.float64)
            if freqs_pb.size
            else np.zeros(0)
        )
        _dir = os.path.dirname(os.path.abspath(dump_npz_path))
        if _dir:
            os.makedirs(_dir, exist_ok=True)
        np.savez_compressed(
            dump_npz_path,
            lift_samples=lift_arr,
            fx_samples=fx_arr,
            sample_lbm_step=step_arr,
            burn_index_samples=int(burn_samp),
            record_every_lbm_steps=int(rec_every),
            freqs_hz_post_burn=freqs_pb,
            power_post_burn=spec_pb,
            band_mask=band,
            guided_gaussian_weight=w_guided,
            f_hz_min=np.array([float(f_hz_min)], dtype=np.float64),
            f_hz_max=np.array([float(f_hz_max)], dtype=np.float64),
            f0_hz_band_mid=np.array([f0_hz], dtype=np.float64),
            St_raw=np.array([float(dual["St_raw"])], dtype=np.float64),
            St_guided=np.array([float(dual["St_guided"])], dtype=np.float64),
            u_max=np.array([float(u_max)], dtype=np.float64),
            diameter_lattice=np.array([float(d_lattice)], dtype=np.float64),
        )

    return {
        "St": float(st_guided),
        "St_guided": float(st_guided),
        "St_raw": float(dual["St_raw"]),
        "f_peak_per_step": float(f_guided),
        "f_peak_raw_per_step": float(dual["f_raw"]),
        "f_peak_guided_per_step": float(f_guided),
        "mean_Cd": float(mean_cd),
        "n_curved": n_curved,
        "fallback_links": fb,
        "low_q_links": lq,
        "rho_min_post_burn": rho_rng[0],
        "rho_max_post_burn": rho_rng[1],
        "n_lift_samples": int(lift_arr.size),
    }


def relative_st_error(st_meas: float, st_target: float) -> float:
    if not np.isfinite(st_meas) or st_target <= 0:
        return float("inf")
    return abs(st_meas - st_target) / st_target


def evaluate_hard_cases(rows: Sequence[Dict[str, Any]], collision: str) -> Dict[str, Any]:
    hard = [
        r
        for r in rows
        if r.get("hard")
        and r.get("collision") == collision
        and "St" in r
        and np.isfinite(r["St"])
    ]
    errs = [relative_st_error(r["St"], r["target_st"]) for r in hard]
    errs_raw = [
        relative_st_error(r["St_raw"], r["target_st"])
        for r in hard
        if "St_raw" in r and np.isfinite(r["St_raw"])
    ]
    finite = [e for e in errs if np.isfinite(e)]
    within5 = sum(1 for e in finite if e <= 0.05)
    worse10 = sum(1 for e in finite if e > 0.10)
    fr = [e for e in errs_raw if np.isfinite(e)]
    return {
        "collision": collision,
        "hard_count": len(hard),
        "hard_within_5pct": int(within5),
        "hard_worse_than_10pct": int(worse10),
        "pass_primary_rule": bool(within5 >= 5 and worse10 == 0),
        "hard_median_abs_rel_err_raw": float(np.median(fr)) if fr else None,
    }


def main() -> int:
    ap = argparse.ArgumentParser(description="Sah04 St validation matrix (9 cases × collisions)")
    ap.add_argument("--collision", default="MRT", help="SRT, TRT, MRT, or all")
    ap.add_argument("--outlet", default="neq_extrap", choices=("neq_extrap", "zero_gradient", "blended"))
    ap.add_argument("--record-every", type=int, default=5)
    ap.add_argument("--case", default="all", help='Case id 1-9 or "all"')
    ap.add_argument("--smoke", action="store_true", help="Short steps/burn for wiring checks")
    ap.add_argument("--json-out", type=str, default=None, help="Write full result rows to JSON")
    ap.add_argument(
        "--steps",
        type=int,
        default=None,
        help="Override matrix LBM steps for each case (ignored with --smoke)",
    )
    ap.add_argument(
        "--burn",
        type=int,
        default=None,
        help="Override matrix burn in LBM steps (ignored with --smoke)",
    )
    ap.add_argument(
        "--dump-npz-dir",
        type=str,
        default=None,
        help="Directory: write case{id}_{COLL}.npz (+ .meta.json) with lift, fx, sample steps, post-burn spectrum",
    )
    ap.add_argument(
        "--final-vorticity-dir",
        type=str,
        default=None,
        help="Directory: write case{id}_{COLL}_laststep.png (omega_z from final macroscopic slice; needs matplotlib)",
    )
    args = ap.parse_args()

    u_max = 0.1
    collisions: List[str]
    if str(args.collision).lower() == "all":
        collisions = ["SRT", "TRT", "MRT"]
    else:
        collisions = [str(args.collision).upper()]
        if collisions[0] not in ("SRT", "TRT", "MRT"):
            print("--collision must be SRT, TRT, MRT, or all", file=sys.stderr)
            return 2

    case_filter: Optional[int] = None
    if str(args.case).lower() != "all":
        case_filter = int(args.case)
        if case_filter < 1 or case_filter > 9:
            print("--case must be 1-9 or all", file=sys.stderr)
            return 2

    cases_to_run = [c for c in MATRIX if case_filter is None or c.case_id == case_filter]

    if args.dump_npz_dir:
        os.makedirs(args.dump_npz_dir, exist_ok=True)
    if args.final_vorticity_dir:
        os.makedirs(args.final_vorticity_dir, exist_ok=True)

    rows: List[Dict[str, Any]] = []
    for coll in collisions:
        for mc in cases_to_run:
            tier = _TIERS[mc.tier]
            ny = int(tier["ny"])
            center = (_CX, float(tier["center_y"]))
            if args.smoke:
                steps = 6000
                burn = 1500
            else:
                steps = int(args.steps) if args.steps is not None else mc.steps
                burn = int(args.burn) if args.burn is not None else mc.burn
            flo, fhi = shedding_freq_band_hz(mc.target_st, u_max, float(_D))

            dump_path: Optional[str] = None
            if args.dump_npz_dir:
                dump_path = os.path.join(args.dump_npz_dir, f"case{mc.case_id}_{coll}.npz")

            vort_path: Optional[str] = None
            vort_title = ""
            if args.final_vorticity_dir:
                vort_path = os.path.join(
                    args.final_vorticity_dir,
                    f"case{mc.case_id}_{coll}_laststep.png",
                )
                vort_title = (
                    f"Sah04 case {mc.case_id} tier={mc.tier} {coll} Re={mc.re} "
                    f"nx={_NX} ny={ny} last LBM step={steps} (omega_z)"
                )

            print(
                f"--- case {mc.case_id} tier={mc.tier} beta~{tier['beta_nom']} "
                f"Re={mc.re} target_St={mc.target_st} {coll} steps={steps} burn={burn} ---",
                flush=True,
            )
            try:
                out = run_one_simulation(
                    collision=coll,
                    outlet=args.outlet,
                    nx=_NX,
                    ny=ny,
                    center=center,
                    re=float(mc.re),
                    d_lattice=float(_D),
                    r_cyl=_R_CYL,
                    u_max=u_max,
                    steps=steps,
                    burn=burn,
                    record_every=int(args.record_every),
                    f_hz_min=flo,
                    f_hz_max=fhi,
                    dump_npz_path=dump_path,
                    final_vorticity_png_path=vort_path,
                    flow_figure_title=vort_title or "Sah04 vorticity (final LBM step)",
                )
            except Exception as e:
                print(f"FAILED: {e}", file=sys.stderr)
                rows.append(
                    {
                        "case_id": mc.case_id,
                        "tier": mc.tier,
                        "collision": coll,
                        "Re": mc.re,
                        "target_st": mc.target_st,
                        "hard": mc.hard,
                        "error": str(e),
                    }
                )
                continue

            rel_err = relative_st_error(out["St"], mc.target_st)
            rel_err_raw = relative_st_error(out["St_raw"], mc.target_st)
            row = {
                "case_id": mc.case_id,
                "tier": mc.tier,
                "beta_nominal": tier["beta_nom"],
                "collision": coll,
                "Re": mc.re,
                "target_st": mc.target_st,
                "hard": mc.hard,
                "St": out["St"],
                "St_raw": out["St_raw"],
                "St_guided": out["St_guided"],
                "rel_err_st": float(rel_err) if np.isfinite(rel_err) else None,
                "rel_err_st_raw": float(rel_err_raw) if np.isfinite(rel_err_raw) else None,
                "within_5pct": bool(np.isfinite(rel_err) and rel_err <= 0.05),
                "worse_10pct": bool(np.isfinite(rel_err) and rel_err > 0.10),
                "mean_Cd": out["mean_Cd"],
                "n_curved": out["n_curved"],
                "fallback_links": out["fallback_links"],
                "low_q_links": out["low_q_links"],
                "steps": steps,
                "burn": burn,
            }
            rows.append(row)
            if dump_path:
                meta_path = os.path.splitext(dump_path)[0] + ".meta.json"
                _write_json(
                    meta_path,
                    {
                        "case_id": mc.case_id,
                        "tier": mc.tier,
                        "collision": coll,
                        "Re": mc.re,
                        "target_st": mc.target_st,
                        "hard": mc.hard,
                        "steps": steps,
                        "burn": burn,
                        "record_every": int(args.record_every),
                        "f_hz_min": flo,
                        "f_hz_max": fhi,
                        "npz_basename": os.path.basename(dump_path),
                        "St_raw": out["St_raw"],
                        "St_guided": out["St_guided"],
                        "mean_Cd": out["mean_Cd"],
                    },
                )
            flag = "OK" if row.get("within_5pct") else ("SOFT" if not mc.hard else "CHECK")
            print(
                f"  St_raw={out['St_raw']:.5f}  St_guided={out['St_guided']:.5f}  "
                f"rel(guided)={100.0 * rel_err:.2f}%  rel(raw)={100.0 * rel_err_raw:.2f}%  "
                f"Cd~{out['mean_Cd']:.4f}  [{flag}]",
                flush=True,
            )

    if not args.smoke and case_filter is None and str(args.case).lower() == "all":
        print("\n=== Hard-case summary per collision (5% gate; need 5/7 & none >10%) ===")
        for coll in collisions:
            ev = evaluate_hard_cases(rows, coll)
            print(json.dumps(ev, indent=2))

    if args.json_out:
        ev_all = {
            coll: evaluate_hard_cases(rows, coll)
            for coll in {r.get("collision") for r in rows if r.get("collision")}
            if coll
        }
        _write_json(args.json_out, {"rows": rows, "evaluation_by_collision": ev_all})

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
