# tests/run_sah04_case9_grid_blockage_compare.py
"""Sah04 matrix case 9: compare inlet-near flow vs grid refinement vs blockage.

Runs three D2Q9 setups anchored to case 9 (high tier, Re=200, ``u_max``=0.1):

1. **baseline** — matrix geometry: ``D=30``, ``nx=80*D+2``, ``ny=35``,
   cylinder ``(40*D+0.5, 17)``, ``r=D/2``.
2. **grid_2x** — double lattice resolution with the same **blockage** ``D/H``
   and the same ``Lx/D=80`` convention: ``D=60``, ``nx=80*D+2``, ``ny=68``,
   ``center_y`` doubled with the channel, ``r=D/2``, ``Re`` still based on ``D``.
3. **radius_half** — same channel height as baseline but **cylinder diameter
   halved** in lattice units: ``D=15``, ``nx=80*D+2``, ``ny=35``, same
   ``center_y``, ``r=D/2``, ``Re`` based on the smaller ``D`` (lower blockage).

Outputs under ``tests/output/sah04_case9_compare/`` by default:

- ``compare_meta.json`` — per-variant grid, ``nu``, estimated ``beta=D/H`` (no
  full arrays; see ``compare_fields.npz``).
- ``compare_fields.npz`` — full ``ux`` / ``uy`` per variant.
- ``ux_vs_y_inlet.png`` — ``u_x`` vs normalized wall-normal ``eta`` at several
  ``x`` indices (all variants overlaid per panel).
- ``omega_z_<variant>.png`` — final-step vorticity ``omega_z = d u_y/dx - d u_x/dy``
  for each variant.

Forward schedule: ``burn`` warm-up steps (may be ``0``), then ``steps`` more
steps; the reported field is the state **after** ``burn + steps`` LBM steps.

Usage::
    conda run -n pycuda_3_10 python tests/run_sah04_case9_grid_blockage_compare.py --smoke
    conda run -n pycuda_3_10 python tests/run_sah04_case9_grid_blockage_compare.py \\
      --steps 20000 --burn 0 --collision MRT

Design::
    Isolates whether an inlet-channel artifact scales with **mesh** (grid_2x)
    or with **blockage** (radius_half) while keeping the Sah04-style confined
    channel layout. Requires **matplotlib** for PNG output.
"""

from __future__ import annotations

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

import numpy as np
import pycuda.driver as cuda

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


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)


@dataclass(frozen=True)
class Variant:
    """One lattice-resolved Sah04-like case (case 9 Re / u_max convention)."""

    key: str
    nx: int
    ny: int
    center: Tuple[float, float]
    d_lattice: float
    r_cyl: float
    re: float
    note: str


def _variants() -> Tuple[Variant, Variant, Variant]:
    d0 = 30
    ny0 = 35
    cy0 = 17.0
    v0 = Variant(
        key="baseline",
        nx=80 * d0 + 2,
        ny=ny0,
        center=(40.0 * d0 + 0.5, cy0),
        d_lattice=float(d0),
        r_cyl=0.5 * d0,
        re=200.0,
        note="Matrix case 9 (high tier): D=30, ny=35.",
    )
    d2 = 60
    ny2 = 2 * ny0 - 2
    v1 = Variant(
        key="grid_2x",
        nx=80 * d2 + 2,
        ny=ny2,
        center=(40.0 * d2 + 0.5, 2.0 * cy0),
        d_lattice=float(d2),
        r_cyl=0.5 * d2,
        re=200.0,
        note="Double D and ny-2; Lx/D=80 unchanged; blockage D/H ~ baseline.",
    )
    dh = 15
    v2 = Variant(
        key="radius_half",
        nx=80 * dh + 2,
        ny=ny0,
        center=(40.0 * dh + 0.5, cy0),
        d_lattice=float(dh),
        r_cyl=0.5 * dh,
        re=200.0,
        note="Half cylinder diameter in lattice units; same ny as baseline.",
    )
    return v0, v1, v2


