# CelerisLab/tests/screening/run_config_sweep.py
"""
Lightweight Kan99b K2 config sweep for the flume-optimisation plan.

Parameterizes collision, streaming, store_precision, ddf_shifting,
inlet_scheme, and D.  Runs K2 (Re=100, alpha=1.0) with reduced steps
(60k total, 20k burn) and reports St, force metrics, rel_err, and
wall-clock speed.

Usage::

    # Single run (declarative)
    python tests/screening/run_config_sweep.py \\
        --collision MRT --streaming double_buffer \\
        --inlet-scheme regularized --D 20

    # Batch -- all 12 core runs (serially on one GPU)
    python tests/screening/run_config_sweep.py \\
        --batch-all --device-id 0

    # Batch -- assign to specific GPU devices for parallel execution
    python tests/screening/run_config_sweep.py --batch-all --gpu-map MR1=0 MR2=0 ...
"""

from __future__ import annotations

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

import numpy as np
import pycuda.driver as cuda

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
sys.path.insert(0, os.path.join(_REPO, "src"))

# ---- Constants matching Kan99b spec -------------------------------------------
U_INF = 0.03
KAN99B_ANCHOR = {
    "St": 0.1655,
    "mean_cl": -2.4881,
    "mean_cd": 1.1040,
    "amp_cl": 0.3631,
    "amp_cd": 0.0993,
}

# ---- Domain specs indexed by D ------------------------------------------------
# Layout: (nx, ny, center_x, center_y)  roughly 45D x 20D
DOMAINS = {
    60: {"nx": 2701, "ny": 1201, "cx": 900.0, "cy": 600.0},
    30: {"nx": 1351, "ny": 601, "cx": 450.0, "cy": 300.0},
    20: {"nx":  901, "ny": 401, "cx": 300.0, "cy": 200.0},
}


@dataclass(frozen=True)
class SweepRun:
    id: str
    collision: str
    streaming: str
    store_precision: str
    ddf_shifting: bool
    inlet_scheme: str
    D: int
    # Optional override for inlet_profile (default uniform)
    inlet_profile: str = "uniform"


# ---- Core matrix (12 runs) ----------------------------------------------------
CORE_RUNS: List[SweepRun] = [
    # MR1–MR7: MRT variants
    SweepRun("MR1", "MRT", "double_buffer", "FP32",  False, "regularized", 30),
    SweepRun("MR2", "MRT", "double_buffer", "FP32",  False, "regularized", 20),
    SweepRun("MR3", "MRT", "esopull",       "FP32",  False, "regularized", 20),
    SweepRun("MR4", "MRT", "double_buffer", "FP16S", True,  "regularized", 20),
    SweepRun("MR5", "MRT", "double_buffer", "FP16S", False, "regularized", 20),
    SweepRun("MR6", "MRT", "double_buffer", "FP32",  False, "zou_he_local", 20),
    SweepRun("MR7", "MRT", "double_buffer", "FP16S", True,  "regularized", 30),
    # SR1–S3: SRT variants
    SweepRun("SR1", "SRT", "double_buffer", "FP32",  False, "equilibrium", 20),
    SweepRun("SR2", "SRT", "esopull",       "FP32",  False, "equilibrium", 20),
    SweepRun( "S3", "SRT", "double_buffer", "FP16S", True,  "equilibrium", 20),
    # TR1–TR2: TRT variants
    SweepRun("TR1", "TRT", "double_buffer", "FP32",  False, "regularized", 20),
    SweepRun("TR2", "TRT", "esopull",       "FP32",  False, "regularized", 20),
]

PERF_RUNS: List[SweepRun] = [
    SweepRun("P1", "MRT", "double_buffer", "FP32",  False, "regularized", 20),
    SweepRun("P2", "MRT", "esopull",       "FP32",  False, "regularized", 20),
    SweepRun("P3", "MRT", "double_buffer", "FP16S", True,  "regularized", 20),
    SweepRun("P4", "SRT", "double_buffer", "FP32",  False, "equilibrium", 20),
]

# Ensure perfs runs use the flume grid size (3000x300)
PERF_GRID = (3000, 300)

