# tests/run_sah04_case9_flow_diagnostics.py
"""Sah04 matrix case 9: velocity profiles at several x and full-field vorticity.

Loads geometry and setup from ``tests/tests/run_sah04_st_matrix.py``, runs one
simulation to the last LBM step, then writes PNG + JSON under
``tests/output/sah04_case9_flow_diag/`` (default).

Usage::
    conda run -n pycuda_3_10 python tests/run_sah04_case9_flow_diagnostics.py
    conda run -n pycuda_3_10 python tests/run_sah04_case9_flow_diagnostics.py --steps 60000 --burn 20000

Design::
    Intended for inlet / channel diagnostics after boundary fixes; not part of
    the St matrix gate.
"""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
import shutil
import sys
import tempfile
from typing import Any, Dict, List, Sequence

import numpy as np
import pycuda.driver as cuda

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
_mtx_path = os.path.join(_REPO, "tests", "tests", "run_sah04_st_matrix.py")
_spec = importlib.util.spec_from_file_location("sah04_mtx", _mtx_path)
if _spec is None or _spec.loader is None:
    raise RuntimeError(f"Cannot load {_mtx_path}")
sah = importlib.util.module_from_spec(_spec)
sys.modules["sah04_mtx"] = sah
_spec.loader.exec_module(sah)


def _inlet_parabolic_ref_u(y: np.ndarray, ny: int, u0_mean: float) -> np.ndarray:
    """Match ``inlet_target_u`` in ``inlet_outlet.cuh`` (parabolic, mean ``u0_mean``)."""
    h = max(float(ny - 2), 1.0)
    yc = np.clip(y.astype(np.float64), 1.0, float(ny - 2))
    eta = (yc - 0.5) / h
    shape = np.maximum(0.0, 4.0 * eta * (1.0 - eta))
    return u0_mean * 1.5 * shape


def _default_x_indices(nx: int) -> List[int]:
    raw = [1, 2, 3, 5, 10, 20, 40, 80, 160, 320, 640, 960, 1200, 1600, 2000, 2200]
    out: List[int] = []
    for x in raw:
        if 1 <= x < nx - 1:
            out.append(int(x))
    mid = max(1, (nx - 2) // 2)
    if mid not in out:
        out.append(mid)
    out = sorted(set(out))
    return out


def _plot_profiles(
    out_png: str,
    *,
    ux: np.ndarray,
    uy: np.ndarray,
    x_indices: Sequence[int],
    ny: int,
    u0_mean: float,
    title: str,
) -> None:
    import matplotlib

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

    y = np.arange(ny, dtype=np.float64)
    u_ref = _inlet_parabolic_ref_u(y, ny, u0_mean)

    fig, axes = plt.subplots(1, 2, figsize=(12, 6), sharey=True)
    for xi in x_indices:
        axes[0].plot(ux[:, xi], y, label=f"x={xi}", linewidth=1.2)
        axes[1].plot(uy[:, xi], y, label=f"x={xi}", linewidth=1.2)
    axes[0].plot(u_ref, y, "k--", linewidth=1.5, label="inlet parabolic ref")
    axes[0].set_xlabel("u_x")
    axes[0].set_ylabel("y (lattice)")
    axes[0].set_title("u_x (y) at selected x")
    axes[0].legend(loc="best", fontsize=7)
    axes[0].grid(True, alpha=0.3)
    axes[1].set_xlabel("u_y")
    axes[1].set_title("u_y (y) at selected x")
    axes[1].legend(loc="best", fontsize=7)
    axes[1].grid(True, alpha=0.3)
    fig.suptitle(title)
    fig.tight_layout()
    os.makedirs(os.path.dirname(os.path.abspath(out_png)) or ".", exist_ok=True)
    fig.savefig(out_png, dpi=150, bbox_inches="tight")
    plt.close(fig)


def main() -> int:
    ap = argparse.ArgumentParser(description="Sah04 case 9 flow: profiles + vorticity")
    ap.add_argument("--collision", default="MRT", help="SRT, TRT, or MRT")
    ap.add_argument("--steps", type=int, default=60_000)
    ap.add_argument("--burn", type=int, default=20_000)
    ap.add_argument("--record-every", type=int, default=5)
    ap.add_argument("--outlet", default="neq_extrap")
    ap.add_argument(
        "--out-dir",
        default=os.path.join(_REPO, "tests", "output", "sah04_case9_flow_diag"),
        help="Output directory for PNG + JSON",
    )
    ap.add_argument(
        "--x-list",
        type=str,
        default="",
        help='Comma-separated x indices (default: built-in list); e.g. "1,2,5,40,200"',
    )
    args = ap.parse_args()

    mc = next(c for c in sah.MATRIX if c.case_id == 9)
    tier = sah._TIERS[mc.tier]
    ny = int(tier["ny"])
    nx = int(sah._NX)
    center = (sah._CX, float(tier["center_y"]))
    u_max = 0.1
    u0_mean = u_max / 1.5
    nu = u_max * float(sah._D) / float(mc.re)

    if args.x_list.strip():
        x_indices = [int(s.strip()) for s in args.x_list.split(",") if s.strip()]
        x_indices = [x for x in x_indices if 1 <= x < nx - 1]
    else:
        x_indices = _default_x_indices(nx)

    lbm_path = sah._DEFAULT_LBM
    cfg = sah._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"] = str(args.collision).upper()
    cfg["method"]["streaming"] = "double_buffer"
    cfg["method"]["les"]["enabled"] = False
    cfg["method"]["outlet"]["mode"] = args.outlet

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

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

    os.makedirs(args.out_dir, exist_ok=True)
    png_profiles = os.path.join(args.out_dir, "case9_profiles_ux_uy.png")
    png_vort = os.path.join(args.out_dir, "case9_omega_z_full.png")
    json_path = os.path.join(args.out_dir, "case9_profiles.json")

    try:
        from CelerisLab import Simulation  # noqa: WPS433

        sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp)
        sim.initialize()
        stream = cuda.Stream()
        rec_every = max(1, int(args.record_every))
        steps = int(args.steps)

        for step in range(1, 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,
            )

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

        y = np.arange(ny, dtype=np.float64)
        u_ref = _inlet_parabolic_ref_u(y, ny, u0_mean)

        title = (
            f"Sah04 case 9 {args.collision} Re={mc.re} ny={ny} "
            f"steps={steps} (last step)"
        )
        _plot_profiles(png_profiles, ux=ux, uy=uy, x_indices=x_indices, ny=ny, u0_mean=u0_mean, title=title)
        sah.save_final_vorticity_png(png_vort, macro["ux"], macro["uy"], title=title + " omega_z")

        profiles: Dict[str, Any] = {
            "case_id": 9,
            "collision": str(args.collision).upper(),
            "Re": mc.re,
            "nx": nx,
            "ny": ny,
            "steps": steps,
            "burn_requested": int(args.burn),
            "u_max": u_max,
            "u0_mean_parabolic": u0_mean,
            "x_indices": list(x_indices),
            "y_lattice": y.tolist(),
            "inlet_parabolic_ref_ux": u_ref.tolist(),
            "profiles": {},
        }
        for xi in x_indices:
            profiles["profiles"][str(xi)] = {
                "ux": ux[:, xi].tolist(),
                "uy": uy[:, xi].tolist(),
            }
        with open(json_path, "w", encoding="utf-8") as f:
            json.dump(profiles, f, indent=2)

        print("Wrote:", png_profiles, flush=True)
        print("Wrote:", png_vort, flush=True)
        print("Wrote:", json_path, flush=True)
    finally:
        shutil.rmtree(tmpd, ignore_errors=True)

    return 0


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