# CelerisLab/tests/run_inlet_scenario_fields.py
"""Three-scenario inlet field export: final or last-stable step at 5000 LBM steps.

Scenarios:
  - empty_channel: zou_he_local × {SRT,MRT} × {free_slip,bounce_back}
  - empty_channel_inlet_matrix: all inlet schemes × {SRT,MRT}, bounce_back only
  - kan99b:        zou_he_local × {SRT,MRT} × {free_slip,bounce_back}, Re=100, domain M
  - sah04_case9:   channel_stabilized × {SRT,MRT}, high-blockage case 9 geometry

Outputs per run (under ``--out-dir``):
  fields/final_{rho,ux,vort}.png, fields/final.npz
  lines/lines_ux_rho.png, lines/lines.npz
  run_meta.json

Usage::

    conda run -n pycuda_3_10 python tests/run_inlet_scenario_fields.py
    conda run -n pycuda_3_10 python tests/run_inlet_scenario_fields.py --scenario empty_channel
    conda run -n pycuda_3_10 python tests/run_inlet_scenario_fields.py --scenario empty_channel_inlet_matrix
"""

from __future__ import annotations

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

import numpy as np

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

# Kan99b lattice contract
_KAN_U_INF = 0.03
_KAN_D = 30.0
_KAN_R = 15.0
_KAN_RE = 100.0
_KAN_ALPHA = 1.0

# Sah04 case 9 (high tier)
_SAH_D = 30
_SAH_NX = 80 * _SAH_D + 2
_SAH_NY = 35
_SAH_CX = 40.0 * _SAH_D + 0.5
_SAH_CY = 17.0
_SAH_RE = 200.0
_SAH_U_MAX = 0.1


@dataclass(frozen=True)
class RunSpec:
    """One simulation run specification."""

    scenario: str
    run_id: str
    label: str
    nx: int
    ny: int
    collision: str
    inlet_scheme: str
    inlet_profile: str
    y_wall_bc: str
    outlet_mode: str
    velocity: float
    viscosity: float
    steps: int
    has_cylinder: bool
    cylinder_center: Tuple[float, float] = (0.0, 0.0)
    cylinder_radius: float = 0.0
    cylinder_omega: float = 0.0


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


def _write_json(path: str, payload: dict) -> None:
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2)