# ---- Diagnostic runs (Part A: FP16S + Part B: EsoPull) -------------------------
DIAG_RUNS: List[SweepRun] = [
    # A1-A5: FP16S diagnosis
    SweepRun("A1", "MRT", "double_buffer", "FP16S", False, "regularized", 30),
    SweepRun("A2", "MRT", "double_buffer", "FP16S", True,  "zou_he_local", 30),
    SweepRun("A3", "MRT", "double_buffer", "FP16S", False, "zou_he_local", 30),
    SweepRun("A4", "MRT", "double_buffer", "FP16S", True,  "regularized", 60),
    SweepRun("A5", "MRT", "double_buffer", "FP16S", True,  "zou_he_local", 60),
    # B1-B4: EsoPull diagnosis
    SweepRun("B1", "MRT", "esopull",       "FP32",  False, "regularized", 30),
    SweepRun("B2", "MRT", "esopull",       "FP32",  False, "regularized", 60),
    SweepRun("B3", "TRT", "esopull",       "FP32",  False, "regularized", 30),
    SweepRun("B4", "MRT", "esopull",       "FP32",  False, "channel_stabilized", 30),
]


# ---- Helpers -------------------------------------------------------------------
def _nu_from_re(re: float, D: float) -> float:
    return U_INF * D / float(re)


def _omega_body(alpha: float, D: float) -> float:
    return 2.0 * float(alpha) * U_INF / D


def _make_config(run: SweepRun, total_steps: int, burn_in: int) -> Dict[str, Any]:
    """Build a full config dict from a SweepRun spec.

    Returns (lbm_config, body_config) as dicts.
    """
    dom = DOMAINS[run.D]
    nu = _nu_from_re(100.0, float(run.D))  # Re=100 for K2
    ob = _omega_body(1.0, float(run.D))    # alpha=1.0

    lbm = {
        "grid": {
            "lattice_model": "D2Q9",
            "nx": dom["nx"],
            "ny": dom["ny"],
            "nz": 1,
        },
        "physics": {
            "data_type": "FP32",
            "viscosity": nu,
            "velocity": U_INF,
            "rho": 1.0,
        },
        "method": {
            "collision": run.collision,
            "streaming": run.streaming,
            "store_precision": run.store_precision,
            "ddf_shifting": run.ddf_shifting,
            "les": {
                "enabled": False,
                "cs": 0.16,
                "closed_form": True,
            },
            "trt": {
                "magic_param": 0.1875,
            },
            "inlet": {
                "profile": run.inlet_profile,
                "scheme": run.inlet_scheme,
                "trt_neq_damp": 0.5,
                "regularized_neq_damp": 0.5,
            },
            "outlet": {
                "mode": "neq_extrap",
                "backflow_clamp": True,
                "blend_alpha": 0.7,
                "srt_neq_damp": 0.5,
            },
            "y_wall_bc": "free_slip",
            "omega_guard": {
                "min": 0.01,
                "max": 1.96,
            },
        },
        "cuda": {
            "threads_per_block": 256,
            "compute_capability": "auto",
        },
    }
    body = {
        "objects": [
            {
                "type": "cylinder",
                "center": [dom["cx"], dom["cy"]],
                "radius": float(run.D) / 2.0,
                "omega": ob,
            }
        ]
    }
    return lbm, body


def _rfft_spectrum(x: np.ndarray, sample_dt: float) -> Tuple[np.ndarray, np.ndarray]:
    arr = np.asarray(x, dtype=np.float64)
    if arr.size < 64:
        return np.zeros(0, dtype=np.float64), np.zeros(0, dtype=np.float64)
    arr = arr - np.mean(arr)
    spec = np.abs(np.fft.rfft(arr * np.hanning(arr.size))) ** 2
    freqs = np.fft.rfftfreq(arr.size, d=float(sample_dt))
    return freqs.astype(np.float64), spec.astype(np.float64)


def _peak_freq_parabolic(freqs: np.ndarray, spec: np.ndarray, idx: int) -> float:
    i = int(np.clip(idx, 0, spec.size - 1))
    if i <= 0 or i + 1 >= spec.size:
        return float(freqs[i])
    y0 = np.log(spec[i - 1] + 1e-30)
    y1 = np.log(spec[i] + 1e-30)
    y2 = np.log(spec[i + 1] + 1e-30)
    den = y0 - 2.0 * y1 + y2
    if abs(den) < 1e-20:
        return float(freqs[i])
    delta = float(np.clip(0.5 * (y0 - y2) / den, -1.0, 1.0))
    return float(freqs[i]) + delta * float(freqs[i + 1] - freqs[i])


