# CelerisLab/tests/run_inlet_ghost_timing_experiment.py
"""Minimal ghost-inlet timing experiments (post field-diagnostic).

Experiment 1 — DDF time series at inlet center vs first interior column:
  (x=0, y=NY/2) and (x=1, y=NY/2), steps 1..N.
  Populations f1,f2,f5,f6,f7,f8 plus rho, ux (macro and sum f).

Experiment 2 — Same channel with ``inlet.collide=false`` vs ``true``:
  When collide is on, inlet ghost nodes undergo collision after Zou-He BC.
  Compare rho_max / ux_max vs step to test ghost-source timing hypothesis.

Usage::

    conda run -n pycuda_3_10 python tests/run_inlet_ghost_timing_experiment.py
    conda run -n pycuda_3_10 python tests/run_inlet_ghost_timing_experiment.py --exp 1 --steps 50
    conda run -n pycuda_3_10 python tests/run_inlet_ghost_timing_experiment.py --exp 2 --steps 500
"""

from __future__ import annotations

import argparse
import csv
import json
import os
import sys
import tempfile
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")

# D2Q9 indices logged (see zou_he_velocity.cuh).
POP_IDX: Tuple[int, ...] = (1, 2, 5, 6, 7, 8)
POP_NAMES: Tuple[str, ...] = tuple(f"f{i}" for i in POP_IDX)

# cx, cy for macroscopic ux, uy from local f (matches descriptors.cuh D2Q9).
_CX = np.array([0, 1, -1, 0, 0, 1, -1, 1, -1], dtype=np.float64)
_CY = np.array([0, 0, 0, 1, -1, 1, -1, -1, 1], dtype=np.float64)


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 _build_cfg(
    base: dict,
    *,
    nx: int,
    ny: int,
    collision: str,
    inlet_collide: bool,
    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"] = "uniform"
    cfg["method"]["inlet"]["scheme"] = "zou_he_local"
    cfg["method"]["inlet"]["collide"] = bool(inlet_collide)
    cfg["method"]["y_wall_bc"] = "free_slip"
    cfg["method"]["outlet"]["mode"] = "neq_extrap"
    return cfg


def _macro_from_f(f: np.ndarray) -> Tuple[float, float, float]:
    f = np.asarray(f, dtype=np.float64)
    rho = float(np.sum(f))
    if abs(rho) < 1e-14:
        return rho, 0.0, 0.0
    ux = float(np.dot(f, _CX) / rho)
    uy = float(np.dot(f, _CY) / rho)
    return rho, ux, uy


def _sample_node(ddf_qnyx: np.ndarray, x: int, y: int) -> Dict[str, float]:
    f = ddf_qnyx[:, y, x].astype(np.float64)
    rho_m, ux_m, uy_m = _macro_from_f(f)
    out: Dict[str, float] = {
        "rho_sum": rho_m,
        "ux_macro": ux_m,
        "uy_macro": uy_m,
    }
    for i, name in zip(POP_IDX, POP_NAMES):
        out[name] = float(f[i])
    return out


def _run_steps(
    cfg: dict,
    *,
    steps: int,
    y_mid: int,
    probe_x: Sequence[int] = (0, 1),
) -> Tuple[List[Dict[str, Any]], Dict[str, Any]]:
    sys.path.insert(0, os.path.join(_REPO, "src"))
    from CelerisLab import Simulation  # noqa: WPS433

    bdoc = {"objects": []}
    tmpd = tempfile.mkdtemp(prefix="ghost_timing_")
    lbm_tmp = os.path.join(tmpd, "config_lbm.json")
    body_tmp = os.path.join(tmpd, "config_body.json")
    with open(lbm_tmp, "w", encoding="utf-8") as f:
        json.dump(cfg, f)
    with open(body_tmp, "w", encoding="utf-8") as f:
        json.dump(bdoc, f)

    sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp)
    sim.initialize()
    nx, ny = sim.lbm_cfg.nx, sim.lbm_cfg.ny
    y_mid = int(np.clip(y_mid, 1, ny - 2))

    rows: List[Dict[str, Any]] = []
    for step in range(1, int(steps) + 1):
        sim.step(1)
        sim.field.download_ddf()
        farr = sim.field.ddf.reshape(sim.lbm_cfg.nq, ny, nx)

        rec: Dict[str, Any] = {"step": step}
        for x in probe_x:
            tag = "inlet" if x == 0 else "interior"
            s = _sample_node(farr, int(x), y_mid)
            for k, v in s.items():
                rec[f"{tag}_{k}"] = v

        # Pull semantics at interior: f[2] is read from stored f[2] at x=0 (same link index).
        rec["cross_f2_match"] = abs(rec["interior_f2"] - rec["inlet_f2"]) < 1e-5
        rec["cross_f1_inlet_to_int_pull"] = float(
            farr[1, y_mid, 1]
        )  # after step, what x=1 holds in f1
        rec["delta_inlet_f1"] = (
            float("nan") if step == 1 else rec["inlet_f1"] - rows[-1]["inlet_f1"]
        )
        rec["delta_inlet_ux"] = (
            float("nan")
            if step == 1
            else rec["inlet_ux_macro"] - rows[-1]["inlet_ux_macro"]
        )

        macro = sim.get_macroscopic()
        rho_f = np.asarray(macro["rho"], dtype=np.float64)
        ux_f = np.asarray(macro["ux"], dtype=np.float64)
        rec["domain_rho_max"] = float(np.nanmax(rho_f))
        rec["domain_rho_min"] = float(np.nanmin(rho_f))
        rec["domain_ux_max"] = float(np.nanmax(np.abs(ux_f)))
        rec["finite"] = bool(np.isfinite(rho_f).all() and np.isfinite(ux_f).all())

        rows.append(rec)

    meta = {
        "nx": nx,
        "ny": ny,
        "y_mid": y_mid,
        "probe_x": list(probe_x),
        "inlet_collide": bool(cfg["method"]["inlet"].get("collide", False)),
        "collision": cfg["method"]["collision"],
        "inlet_scheme": cfg["method"]["inlet"]["scheme"],
        "y_wall_bc": cfg["method"]["y_wall_bc"],
        "outlet_mode": cfg["method"]["outlet"]["mode"],
        "velocity": cfg["physics"]["velocity"],
        "viscosity": cfg["physics"]["viscosity"],
    }
    sim.close()
    return rows, meta