def _vorticity_z_from_velocity(ux: np.ndarray, uy: np.ndarray) -> np.ndarray:
    """``omega_z = d u_y/dx - d u_x/dy`` on a 2D ``(ny, nx)`` slice (unit lattice spacing)."""
    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_vorticity_png(
    path: str,
    ux: np.ndarray,
    uy: np.ndarray,
    *,
    title: str,
) -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    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()
    os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def _run_variant(
    v: Variant,
    *,
    collision: str,
    outlet: str,
    u_max: float,
    burn: int,
    steps: int,
) -> Dict[str, Any]:
    u0_mean = u_max / 1.5
    nu = u_max * float(v.d_lattice) / float(v.re)

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

    cfg = _load_json(_DEFAULT_LBM)
    cfg["grid"]["nx"] = int(v.nx)
    cfg["grid"]["ny"] = int(v.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(v.center),
                "radius": float(v.r_cyl),
            }
        ]
    }

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

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

    for _step in range(1, total + 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,
        )

    stream.synchronize()
    macro = sim.get_macroscopic()
    ux = np.asarray(macro["ux"], dtype=np.float64).reshape(v.ny, v.nx)
    uy = np.asarray(macro["uy"], dtype=np.float64).reshape(v.ny, v.nx)
    rho = np.asarray(macro["rho"], dtype=np.float64).reshape(v.ny, v.nx)
    sim.close()

    h_fluid = max(int(v.ny) - 2, 1)
    beta_est = float(v.d_lattice) / float(h_fluid)

    return {
        "key": v.key,
        "nx": int(v.nx),
        "ny": int(v.ny),
        "center": [float(v.center[0]), float(v.center[1])],
        "d_lattice": float(v.d_lattice),
        "r_cyl": float(v.r_cyl),
        "re": float(v.re),
        "nu": float(nu),
        "u_max": float(u_max),
        "u0_mean": float(u0_mean),
        "beta_est": beta_est,
        "note": v.note,
        "burn": int(burn),
        "steps": int(steps),
        "total_lbm_steps": int(total),
        "ux_shape": list(ux.shape),
        "rho_min": float(np.min(rho)),
        "rho_max": float(np.max(rho)),
        "ux": ux.tolist(),
        "uy": uy.tolist(),
    }


def _eta_coords(ny: int) -> np.ndarray:
    """Wall-normal fractional coordinate in (0,1) for fluid rows y=1..ny-2."""
    y = np.arange(ny, dtype=np.float64)
    h = max(float(ny - 2), 1.0)
    out = np.zeros_like(y)
    for yi in range(1, max(ny - 1, 1)):
        out[yi] = (float(yi) - 0.5) / h
    return out


def _plot_ux_vs_y_inlet(
    path: str,
    results: List[Dict[str, Any]],
    x_indices: List[int],
) -> None:
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    n_x = len(x_indices)
    fig, axes = plt.subplots(1, n_x, figsize=(4.2 * n_x, 5.0), sharey=True)
    if n_x == 1:
        axes = [axes]
    colors = {"baseline": "C0", "grid_2x": "C1", "radius_half": "C2"}

    for ax, xi in zip(axes, x_indices):
        for r in results:
            nx = int(r["nx"])
            if xi < 1 or xi >= nx - 1:
                continue
            ux = np.asarray(r["ux"], dtype=np.float64).reshape(r["ny"], r["nx"])
            eta = _eta_coords(int(r["ny"]))
            y = np.arange(int(r["ny"]))
            ax.plot(ux[:, xi], eta, label=r["key"], color=colors.get(r["key"], "k"))
        ax.set_title(f"x = {xi}")
        ax.set_xlabel("u_x")
        ax.grid(True, alpha=0.3)
    axes[0].set_ylabel("eta = (y-0.5)/H  (fluid band)")
    fig.suptitle("u_x vs wall-normal coordinate near inlet (several x)")
    axes[0].set_ylim(0.0, 1.0)
    axes[0].legend(loc="best", fontsize=8)
    fig.tight_layout()
    os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True)
    fig.savefig(path, dpi=150, bbox_inches="tight")
    plt.close(fig)