def _st_from_lift(lift: np.ndarray, sample_dt: float, D: float) -> float:
    freqs, spec = _rfft_spectrum(lift, sample_dt=sample_dt)
    if freqs.size <= 1:
        return float("nan")
    idx = int(np.argmax(spec[1:])) + 1
    f_peak = _peak_freq_parabolic(freqs, spec, idx)
    return float(f_peak * D / U_INF)


def _cycle_half_p2p(y: np.ndarray) -> float:
    arr = np.asarray(y, dtype=np.float64)
    if arr.size < 8:
        return float("nan")
    centered = arr - np.mean(arr)
    crossing = np.where((centered[:-1] <= 0.0) & (centered[1:] > 0.0))[0]
    if crossing.size >= 2:
        amps: List[float] = []
        for i in range(crossing.size - 1):
            seg = arr[crossing[i] + 1: crossing[i + 1] + 1]
            if seg.size >= 3:
                amps.append(0.5 * (float(np.max(seg)) - float(np.min(seg))))
        if amps:
            return float(np.mean(amps))
    return 0.5 * (float(np.max(arr)) - float(np.min(arr)))


# ---- Run one sweep configuration -----------------------------------------------
def run_sweep(
    run: SweepRun,
    *,
    total_steps: int = 60000,
    burn_in: int = 20000,
    record_every: int = 100,
    device_id: int = 0,
    perf_timing_steps: int = 0,
    out_dir: str = "",
) -> Dict[str, Any]:
    """Execute one K2 sweep run and return metrics dict."""
    from CelerisLab import Simulation

    use_perf_grid = (perf_timing_steps > 0)
    if use_perf_grid:
        # Build config for the big flume grid for pure timing (no body for simplicity)
        dom = DOMAINS[run.D]
        nu = _nu_from_re(100.0, float(run.D))
        lbm = {
            "grid": {"lattice_model": "D2Q9", "nx": PERF_GRID[0],
                     "ny": PERF_GRID[1], "nz": 1},
            "physics": {"data_type": "FP32", "viscosity": nu,
                        "velocity": U_INF, "rho": 1.0},
            "method": {
                "collision": run.collision,
                "streaming": run.streaming,
                "store_precision": run.store_precision,
                "ddf_shifting": run.ddf_shifting,
                "les": {"enabled": False, "cs": 0.16, "closed_form": True},
                "trt": {"magic_param": 0.1875},
                "inlet": {"profile": "uniform", "scheme": run.inlet_scheme,
                          "trt_neq_damp": 0.5, "regularized_neq_damp": 0.5},
                "outlet": {"mode": "neq_extrap", "backflow_clamp": True,
                           "blend_alpha": 0.7, "srt_neq_damp": 0.5},
                "y_wall_bc": "free_slip",
                "omega_guard": {"min": 0.01, "max": 1.96},
            },
            "cuda": {"threads_per_block": 256, "compute_capability": "auto"},
        }
        body = {"objects": []}
        tmpd = tempfile.mkdtemp(prefix="celeris_sweep_perf_")
        lbm_tmp = os.path.join(tmpd, "config_lbm.json")
        body_tmp = os.path.join(tmpd, "config_body.json")
        with open(lbm_tmp, "w") as f:
            json.dump(lbm, f, indent=2)
        with open(body_tmp, "w") as f:
            json.dump(body, f, indent=2)
        sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp,
                          device_id=device_id)
        sim.initialize()

        stream = cuda.Stream()
        # Warmup
        sim.run(5000, stream=stream)

        # Timed loop
        t0 = time.perf_counter()
        sim.run(perf_timing_steps, stream=stream)
        t1 = time.perf_counter()
        elapsed = t1 - t0
        sim.close()

        n_cells = PERF_GRID[0] * PERF_GRID[1]
        mlups = n_cells * perf_timing_steps / elapsed / 1e6

        return {
            "run_id": f"{run.id}_perf",
            "collision": run.collision,
            "streaming": run.streaming,
            "store_precision": run.store_precision,
            "ddf_shifting": run.ddf_shifting,
            "inlet_scheme": run.inlet_scheme,
            "D": run.D,
            "grid": f"{PERF_GRID[0]}x{PERF_GRID[1]}",
            "perf_timing_steps": perf_timing_steps,
            "wall_clock_s": round(elapsed, 4),
            "mlups": round(mlups, 2),
            "us_per_step": round(elapsed / perf_timing_steps * 1e6, 2),
            "n_cells": n_cells,
        }

    # ---- Normal K2 accuracy run -----------------------------------------------
    lbm_cfg, body_cfg = _make_config(run, total_steps, burn_in)
    tmpd = tempfile.mkdtemp(prefix="celeris_sweep_")
    lbm_tmp = os.path.join(tmpd, "config_lbm.json")
    body_tmp = os.path.join(tmpd, "config_body.json")
    with open(lbm_tmp, "w") as f:
        json.dump(lbm_cfg, f, indent=2)
    with open(body_tmp, "w") as f:
        json.dump(body_cfg, f, indent=2)

    from CelerisLab import Simulation

    sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp,
                      device_id=device_id)
    if sim.bodies.count < 1:
        sim.close()
        raise RuntimeError("Expected one cylinder in body config.")

    # Set rotation and verify
    ob = _omega_body(1.0, float(run.D))
    obj = sim.bodies.get(0)
    ob_f32 = np.float32(ob)
    print(f"  D={run.D} omega_body_set={float(ob_f32):.6f}  "
          f"(pre-init state.omega={float(obj.state.omega):.6f})")
    obj.state.omega = ob_f32
    sim.initialize()
    # Verify action buffer contains omega
    dim = sim.lbm_cfg.dim
    slot = 3 * dim
    action_omega = float(sim.bodies.action[slot - 1])
    print(f"  action_gpu[omega_slot]={action_omega:.6f}  "
          f"(expected {float(ob_f32):.6f}) match={abs(action_omega - float(ob_f32)) < 1e-8}")

    stream = cuda.Stream()
    total = int(burn_in) + int(total_steps)
    if total < 1:
        sim.close()
        raise ValueError("burn + steps must be >= 1")

    step_hist: List[int] = []
    fx_hist: List[float] = []
    fy_hist: List[float] = []

    t0 = time.perf_counter()
    for step in range(1, total + 1):
        sim.bodies.zero_obs_async(stream)
        sim.stepper.step(
            1,
            action_gpu=sim.bodies.action_gpu,
            obs_gpu=sim.bodies.obs_gpu,
            stream=stream,
        )
        if step % record_every == 0 or step == total:
            stream.synchronize()
            sim.bodies.download_obs_full_async(stream)
            stream.synchronize()
            force = sim.bodies.read_force(0, normalize=False)
            fx = float(force[0])
            fy = float(force[1])
            if not np.isfinite(fx) or not np.isfinite(fy):
                sim.close()
                raise RuntimeError(f"NaN/Inf force at step {step}")
            step_hist.append(step)
            fx_hist.append(fx)
            fy_hist.append(fy)
    t1 = time.perf_counter()
    sim.close()

    step_arr = np.asarray(step_hist, dtype=np.int64)
    fx_arr = np.asarray(fx_hist, dtype=np.float64)
    fy_arr = np.asarray(fy_hist, dtype=np.float64)
    burn_mask = step_arr >= int(burn_in)
    if not np.any(burn_mask):
        burn_mask = np.ones_like(step_arr, dtype=bool)

    D_val = float(run.D)
    cl = 2.0 * fy_arr / (U_INF**2 * D_val)
    cd = 2.0 * fx_arr / (U_INF**2 * D_val)
    cl_tail = cl[burn_mask]
    cd_tail = cd[burn_mask]
    st = _st_from_lift(cl_tail, sample_dt=float(record_every), D=D_val)
    amp_cl = _cycle_half_p2p(cl_tail)
    amp_cd = _cycle_half_p2p(cd_tail)
    mean_cl = float(np.mean(cl_tail))
    mean_cd = float(np.mean(cd_tail))
    wall_s = t1 - t0

    n_cells = DOMAINS[run.D]["nx"] * DOMAINS[run.D]["ny"]
    mlups = n_cells * total / wall_s / 1e6

    # Relative errors vs Kan99b anchor
    def _relerr(meas: float, ref: float) -> Optional[float]:
        if not np.isfinite(meas) or ref == 0.0:
            return None
        return abs(float(meas) - float(ref)) / abs(float(ref))

    metrics = {
        "run_id": run.id,
        "collision": run.collision,
        "streaming": run.streaming,
        "store_precision": run.store_precision,
        "ddf_shifting": run.ddf_shifting,
        "inlet_scheme": run.inlet_scheme,
        "inlet_profile": run.inlet_profile,
        "D": int(run.D),
        "grid": f"{DOMAINS[run.D]['nx']}x{DOMAINS[run.D]['ny']}",
        "total_steps": int(total),
        "burn_in": int(burn_in),
        "record_every": int(record_every),
        "n_samples": int(step_arr.size),
        "n_stat_samples": int(np.sum(burn_mask)),
        "wall_clock_s": round(wall_s, 4),
        "mlups": round(mlups, 2),
        "St": float(st),
        "mean_cl": float(mean_cl),
        "mean_cd": float(mean_cd),
        "amp_cl": float(amp_cl),
        "amp_cd": float(amp_cd),
        "err_St": round(_relerr(st, KAN99B_ANCHOR["St"]) * 100, 2) if _relerr(st, KAN99B_ANCHOR["St"]) is not None else None,
        "err_mean_cl": round(_relerr(mean_cl, KAN99B_ANCHOR["mean_cl"]) * 100, 2) if _relerr(mean_cl, KAN99B_ANCHOR["mean_cl"]) is not None else None,
        "err_mean_cd": round(_relerr(mean_cd, KAN99B_ANCHOR["mean_cd"]) * 100, 2) if _relerr(mean_cd, KAN99B_ANCHOR["mean_cd"]) is not None else None,
        "err_amp_cl": round(_relerr(amp_cl, KAN99B_ANCHOR["amp_cl"]) * 100, 2) if _relerr(amp_cl, KAN99B_ANCHOR["amp_cl"]) is not None else None,
        "err_amp_cd": round(_relerr(amp_cd, KAN99B_ANCHOR["amp_cd"]) * 100, 2) if _relerr(amp_cd, KAN99B_ANCHOR["amp_cd"]) is not None else None,
    }

    # Save per-run JSON and CSV to isolated output directory (if out_dir set)
    if out_dir:
        run_out_dir = os.path.join(out_dir, run.id)
        os.makedirs(run_out_dir, exist_ok=True)
        json_path = os.path.join(run_out_dir, "summary.json")
        with open(json_path, "w") as f:
            json.dump(metrics, f, indent=2)
        csv_path = os.path.join(run_out_dir, "force_hist.csv")
        with open(csv_path, "w", newline="") as f_c:
            w_csv = csv.writer(f_c)
            w_csv.writerow(["step", "fx", "fy", "cd", "cl"])
            for i in range(len(step_hist)):
                w_csv.writerow([step_hist[i], fx_hist[i], fy_hist[i],
                                cd[i], cl[i]])
    return metrics