def _write_csv(path: str, rows: Sequence[Dict[str, Any]]) -> None:
    if not rows:
        return
    keys: List[str] = []
    for r in rows:
        for k in r:
            if k not in keys:
                keys.append(k)
    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)
        w.writeheader()
        w.writerows(rows)


def _plot_exp1(out_dir: str, rows: Sequence[Dict[str, Any]], y_mid: int) -> List[str]:
    try:
        import matplotlib

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

    steps = [int(r["step"]) for r in rows]
    paths: List[str] = []

    def _ts(key: str, label: str, ax, **kw):
        ax.plot(steps, [r[key] for r in rows], label=label, **kw)

    # Populations
    fig, axes = plt.subplots(2, 1, figsize=(11, 7), sharex=True)
    for name in POP_NAMES:
        _ts(f"inlet_{name}", f"inlet {name}", axes[0], linewidth=1.0)
    axes[0].set_ylabel("f at x=0")
    axes[0].legend(ncol=3, fontsize=7, loc="best")
    axes[0].grid(True, alpha=0.3)
    for name in POP_NAMES:
        _ts(f"interior_{name}", f"int {name}", axes[1], linewidth=1.0)
    axes[1].set_ylabel("f at x=1")
    axes[1].set_xlabel("step")
    axes[1].legend(ncol=3, fontsize=7, loc="best")
    axes[1].grid(True, alpha=0.3)
    fig.suptitle(f"Exp1 populations (y={y_mid})")
    fig.tight_layout()
    p1 = os.path.join(out_dir, "exp1_populations.png")
    fig.savefig(p1, dpi=150, bbox_inches="tight")
    plt.close(fig)
    paths.append(p1)

    # rho / ux + step-to-step deltas
    fig, axes = plt.subplots(3, 1, figsize=(11, 8), sharex=True)
    _ts("inlet_ux_macro", "inlet ux", axes[0])
    _ts("interior_ux_macro", "interior ux", axes[0])
    axes[0].axhline(0.03, color="k", ls="--", lw=0.8, label="U0")
    axes[0].set_ylabel("u_x")
    axes[0].legend(fontsize=8)
    axes[0].grid(True, alpha=0.3)
    _ts("inlet_rho_sum", "inlet rho", axes[1])
    _ts("interior_rho_sum", "interior rho", axes[1])
    axes[1].set_ylabel("rho")
    axes[1].grid(True, alpha=0.3)
    _ts("delta_inlet_f1", "|Δf1| inlet", axes[2])
    axes[2].set_ylabel("Δf1")
    axes[2].set_xlabel("step")
    axes[2].grid(True, alpha=0.3)
    fig.suptitle(f"Exp1 macro / inlet f1 increment (y={y_mid})")
    fig.tight_layout()
    p2 = os.path.join(out_dir, "exp1_macro_delta.png")
    fig.savefig(p2, dpi=150, bbox_inches="tight")
    plt.close(fig)
    paths.append(p2)

    return paths