def main() -> int:
    ap = argparse.ArgumentParser(
        description="Case 9 baseline vs finer grid vs half-radius: inlet-near u_x",
    )
    ap.add_argument("--collision", default="MRT", help="SRT, TRT, or MRT")
    ap.add_argument("--outlet", default="neq_extrap")
    ap.add_argument("--u-max", type=float, default=0.1, dest="u_max")
    ap.add_argument(
        "--burn",
        type=int,
        default=0,
        help="Warm-up LBM steps before the --steps production segment (0 is valid).",
    )
    ap.add_argument(
        "--steps",
        type=int,
        default=20_000,
        help="LBM steps after burn-in; final field is after burn+steps.",
    )
    ap.add_argument(
        "--smoke",
        action="store_true",
        help="Short run: burn=0, steps=2500 (overrides --burn and --steps)",
    )
    ap.add_argument(
        "--out-dir",
        default=os.path.join(_REPO, "tests", "output", "sah04_case9_compare"),
    )
    ap.add_argument(
        "--inlet-x",
        type=str,
        default="1,2,3,5,8",
        help="Comma-separated x indices for u_x(y) panels (must exist for all nx)",
    )
    args = ap.parse_args()

    if args.smoke:
        burn = 0
        steps = 2500
    else:
        burn = max(0, int(args.burn))
        steps = max(1, int(args.steps))

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

    v0, v1, v2 = _variants()
    variants = (v0, v1, v2)

    x_list = [int(s.strip()) for s in args.inlet_x.split(",") if s.strip()]
    min_nx = min(v.nx for v in variants)
    x_list = [x for x in x_list if 1 <= x < min_nx - 1]
    if not x_list:
        print("No valid --inlet-x indices common to all variants.", file=sys.stderr)
        return 2

    results: List[Dict[str, Any]] = []
    for v in variants:
        print(f"--- {v.key}: nx={v.nx} ny={v.ny} D={v.d_lattice} r={v.r_cyl} ---", flush=True)
        row = _run_variant(
            v,
            collision=str(args.collision).upper(),
            outlet=str(args.outlet),
            u_max=float(args.u_max),
            burn=burn,
            steps=steps,
        )
        results.append(row)

    meta_path = os.path.join(out_dir, "compare_meta.json")
    slim = []
    for r in results:
        slim.append({k: v for k, v in r.items() if k not in ("ux", "uy")})
    _write_json(
        meta_path,
        {
            "variants_summary": slim,
            "inlet_x_used": x_list,
            "burn": burn,
            "steps": steps,
            "total_lbm_steps": burn + steps,
            "collision": str(args.collision).upper(),
        },
    )

    try:
        _plot_ux_vs_y_inlet(
            os.path.join(out_dir, "ux_vs_y_inlet.png"),
            results,
            x_list,
        )
        for r in results:
            key = str(r["key"])
            ux = np.asarray(r["ux"], dtype=np.float64).reshape(r["ny"], r["nx"])
            uy = np.asarray(r["uy"], dtype=np.float64).reshape(r["ny"], r["nx"])
            _save_vorticity_png(
                os.path.join(out_dir, f"omega_z_{key}.png"),
                ux,
                uy,
                title=f"omega_z final ({key}) nx={r['nx']} ny={r['ny']} "
                f"D={r['d_lattice']} burn={r['burn']} steps={r['steps']}",
            )
    except ImportError:
        print(
            "matplotlib not installed; skipped PNG. Install matplotlib for figures.",
            file=sys.stderr,
        )

    full_npz = os.path.join(out_dir, "compare_fields.npz")
    np.savez_compressed(
        full_npz,
        keys=np.array([r["key"] for r in results]),
        **{f"ux_{r['key']}": np.asarray(r["ux"], dtype=np.float64) for r in results},
        **{f"uy_{r['key']}": np.asarray(r["uy"], dtype=np.float64) for r in results},
        **{
            f"omega_z_{r['key']}": _vorticity_z_from_velocity(
                np.asarray(r["ux"], dtype=np.float64).reshape(r["ny"], r["nx"]),
                np.asarray(r["uy"], dtype=np.float64).reshape(r["ny"], r["nx"]),
            )
            for r in results
        },
    )
    print(f"Wrote: {meta_path}", flush=True)
    print(f"Wrote: {full_npz}", flush=True)
    print(f"Wrote: {os.path.join(out_dir, 'ux_vs_y_inlet.png')}", flush=True)
    for r in results:
        oz_path = os.path.join(out_dir, f"omega_z_{r['key']}.png")
        print(f"Wrote: {oz_path}", flush=True)
    return 0


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