def _format_err(val: Optional[float], band5: float, band10: float) -> str:
    """Format error with colour indicator: pass / flag / fail."""
    if val is None:
        return "  N/A   "
    if val <= band5:
        return f" {val:6.2f}% "  # pass (no ANSI in terminal)
    if val <= band10:
        return f"*{val:6.2f}%*"  # flag
    return f"!{val:6.2f}%!"  # fail


def print_summary(rows: List[Dict[str, Any]]) -> None:
    """Pretty-print the sweep results."""
    cl_lbl = "C'L"
    cd_lbl = "C'D"
    print()
    print("=" * 120)
    print(f"{'Run':>6} {'Coll':>5} {'Stream':>12} {'Store/DDF':>14} {'Inlet':>14} "
          f"{'D':>3} {'Grid':>11} {'St':>8} {'mCL':>8} {'mCD':>8} "
          f"{cl_lbl:>7} {cd_lbl:>7} {'Wall(s)':>8} {'MLUPS':>7}")
    print("-" * 120)
    for r in rows:
        if "error" in r and "grid" not in r:
            print(f"{r['run_id']:>6} {r.get('collision','?'):>5} "
                  f"{r.get('streaming','?'):>12} "
                  f"{r.get('store_precision','?'):>7}/{'S' if r.get('ddf_shifting',False) else 'N':>1}  "
                  f"{r.get('inlet_scheme','?'):>14} "
                  f"{r.get('D','?'):>3} {'?':>11}  --- FAILED: {r.get('error','?')[:60]}")
            continue
        if "perf" in r.get("run_id", ""):
            # Perf row
            print(f"{r['run_id']:>6} {r['collision']:>5} {r['streaming']:>12} "
                  f"{r['store_precision']:>7}/{'S' if r['ddf_shifting'] else 'N':>1}  "
                  f"{r['inlet_scheme']:>14} "
                  f"{r['D']:>3} {r['grid']:>11} {'':>8} {'':>8} {'':>8} "
                  f"{'':>7} {'':>7} "
                  f"{r['wall_clock_s']:>8.4f} {r['mlups']:>7.2f}")
        else:
            e_St = r.get("err_St")
            e_mcl = r.get("err_mean_cl")
            e_mcd = r.get("err_mean_cd")
            e_acl = r.get("err_amp_cl")
            e_acd = r.get("err_amp_cd")
            print(f"{r['run_id']:>6} {r['collision']:>5} {r['streaming']:>12} "
                  f"{r['store_precision']:>7}/{'S' if r['ddf_shifting'] else 'N':>1}  "
                  f"{r['inlet_scheme']:>14} "
                  f"{r['D']:>3} {r['grid']:>11} {r['St']:>8.5f} {r['mean_cl']:>8.4f} {r['mean_cd']:>8.4f} "
                  f"{r['amp_cl']:>7.4f} {r['amp_cd']:>7.4f} "
                  f"{r['wall_clock_s']:>8.4f} {r['mlups']:>7.2f}")
            print(f"{'':>6} {'':>5} {'':>12} {'':>14} {'':>14} "
                  f"{'':>3} {'':>11} "
                  f"{_format_err(e_St, 3, 5)} "
                  f"{_format_err(e_mcl, 4, 8)} "
                  f"{_format_err(e_mcd, 5, 10)} "
                  f"{_format_err(e_acl, 8, 12)} "
                  f"{_format_err(e_acd, 10, 15)} "
                  f"{'':>8} {'':>7}")
    print("=" * 120)
    print("Format: plain=pass, *flag* = outside preferred band, !fail! = outside acceptable band")
    print()


