# CelerisLab/tests/run_inlet_channel_diagnostic.py
"""Empty-channel inlet diagnostic: field snapshots and line profiles.

Runs no-cylinder channel flows to isolate inlet / wall / outlet effects before
adding a body.  See user matrix in module docstring sections A–C.

Usage::

    conda run -n pycuda_3_10 python tests/run_inlet_channel_diagnostic.py --part all
    conda run -n pycuda_3_10 python tests/run_inlet_channel_diagnostic.py --part a
    conda run -n pycuda_3_10 python tests/run_inlet_channel_diagnostic.py --part b --nx 401 --ny 201

Output (default ``tests/output/inlet_channel_diag/``)::

    A_baseline/{SRT,MRT}/snapshots/step_XXXXXX.{npz,png}
    B_matrix/caseNN_.../snapshots/...
    B_matrix/caseNN_.../lines/ux_lines_stepXXXXXX.png
    summary.csv
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import sys
import tempfile
from dataclasses import dataclass
from typing import Any, Dict, Iterable, 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")

# Default snapshot steps for parts A and B.
DEFAULT_SNAPSHOT_STEPS: Tuple[int, ...] = (100, 500, 1000, 1500, 2000)


@dataclass(frozen=True)
class CaseSpec:
    """One empty-channel configuration."""

    case_id: str
    label: str
    inlet_scheme: str
    y_wall_bc: str
    outlet_mode: str
    collision: str = "MRT"


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:
    """ωz = ∂uy/∂x − ∂ux/∂y on (ny, nx) arrays."""
    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}"),
    ]


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


def _field_stats(rho: np.ndarray, ux: np.ndarray, vort: np.ndarray) -> Dict[str, float]:
    def _f(a: np.ndarray) -> float:
        b = a[np.isfinite(a)]
        return float("nan") if b.size == 0 else float(np.max(np.abs(b)))

    return {
        "rho_min": float(np.nanmin(rho)) if np.isfinite(rho).any() else float("nan"),
        "rho_max": float(np.nanmax(rho)) if np.isfinite(rho).any() else float("nan"),
        "ux_max": _f(ux),
        "vort_max": _f(vort),
        "finite": bool(np.isfinite(rho).all() and np.isfinite(ux).all()),
    }


def _save_snapshot_npz(
    path: str,
    *,
    step: int,
    rho: np.ndarray,
    ux: np.ndarray,
    uy: np.ndarray,
    vort: np.ndarray,
    meta: dict,
) -> None:
    os.makedirs(os.path.dirname(path), exist_ok=True)
    np.savez_compressed(
        path,
        rho=np.asarray(rho, dtype=np.float32),
        ux=np.asarray(ux, dtype=np.float32),
        uy=np.asarray(uy, dtype=np.float32),
        vort=np.asarray(vort, dtype=np.float32),
        step=np.int32(step),
        meta_json=np.asarray(json.dumps(meta)),
    )


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
        fig, ax = plt.subplots(figsize=(min(16.0, max(8.0, nx / 80.0)), min(8.0, max(3.0, ny / 50.0))))
        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=140, 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,
    case_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 / 60.0)), 7.0), sharex=True)
    for y_idx, y_lab in y_rows:
        y_idx = int(np.clip(y_idx, 0, ny - 1))
        axes[0].plot(x, ux[y_idx, :], label=y_lab, linewidth=1.2)
        axes[1].plot(x, rho[y_idx, :], label=y_lab, linewidth=1.2)

    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"{case_label} — line profiles 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 _run_channel(
    case: CaseSpec,
    *,
    base_cfg: dict,
    nx: int,
    ny: int,
    velocity: float,
    viscosity: float,
    out_root: str,
    snapshot_steps: Sequence[int],
    max_step: int,
    save_png: bool,
    save_lines: bool,
    stop_on_nan: bool,
) -> List[Dict[str, Any]]:
    """Run one case; write snapshots and optional line plots. Return summary rows."""
    sys.path.insert(0, os.path.join(_REPO, "src"))
    from CelerisLab import Simulation  # noqa: WPS433

    cfg = _build_cfg(
        base_cfg,
        nx=nx,
        ny=ny,
        collision=case.collision,
        inlet_scheme=case.inlet_scheme,
        inlet_profile="uniform",
        y_wall_bc=case.y_wall_bc,
        outlet_mode=case.outlet_mode,
        velocity=velocity,
        viscosity=viscosity,
    )
    bdoc = {"objects": []}

    run_dir = os.path.join(out_root, case.case_id)
    snap_dir = os.path.join(run_dir, "snapshots")
    line_dir = os.path.join(run_dir, "lines")
    os.makedirs(snap_dir, exist_ok=True)
    if save_lines:
        os.makedirs(line_dir, exist_ok=True)

    tmpd = tempfile.mkdtemp(prefix="inlet_diag_")
    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, bdoc)

    meta_base = {
        "case_id": case.case_id,
        "label": case.label,
        "nx": nx,
        "ny": ny,
        "inlet_scheme": case.inlet_scheme,
        "inlet_profile": "uniform",
        "y_wall_bc": case.y_wall_bc,
        "outlet_mode": case.outlet_mode,
        "collision": case.collision,
        "velocity": velocity,
        "viscosity": viscosity,
    }
    _write_json(os.path.join(run_dir, "case_meta.json"), meta_base)

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

    y_rows = _line_y_indices(ny)
    want_steps = sorted({int(s) for s in snapshot_steps if 0 < int(s) <= max_step})
    next_snap = 0
    rows: List[Dict[str, Any]] = []

    for step in range(1, max_step + 1):
        sim.step(1)
        if next_snap < len(want_steps) and step == want_steps[next_snap]:
            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)
            stats = _field_stats(rho, ux, vort)
            stem = f"step_{step:06d}"
            meta = {**meta_base, "step": step, **stats}
            npz_path = os.path.join(snap_dir, f"{stem}.npz")
            _save_snapshot_npz(
                npz_path,
                step=step,
                rho=rho,
                ux=ux,
                uy=uy,
                vort=vort,
                meta=meta,
            )
            if save_png:
                _save_field_pngs(
                    snap_dir,
                    stem,
                    rho=rho,
                    ux=ux,
                    vort=vort,
                    title=f"{case.label} step {step}",
                )
            if save_lines:
                line_png = os.path.join(line_dir, f"lines_{stem}.png")
                _save_line_plots(
                    line_png,
                    rho=rho,
                    ux=ux,
                    step=step,
                    case_label=case.label,
                    y_rows=y_rows,
                )
                # Also save raw 1D data for replotting.
                line_npz = os.path.join(line_dir, f"lines_{stem}.npz")
                payload = {"x": np.arange(nx, dtype=np.float32)}
                for y_idx, y_lab in y_rows:
                    payload[f"ux_{y_lab}"] = ux[y_idx, :].astype(np.float32)
                    payload[f"rho_{y_lab}"] = rho[y_idx, :].astype(np.float32)
                payload["step"] = np.int32(step)
                np.savez_compressed(line_npz, **payload)

            rows.append(
                {
                    "case_id": case.case_id,
                    "label": case.label,
                    "collision": case.collision,
                    "inlet_scheme": case.inlet_scheme,
                    "y_wall_bc": case.y_wall_bc,
                    "outlet_mode": case.outlet_mode,
                    "step": step,
                    **stats,
                    "npz": npz_path,
                }
            )
            if stop_on_nan and not stats["finite"]:
                sim.close()
                rows[-1]["stopped_early"] = True
                return rows
            next_snap += 1

    sim.close()
    return rows


def _part_a_cases() -> List[CaseSpec]:
    # Kan99b-style baseline: zou_he + free_slip + neq_extrap; SRT and MRT.
    base = CaseSpec(
        case_id="",
        label="",
        inlet_scheme="zou_he_local",
        y_wall_bc="free_slip",
        outlet_mode="neq_extrap",
    )
    out: List[CaseSpec] = []
    for coll in ("SRT", "MRT"):
        cid = f"A_{coll.lower()}_zouhe_fs_neq"
        out.append(
            CaseSpec(
                case_id=cid,
                label=f"A baseline {coll} zou_he / free_slip / neq_extrap",
                inlet_scheme=base.inlet_scheme,
                y_wall_bc=base.y_wall_bc,
                outlet_mode=base.outlet_mode,
                collision=coll,
            )
        )
    return out


def _part_b_cases() -> List[CaseSpec]:
    return [
        CaseSpec(
            case_id="B_case01_zouhe_fs_neq",
            label="1 zou_he / free_slip / neq_extrap",
            inlet_scheme="zou_he_local",
            y_wall_bc="free_slip",
            outlet_mode="neq_extrap",
            collision="MRT",
        ),
        CaseSpec(
            case_id="B_case02_zouhe_bb_neq",
            label="2 zou_he / bounce_back / neq_extrap",
            inlet_scheme="zou_he_local",
            y_wall_bc="bounce_back",
            outlet_mode="neq_extrap",
            collision="MRT",
        ),
        CaseSpec(
            case_id="B_case03_zouhe_fs_zgrad",
            label="3 zou_he / free_slip / zero_gradient",
            inlet_scheme="zou_he_local",
            y_wall_bc="free_slip",
            outlet_mode="zero_gradient",
            collision="MRT",
        ),
        CaseSpec(
            case_id="B_case04_stab_fs_neq",
            label="4 channel_stabilized / free_slip / neq_extrap",
            inlet_scheme="channel_stabilized",
            y_wall_bc="free_slip",
            outlet_mode="neq_extrap",
            collision="MRT",
        ),
    ]


def _write_summary_csv(path: str, rows: Sequence[Dict[str, Any]]) -> None:
    if not rows:
        return
    keys = [
        "case_id",
        "label",
        "collision",
        "inlet_scheme",
        "y_wall_bc",
        "outlet_mode",
        "step",
        "finite",
        "rho_min",
        "rho_max",
        "ux_max",
        "vort_max",
        "stopped_early",
    ]
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    with open(path, "w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=keys, extrasaction="ignore")
        w.writeheader()
        for r in rows:
            w.writerow(r)


def main() -> int:
    ap = argparse.ArgumentParser(description="Empty-channel inlet diagnostic (fields + line profiles)")
    ap.add_argument(
        "--part",
        choices=("a", "b", "all"),
        default="all",
        help="A=baseline SRT/MRT; B=4-case matrix; all=both",
    )
    ap.add_argument("--nx", type=int, default=401, help="Channel length (lattice)")
    ap.add_argument("--ny", type=int, default=201, help="Channel height (lattice)")
    ap.add_argument("--velocity", type=float, default=0.03, help="Uniform inlet U0")
    ap.add_argument("--viscosity", type=float, default=0.009, help="Kinematic viscosity")
    ap.add_argument(
        "--out-dir",
        type=str,
        default=os.path.join(_REPO, "tests", "output", "inlet_channel_diag"),
    )
    ap.add_argument(
        "--steps",
        type=str,
        default="",
        help="Comma-separated snapshot steps (default: 100,500,1000,1500,2000)",
    )
    ap.add_argument("--no-png", action="store_true", help="Skip rho/ux/vort PNG maps")
    ap.add_argument("--no-lines", action="store_true", help="Skip ux/rho line-profile plots")
    ap.add_argument(
        "--continue-on-nan",
        action="store_true",
        help="Keep stepping after non-finite fields (default: stop case early)",
    )
    args = ap.parse_args()

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

    if args.steps.strip():
        snap_steps = tuple(int(s.strip()) for s in args.steps.split(",") if s.strip())
    else:
        snap_steps = DEFAULT_SNAPSHOT_STEPS
    max_step = max(snap_steps)

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

    cases: List[CaseSpec] = []
    if args.part in ("a", "all"):
        cases.extend(_part_a_cases())
    if args.part in ("b", "all"):
        cases.extend(_part_b_cases())

    save_png = not args.no_png
    # Part A: field maps only; Part B: fields + line plots.
    all_rows: List[Dict[str, Any]] = []

    for case in cases:
        part = "A_baseline" if case.case_id.startswith("A_") else "B_matrix"
        root = os.path.join(out_dir, part)
        save_lines = not args.no_lines and part == "B_matrix"
        print(f"--- {case.case_id}: {case.label} ({case.collision}) ---", flush=True)
        rows = _run_channel(
            case,
            base_cfg=base_cfg,
            nx=args.nx,
            ny=args.ny,
            velocity=args.velocity,
            viscosity=args.viscosity,
            out_root=root,
            snapshot_steps=snap_steps,
            max_step=max_step,
            save_png=save_png,
            save_lines=save_lines,
            stop_on_nan=not args.continue_on_nan,
        )
        all_rows.extend(rows)
        for r in rows:
            fin = "OK" if r.get("finite") else "NONFINITE"
            print(
                f"  step {r['step']:5d}  {fin}  rho=[{r['rho_min']:.4f},{r['rho_max']:.4f}]  "
                f"ux_max={r['ux_max']:.4f}  vort_max={r['vort_max']:.4f}",
                flush=True,
            )
        if rows and rows[-1].get("stopped_early"):
            print("  (stopped early: non-finite fields)", flush=True)

    summary_path = os.path.join(out_dir, "summary.csv")
    _write_summary_csv(summary_path, all_rows)
    manifest = {
        "snapshot_steps": list(snap_steps),
        "nx": args.nx,
        "ny": args.ny,
        "velocity": args.velocity,
        "viscosity": args.viscosity,
        "line_y_indices": [{"y": y, "label": lab} for y, lab in _line_y_indices(args.ny)],
        "cases": [c.case_id for c in cases],
    }
    _write_json(os.path.join(out_dir, "manifest.json"), manifest)

    print(f"Wrote: {summary_path}", flush=True)
    print(f"Wrote: {os.path.join(out_dir, 'manifest.json')}", flush=True)
    print(f"Output root: {out_dir}", flush=True)
    return 0


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