def vorticity_z(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    ux = np.asarray(ux, dtype=np.float64)
    uy = np.asarray(uy, dtype=np.float64)
    return np.gradient(uy, axis=1) - np.gradient(ux, axis=0)


def _line_y_indices(ny: int) -> List[Tuple[int, str]]:
    return [(1, "y1"), (ny // 2, f"y{ny // 2}"), (ny - 2, f"y{ny - 2}")]


_INLET_SCHEMES = (
    "zou_he_local",
    "channel_stabilized",
    "equilibrium",
    "regularized",
)


def _empty_channel_inlet_matrix_specs() -> List[RunSpec]:
    """All inlet schemes on empty channel, bounce_back, SRT/MRT (5000-step field export)."""
    specs: List[RunSpec] = []
    for coll in ("SRT", "MRT"):
        for scheme in _INLET_SCHEMES:
            run_id = f"{coll.lower()}_{scheme}"
            specs.append(
                RunSpec(
                    scenario="empty_channel_inlet_matrix",
                    run_id=run_id,
                    label=f"empty {scheme} {coll} bounce_back",
                    nx=401,
                    ny=201,
                    collision=coll,
                    inlet_scheme=scheme,
                    inlet_profile="uniform",
                    y_wall_bc="bounce_back",
                    outlet_mode="neq_extrap",
                    velocity=0.03,
                    viscosity=0.009,
                    steps=5000,
                    has_cylinder=False,
                )
            )
    return specs


def _all_specs() -> List[RunSpec]:
    specs: List[RunSpec] = []
    for coll in ("SRT", "MRT"):
        for wall in ("free_slip", "bounce_back"):
            wid = f"{coll.lower()}_{wall}"
            specs.append(
                RunSpec(
                    scenario="empty_channel",
                    run_id=wid,
                    label=f"empty zou_he {coll} {wall}",
                    nx=401,
                    ny=201,
                    collision=coll,
                    inlet_scheme="zou_he_local",
                    inlet_profile="uniform",
                    y_wall_bc=wall,
                    outlet_mode="neq_extrap",
                    velocity=0.03,
                    viscosity=0.009,
                    steps=5000,
                    has_cylinder=False,
                )
            )
    dom_m = (1351, 601, (450.0, 300.0))
    nu_k = _KAN_U_INF * _KAN_D / _KAN_RE
    omega = 2.0 * _KAN_ALPHA * _KAN_U_INF / _KAN_D
    for coll in ("SRT", "MRT"):
        for wall in ("free_slip", "bounce_back"):
            wid = f"{coll.lower()}_{wall}"
            specs.append(
                RunSpec(
                    scenario="kan99b",
                    run_id=wid,
                    label=f"kan99b zou_he {coll} {wall}",
                    nx=dom_m[0],
                    ny=dom_m[1],
                    collision=coll,
                    inlet_scheme="zou_he_local",
                    inlet_profile="uniform",
                    y_wall_bc=wall,
                    outlet_mode="neq_extrap",
                    velocity=_KAN_U_INF,
                    viscosity=nu_k,
                    steps=5000,
                    has_cylinder=True,
                    cylinder_center=dom_m[2],
                    cylinder_radius=_KAN_R,
                    cylinder_omega=omega,
                )
            )
    u0_mean = _SAH_U_MAX / 1.5
    nu_s = _SAH_U_MAX * _SAH_D / _SAH_RE
    for coll in ("SRT", "MRT"):
        wid = f"{coll.lower()}_channel_stab"
        specs.append(
            RunSpec(
                scenario="sah04_case9",
                run_id=wid,
                label=f"sah04 case9 channel_stab {coll}",
                nx=_SAH_NX,
                ny=_SAH_NY,
                collision=coll,
                inlet_scheme="channel_stabilized",
                inlet_profile="parabolic",
                y_wall_bc="bounce_back",
                outlet_mode="neq_extrap",
                velocity=u0_mean,
                viscosity=nu_s,
                steps=5000,
                has_cylinder=True,
                cylinder_center=(_SAH_CX, _SAH_CY),
                cylinder_radius=0.5 * _SAH_D,
                cylinder_omega=0.0,
            )
        )
    return specs


def _build_cfg(base: dict, spec: RunSpec) -> dict:
    cfg = json.loads(json.dumps(base))
    cfg["grid"]["nx"] = spec.nx
    cfg["grid"]["ny"] = spec.ny
    cfg["grid"]["nz"] = 1
    cfg["physics"]["velocity"] = float(spec.velocity)
    cfg["physics"]["viscosity"] = float(spec.viscosity)
    cfg["physics"]["rho"] = 1.0
    cfg["method"]["collision"] = spec.collision.upper()
    cfg["method"]["streaming"] = "double_buffer"
    cfg["method"]["store_precision"] = "FP32"
    cfg["method"]["les"]["enabled"] = False
    cfg["method"]["inlet"]["profile"] = spec.inlet_profile
    cfg["method"]["inlet"]["scheme"] = spec.inlet_scheme
    cfg["method"]["y_wall_bc"] = spec.y_wall_bc
    cfg["method"]["outlet"]["mode"] = spec.outlet_mode
    return cfg


def _body_doc(spec: RunSpec) -> dict:
    if not spec.has_cylinder:
        return {"objects": []}
    return {
        "objects": [
            {
                "type": "cylinder",
                "center": [float(spec.cylinder_center[0]), float(spec.cylinder_center[1])],
                "radius": float(spec.cylinder_radius),
                "omega": float(spec.cylinder_omega),
            }
        ]
    }


def _save_field_pngs(
    out_dir: str,
    prefix: str,
    *,
    rho: np.ndarray,
    ux: np.ndarray,
    vort: np.ndarray,
    title: str,
) -> List[str]:
    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        return []

    os.makedirs(out_dir, exist_ok=True)
    ny, nx = rho.shape
    extent = (0, nx - 1, 0, ny - 1)
    paths: List[str] = []

    def _one(arr: np.ndarray, name: str, cmap: str, sym: bool) -> None:
        a = np.asarray(arr, dtype=np.float64)
        fin = a[np.isfinite(a)]
        if fin.size == 0:
            vmin, vmax = -1.0, 1.0
        elif sym:
            v = float(np.percentile(np.abs(fin), 99.5)) or 1.0
            vmin, vmax = -v, v
        else:
            vmin = float(np.percentile(fin, 0.5))
            vmax = float(np.percentile(fin, 99.5))
            if vmax <= vmin:
                vmax = vmin + 1.0
        fw = min(18.0, max(8.0, nx / 70.0))
        fh = min(10.0, max(3.0, ny / 45.0))
        fig, ax = plt.subplots(figsize=(fw, fh))
        im = ax.imshow(a, origin="lower", aspect="auto", cmap=cmap, vmin=vmin, vmax=vmax, extent=extent)
        ax.set_xlabel("x")
        ax.set_ylabel("y")
        ax.set_title(f"{title} — {name}")
        fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
        fig.tight_layout()
        p = os.path.join(out_dir, f"{prefix}_{name}.png")
        fig.savefig(p, dpi=150, bbox_inches="tight")
        plt.close(fig)
        paths.append(p)

    _one(rho, "rho", "viridis", sym=False)
    _one(ux, "ux", "RdBu_r", sym=True)
    _one(vort, "vort", "RdBu_r", sym=True)
    return paths


def _save_line_plots(
    path: str,
    *,
    rho: np.ndarray,
    ux: np.ndarray,
    step: int,
    label: str,
    y_rows: Sequence[Tuple[int, str]],
) -> None:
    try:
        import matplotlib

        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        return

    ny, nx = rho.shape
    x = np.arange(nx, dtype=np.float64)
    fig, axes = plt.subplots(2, 1, figsize=(min(14.0, max(8.0, nx / 55.0)), 7.0), sharex=True)
    for y_idx, y_lab in y_rows:
        yi = int(np.clip(y_idx, 0, ny - 1))
        axes[0].plot(x, ux[yi, :], label=y_lab, linewidth=1.0)
        axes[1].plot(x, rho[yi, :], label=y_lab, linewidth=1.0)
    axes[0].set_ylabel("u_x")
    axes[0].legend(loc="best", fontsize=8)
    axes[0].grid(True, alpha=0.3)
    axes[1].set_ylabel("rho")
    axes[1].set_xlabel("x (lattice)")
    axes[1].legend(loc="best", fontsize=8)
    axes[1].grid(True, alpha=0.3)
    fig.suptitle(f"{label} — ux/rho lines at step {step}")
    fig.tight_layout()
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def _snapshot_from_sim(sim) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    macro = sim.get_macroscopic()
    rho = np.asarray(macro["rho"], dtype=np.float64)
    ux = np.asarray(macro["ux"], dtype=np.float64)
    uy = np.asarray(macro["uy"], dtype=np.float64)
    vort = vorticity_z(ux, uy)
    return rho, ux, uy, vort


def _is_stable_fields(
    rho: np.ndarray,
    ux: np.ndarray,
    *,
    rho_lo: float = 0.85,
    rho_hi: float = 1.25,
    ux_cap: float = 0.15,
) -> bool:
    """Finite fields within a physically plausible band (reject pre-blow-up states)."""
    if not (np.isfinite(rho).all() and np.isfinite(ux).all()):
        return False
    r0 = float(np.min(rho))
    r1 = float(np.max(rho))
    umax = float(np.max(np.abs(ux)))
    return (rho_lo <= r0) and (r1 <= rho_hi) and (umax <= ux_cap)


def run_one(spec: RunSpec, base_cfg: dict, out_root: str) -> Dict[str, Any]:
    sys.path.insert(0, os.path.join(_REPO, "src"))
    import pycuda.driver as cuda
    from CelerisLab import Simulation  # noqa: WPS433

    run_dir = os.path.join(out_root, spec.scenario, spec.run_id)
    field_dir = os.path.join(run_dir, "fields")
    line_dir = os.path.join(run_dir, "lines")
    os.makedirs(field_dir, exist_ok=True)
    os.makedirs(line_dir, exist_ok=True)

    cfg = _build_cfg(base_cfg, spec)
    tmpd = tempfile.mkdtemp(prefix="inlet_scenario_")
    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(spec))

    sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp)
    if spec.has_cylinder and spec.cylinder_omega != 0.0:
        sim.bodies.get(0).state.omega = np.float32(spec.cylinder_omega)
    sim.initialize()

    stream = cuda.Stream()
    y_rows = _line_y_indices(spec.ny)

    last_good: Optional[Dict[str, Any]] = None
    first_bad_step: Optional[int] = None
    force_bad_step: Optional[int] = None

    print(f"  [{spec.scenario}/{spec.run_id}] {spec.label}  steps={spec.steps}", flush=True)

    for step in range(1, spec.steps + 1):
        if spec.has_cylinder:
            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 % 100 == 0 or step == spec.steps:
                stream.synchronize()
                sim.bodies.download_obs_full_async(stream)
                stream.synchronize()
                fvec = sim.bodies.read_force(0)
                if not (np.isfinite(fvec[0]) and np.isfinite(fvec[1])):
                    if force_bad_step is None:
                        force_bad_step = step
        else:
            sim.step(1)

        rho, ux, uy, vort = _snapshot_from_sim(sim)
        if _is_stable_fields(rho, ux):
            last_good = {
                "step": step,
                "rho": rho.copy(),
                "ux": ux.copy(),
                "uy": uy.copy(),
                "vort": vort.copy(),
            }
        elif first_bad_step is None:
            first_bad_step = step

    sim.close()

    if last_good is None:
        raise RuntimeError(f"No finite snapshot for {spec.run_id}")

    out_step = int(last_good["step"])
    rho = last_good["rho"]
    ux = last_good["ux"]
    uy = last_good["uy"]
    vort = last_good["vort"]
    requested_final = spec.steps
    used_last_stable = out_step < requested_final

    meta = {
        "scenario": spec.scenario,
        "run_id": spec.run_id,
        "label": spec.label,
        "nx": spec.nx,
        "ny": spec.ny,
        "collision": spec.collision,
        "inlet_scheme": spec.inlet_scheme,
        "inlet_profile": spec.inlet_profile,
        "y_wall_bc": spec.y_wall_bc,
        "outlet_mode": spec.outlet_mode,
        "velocity": spec.velocity,
        "viscosity": spec.viscosity,
        "requested_steps": requested_final,
        "output_step": out_step,
        "used_last_stable": used_last_stable,
        "first_nonfinite_step": first_bad_step,
        "first_force_nonfinite_step": force_bad_step,
        "rho_min": float(np.min(rho)),
        "rho_max": float(np.max(rho)),
        "ux_max": float(np.max(np.abs(ux))),
        "vort_max": float(np.max(np.abs(vort[np.isfinite(vort)]))) if np.isfinite(vort).any() else float("nan"),
    }
    _write_json(os.path.join(run_dir, "run_meta.json"), meta)

    stem = f"step_{out_step:06d}"
    np.savez_compressed(
        os.path.join(field_dir, "final.npz"),
        rho=rho.astype(np.float32),
        ux=ux.astype(np.float32),
        uy=uy.astype(np.float32),
        vort=vort.astype(np.float32),
        step=np.int32(out_step),
    )

    title = f"{spec.label} (step {out_step}" + (", last stable" if used_last_stable else ", final") + ")"
    pngs = _save_field_pngs(field_dir, "final", rho=rho, ux=ux, vort=vort, title=title)
    _save_line_plots(
        os.path.join(line_dir, "lines_ux_rho.png"),
        rho=rho,
        ux=ux,
        step=out_step,
        label=spec.label,
        y_rows=y_rows,
    )
    line_payload: Dict[str, Any] = {"x": np.arange(spec.nx, dtype=np.float32), "step": np.int32(out_step)}
    for y_idx, y_lab in y_rows:
        yi = int(np.clip(y_idx, 0, spec.ny - 1))
        line_payload[f"ux_{y_lab}"] = ux[yi, :].astype(np.float32)
        line_payload[f"rho_{y_lab}"] = rho[yi, :].astype(np.float32)
    np.savez_compressed(os.path.join(line_dir, "lines.npz"), **line_payload)

    status = "last_stable" if used_last_stable else "final"
    print(
        f"    -> {status} step {out_step}  rho=[{meta['rho_min']:.4f},{meta['rho_max']:.4f}]  "
        f"ux_max={meta['ux_max']:.4f}  force_bad={force_bad_step}",
        flush=True,
    )
    return {**meta, "field_pngs": pngs, "run_dir": run_dir}


def main() -> int:
    ap = argparse.ArgumentParser(description="Three-scenario inlet field export (5000 steps)")
    ap.add_argument(
        "--scenario",
        choices=(
            "empty_channel",
            "empty_channel_inlet_matrix",
            "kan99b",
            "sah04_case9",
            "all",
        ),
        default="all",
    )
    ap.add_argument(
        "--collision",
        default="",
        help="Optional filter: SRT or MRT only",
    )
    ap.add_argument("--steps", type=int, default=5000)
    ap.add_argument(
        "--out-dir",
        default=os.path.join(_REPO, "tests", "output", "inlet_scenario_fields"),
    )
    args = ap.parse_args()

    if not os.path.isfile(_DEFAULT_LBM):
        print(f"Missing {_DEFAULT_LBM}", file=sys.stderr)
        return 2

    base = _load_json(_DEFAULT_LBM)
    if args.scenario == "empty_channel_inlet_matrix":
        specs = _empty_channel_inlet_matrix_specs()
    elif args.scenario == "all":
        specs = _all_specs() + _empty_channel_inlet_matrix_specs()
    else:
        specs = _all_specs()
        if args.scenario != "all":
            specs = [s for s in specs if s.scenario == args.scenario]
    if args.collision.strip():
        coll = args.collision.strip().upper()
        specs = [s for s in specs if s.collision.upper() == coll]
    specs = [replace(s, steps=int(args.steps)) for s in specs]

    out_dir = os.path.abspath(args.out_dir)
    os.makedirs(out_dir, exist_ok=True)

    rows: List[Dict[str, Any]] = []
    for spec in specs:
        try:
            row = run_one(spec, base, out_dir)
            rows.append(row)
        except Exception as e:  # noqa: BLE001
            print(f"  FAILED {spec.scenario}/{spec.run_id}: {e}", flush=True)
            rows.append(
                {
                    "scenario": spec.scenario,
                    "run_id": spec.run_id,
                    "label": spec.label,
                    "error": str(e),
                }
            )

    summary_path = os.path.join(out_dir, "summary.csv")
    if rows:
        keys: List[str] = []
        for r in rows:
            for k in r:
                if k not in keys and k != "field_pngs":
                    keys.append(k)
        with open(summary_path, "w", encoding="utf-8", newline="") as f:
            w = csv.DictWriter(f, fieldnames=keys, extrasaction="ignore")
            w.writeheader()
            w.writerows(rows)

    _write_json(
        os.path.join(out_dir, "manifest.json"),
        {"steps": args.steps, "scenario_filter": args.scenario, "runs": [s.run_id for s in specs]},
    )
    print(f"Wrote: {summary_path}", flush=True)
    print(f"Output: {out_dir}", flush=True)
    return 0


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