def main() -> int:
    ap = argparse.ArgumentParser(description="Kan99b K2 config sweep")
    ap.add_argument("--run-id", type=str, default="",
                    help="Run a single sweep by id (e.g. MR1, SR1).")
    ap.add_argument("--batch-all", action="store_true",
                    help="Run all 12 core sweeps serially.")
    ap.add_argument("--run-perf", action="store_true",
                    help="Run the 4 perf-timing sweeps on the 3000x300 grid.")
    ap.add_argument("--run-diag", action="store_true",
                    help="Run all diagnostic sweeps (A1-A5 FP16S + B1-B4 EsoPull).")
    ap.add_argument("--collision", type=str, default="MRT",
                    choices=("SRT", "TRT", "MRT"))
    ap.add_argument("--streaming", type=str, default="double_buffer",
                    choices=("double_buffer", "esopull"))
    ap.add_argument("--store-precision", type=str, default="FP32",
                    choices=("FP32", "FP16S"))
    ap.add_argument("--ddf-shifting", action="store_true")
    ap.add_argument("--inlet-scheme", type=str, default="regularized",
                    choices=("zou_he_local", "channel_stabilized",
                             "equilibrium", "regularized"))
    ap.add_argument("--D", type=int, default=20, choices=(20, 30, 60))
    ap.add_argument("--steps", type=int, default=60000)
    ap.add_argument("--burn", type=int, default=20000)
    ap.add_argument("--record-every", type=int, default=100)
    ap.add_argument("--device-id", type=int, default=0)
    ap.add_argument("--out-dir", type=str, default="",
                    help="Output dir for CSV + summary JSON.  Default: tests/output/screening/")
    ap.add_argument("--perf-steps", type=int, default=10000,
                    help="Timing steps for perf runs (after 5000 warmup).")
    args = ap.parse_args()

    out_dir = args.out_dir
    if not out_dir:
        out_dir = os.path.join(_REPO, "tests", "output", "screening")
    os.makedirs(out_dir, exist_ok=True)

    runs_to_do: List[SweepRun] = []
    is_perf = False

    if args.run_id:
        needle = args.run_id.upper()
        for r in CORE_RUNS:
            if r.id == needle:
                runs_to_do = [r]
                break
        if not runs_to_do:
            for r in PERF_RUNS:
                if r.id.upper() == needle:
                    runs_to_do = [r]
                    is_perf = True
                    break
        if not runs_to_do:
            for r in DIAG_RUNS:
                if r.id.upper() == needle:
                    runs_to_do = [r]
                    break
        if not runs_to_do:
            print(f"Unknown run id: {needle}")
            return 1
    elif args.batch_all:
        runs_to_do = list(CORE_RUNS)
    elif args.run_perf:
        runs_to_do = list(PERF_RUNS)
        is_perf = True
    elif args.run_diag:
        runs_to_do = list(DIAG_RUNS)
    else:
        # Single custom run from CLI args
        runs_to_do = [
            SweepRun("custom", args.collision, args.streaming,
                      args.store_precision, args.ddf_shifting,
                      args.inlet_scheme, args.D)
        ]

    rows: List[Dict[str, Any]] = []
    for run in runs_to_do:
        print(f"\n--- {run.id}: {run.collision} {run.streaming} "
              f"{run.store_precision}/{'S' if run.ddf_shifting else 'N'} "
              f"{run.inlet_scheme} D={run.D} ---")
        try:
            if is_perf:
                row = run_sweep(run, device_id=args.device_id,
                                perf_timing_steps=args.perf_steps,
                                out_dir=out_dir)
                print(f"  perf: {row['mlups']} MLUPS, "
                      f"{row['us_per_step']} us/step")
            else:
                row = run_sweep(run, total_steps=args.steps, burn_in=args.burn,
                                record_every=args.record_every,
                                device_id=args.device_id,
                                out_dir=out_dir)
                print(f"  St={row['St']:.5f} mean_CL={row['mean_cl']:.4f} "
                      f"mean_CD={row['mean_cd']:.4f} "
                      f"C'L={row['amp_cl']:.4f} C'D={row['amp_cd']:.4f}")
            rows.append(row)
        except Exception as exc:
            print(f"FAILED: {exc}")
            rows.append({
                "run_id": run.id,
                "collision": run.collision,
                "streaming": run.streaming,
                "store_precision": run.store_precision,
                "ddf_shifting": run.ddf_shifting,
                "inlet_scheme": run.inlet_scheme,
                "D": run.D,
                "error": str(exc),
            })

    # Summary table
    print_summary(rows)

    # Save summary JSON
    summary = {
        "contract": {
            "U_inf": U_INF,
            "Kan99b_anchor": KAN99B_ANCHOR,
        },
        "runs": rows,
    }
    json_path = os.path.join(out_dir, "screening_summary.json")
    with open(json_path, "w") as f:
        json.dump(summary, f, indent=2)
    print(f"Summary: {json_path}")

    # CSV with key fields
    csv_path = os.path.join(out_dir, "screening_summary.csv")
    csv_keys = [
        "run_id", "collision", "streaming", "store_precision", "ddf_shifting",
        "inlet_scheme", "inlet_profile", "D", "grid",
        "total_steps", "burn_in", "n_stat_samples",
        "wall_clock_s", "mlups",
        "St", "mean_cl", "mean_cd", "amp_cl", "amp_cd",
        "err_St", "err_mean_cl", "err_mean_cd", "err_amp_cl", "err_amp_cd",
        "error",
    ]
    with open(csv_path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=csv_keys)
        w.writeheader()
        for r in rows:
            w.writerow({k: r.get(k, "") for k in csv_keys})
    print(f"CSV:    {csv_path}")
    return 0


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