def _plot_exp2(out_dir: str, rows_a: Sequence[Dict[str, Any]], rows_b: Sequence[Dict[str, Any]]) -> List[str]:
    try:
        import matplotlib

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

    fig, axes = plt.subplots(2, 1, figsize=(10, 6), sharex=True)
    for rows, lab, c in (
        (rows_a, "ghost (no collide)", "C0"),
        (rows_b, "inlet collide", "C1"),
    ):
        steps = [int(r["step"]) for r in rows]
        axes[0].plot(steps, [r["domain_rho_max"] for r in rows], label=lab, color=c)
        axes[1].plot(steps, [r["domain_ux_max"] for r in rows], label=lab, color=c)
    axes[0].set_ylabel("rho_max")
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    axes[1].set_ylabel("|ux|_max")
    axes[1].set_xlabel("step")
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    fig.suptitle("Exp2: ghost vs inlet collide")
    fig.tight_layout()
    p = os.path.join(out_dir, "exp2_stability_compare.png")
    fig.savefig(p, dpi=150, bbox_inches="tight")
    plt.close(fig)
    return [p]


def _oscillation_summary(rows: Sequence[Dict[str, Any]], *, last_n: int = 20) -> Dict[str, float]:
    """High-frequency proxy: std of inlet f1 and ux over the last *last_n* steps."""
    if len(rows) < 2:
        return {}
    tail = rows[-min(last_n, len(rows)) :]
    f1 = np.array([r["inlet_f1"] for r in tail], dtype=np.float64)
    ux = np.array([r["inlet_ux_macro"] for r in tail], dtype=np.float64)
    d1 = np.array([r["delta_inlet_f1"] for r in tail if np.isfinite(r["delta_inlet_f1"])], dtype=np.float64)
    return {
        "std_inlet_f1_last": float(np.std(f1)),
        "std_inlet_ux_last": float(np.std(ux)),
        "mean_abs_delta_f1_last": float(np.mean(np.abs(d1))) if d1.size else float("nan"),
    }


def run_exp1(
    base: dict,
    *,
    out_dir: str,
    nx: int,
    ny: int,
    steps: int,
    collision: str,
    velocity: float,
    viscosity: float,
) -> None:
    y_mid = ny // 2
    cfg = _build_cfg(
        base,
        nx=nx,
        ny=ny,
        collision=collision,
        inlet_collide=False,
        velocity=velocity,
        viscosity=viscosity,
    )
    print(f"Exp1: zou_he ghost inlet, y_mid={y_mid}, steps={steps}", flush=True)
    rows, meta = _run_steps(cfg, steps=steps, y_mid=y_mid)
    meta["oscillation"] = _oscillation_summary(rows)

    exp_dir = os.path.join(out_dir, "exp1_ddf_timeseries")
    os.makedirs(exp_dir, exist_ok=True)
    _write_csv(os.path.join(exp_dir, "timeseries.csv"), rows)
    _write_json(os.path.join(exp_dir, "meta.json"), meta)
    plots = _plot_exp1(exp_dir, rows, y_mid)
    for p in plots:
        print(f"  plot: {p}", flush=True)

    # Console summary for quick read
    print("  last 5 steps (inlet center):", flush=True)
    for r in rows[-5:]:
        print(
            f"    step {r['step']:3d}  f1={r['inlet_f1']:.6f}  ux={r['inlet_ux_macro']:.6f}  "
            f"Δf1={r['delta_inlet_f1']:.2e}  rho_max={r['domain_rho_max']:.4f}  finite={r['finite']}",
            flush=True,
        )
    print(f"  oscillation: {meta['oscillation']}", flush=True)
    print(f"Wrote: {exp_dir}/timeseries.csv", flush=True)


def run_exp2(
    base: dict,
    *,
    out_dir: str,
    nx: int,
    ny: int,
    steps: int,
    collision: str,
    velocity: float,
    viscosity: float,
) -> None:
    y_mid = ny // 2
    exp_dir = os.path.join(out_dir, "exp2_inlet_collide")
    os.makedirs(exp_dir, exist_ok=True)

    summaries: Dict[str, Any] = {}
    all_rows: Dict[str, List[Dict[str, Any]]] = {}

    for collide, tag in ((False, "ghost_no_collide"), (True, "ghost_with_collide")):
        cfg = _build_cfg(
            base,
            nx=nx,
            ny=ny,
            collision=collision,
            inlet_collide=collide,
            velocity=velocity,
            viscosity=viscosity,
        )
        print(f"Exp2 [{tag}]: inlet.collide={collide}, steps={steps}", flush=True)
        rows, meta = _run_steps(cfg, steps=steps, y_mid=y_mid)
        all_rows[tag] = rows
        _write_csv(os.path.join(exp_dir, f"{tag}.csv"), rows)

        last_finite = next(
            (int(r["step"]) for r in rows if not r.get("finite", True)),
            None,
        )
        summaries[tag] = {
            **meta,
            "first_nonfinite_step": last_finite,
            "final_rho_max": rows[-1]["domain_rho_max"] if rows else None,
            "final_finite": rows[-1].get("finite") if rows else None,
            "oscillation": _oscillation_summary(rows),
        }
        print(
            f"  final rho_max={summaries[tag]['final_rho_max']:.4f}  "
            f"finite={summaries[tag]['final_finite']}  "
            f"first_nonfinite={summaries[tag]['first_nonfinite_step']}",
            flush=True,
        )

    _write_json(os.path.join(exp_dir, "summary.json"), summaries)
    plots = _plot_exp2(exp_dir, all_rows["ghost_no_collide"], all_rows["ghost_with_collide"])
    for p in plots:
        print(f"  plot: {p}", flush=True)
    print(f"Wrote: {exp_dir}/summary.json", flush=True)


def main() -> int:
    ap = argparse.ArgumentParser(description="Ghost inlet timing experiments")
    ap.add_argument("--exp", choices=("1", "2", "all"), default="all")
    ap.add_argument("--steps", type=int, default=50, help="Steps for exp1 (default 50)")
    ap.add_argument("--steps2", type=int, default=500, help="Steps for exp2 (default 500)")
    ap.add_argument("--nx", type=int, default=401)
    ap.add_argument("--ny", type=int, default=201)
    ap.add_argument("--collision", default="MRT", choices=("SRT", "TRT", "MRT"))
    ap.add_argument("--velocity", type=float, default=0.03)
    ap.add_argument("--viscosity", type=float, default=0.009)
    ap.add_argument(
        "--out-dir",
        default=os.path.join(_REPO, "tests", "output", "inlet_ghost_timing"),
    )
    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)
    out_dir = os.path.abspath(args.out_dir)
    os.makedirs(out_dir, exist_ok=True)

    if args.exp in ("1", "all"):
        run_exp1(
            base,
            out_dir=out_dir,
            nx=args.nx,
            ny=args.ny,
            steps=args.steps,
            collision=args.collision,
            velocity=args.velocity,
            viscosity=args.viscosity,
        )
    if args.exp in ("2", "all"):
        run_exp2(
            base,
            out_dir=out_dir,
            nx=args.nx,
            ny=args.ny,
            steps=args.steps2,
            collision=args.collision,
            velocity=args.velocity,
            viscosity=args.viscosity,
        )
    return 0


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