"""Performance baseline for minimal host-interference LBM stepping.

This script builds a temporary config, runs warmup + measured batches, and
reports MLUPS under a FluidX3D-like benchmark mindset:

- keep the main loop on GPU (`stepper.step`)
- avoid host downloads by default
- selectively enable host-touch paths to quantify overhead
- sweep `inlet.scheme` to check stability-sensitive combinations

Usage::
    python tests/run_perf_baseline.py
    python tests/run_perf_baseline.py --lattice-model D2Q9 --nx 384 --ny 192 --steps 30000
    python tests/run_perf_baseline.py --macro-every 500 --ddf-every 1000
    python tests/run_perf_baseline.py --with-cylinder --obs-every 50
"""

from __future__ import annotations

import argparse
import json
import os
import tempfile
import time
from typing import Any, Dict, List

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, payload: dict) -> None:
    with open(path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2)


def _build_lbm_cfg(base: dict, args: argparse.Namespace) -> dict:
    cfg = json.loads(json.dumps(base))
    cfg["grid"]["lattice_model"] = args.lattice_model
    cfg["grid"]["nx"] = int(args.nx)
    cfg["grid"]["ny"] = int(args.ny)
    cfg["grid"]["nz"] = int(args.nz)

    cfg["physics"]["viscosity"] = float(args.viscosity)
    cfg["physics"]["velocity"] = float(args.velocity)
    cfg["physics"]["rho"] = float(args.rho)

    cfg["method"]["collision"] = str(args.collision).upper()
    cfg["method"]["streaming"] = str(args.streaming)
    cfg["method"]["store_precision"] = str(args.store_precision).upper()
    cfg["method"]["ddf_shifting"] = bool(args.ddf_shifting)
    cfg["method"]["les"]["enabled"] = bool(args.enable_les)
    cfg["method"]["outlet"]["mode"] = str(args.outlet_mode)
    cfg["method"]["inlet"]["profile"] = str(args.inlet_profile)
    # Expose inlet scheme as a benchmark axis; useful when stability depends
    # on collision/inlet coupling.
    cfg["method"]["inlet"]["scheme"] = str(args.inlet_scheme)
    cfg["method"]["y_wall_bc"] = str(args.y_wall_bc)

    cfg["cuda"]["threads_per_block"] = int(args.threads_per_block)
    cfg["cuda"]["compute_capability"] = str(args.compute_capability)
    return cfg


def _build_body_cfg(args: argparse.Namespace) -> dict:
    if not args.with_cylinder:
        return {"objects": []}
    cx = 0.5 * float(args.nx)
    cy = 0.5 * float(args.ny)
    radius = max(2.0, min(float(args.nx), float(args.ny)) * 0.08)
    return {
        "objects": [
            {
                "type": "cylinder",
                "center": [cx, cy],
                "radius": radius,
            }
        ]
    }


def _maybe_probe_macroscopic(sim: Any, step: int, every: int, repeat: int) -> None:
    if every > 0 and step % every == 0:
        for _ in range(max(1, int(repeat))):
            _ = sim.get_macroscopic()


def _maybe_probe_ddf(sim: Any, step: int, every: int, repeat: int) -> None:
    if every > 0 and step % every == 0:
        for _ in range(max(1, int(repeat))):
            _ = sim.get_ddf()


def _maybe_checkpoint(sim: Any, step: int, every: int, out_dir: str) -> None:
    if every > 0 and step % every == 0:
        path = os.path.join(out_dir, f"checkpoint_step_{step:09d}.h5")
        sim.save_checkpoint(path)


def _maybe_probe_obs(sim: Any, stream: cuda.Stream, step: int, every: int) -> None:
    if every <= 0 or step % every != 0:
        return
    if sim.bodies.count <= 0:
        return
    sim.bodies.download_obs_full_async(stream)
    stream.synchronize()
    _ = sim.bodies.read_force(0)


def run(args: argparse.Namespace) -> Dict[str, Any]:
    if not os.path.isfile(_DEFAULT_LBM):
        raise FileNotFoundError(f"Base config missing: {_DEFAULT_LBM}")

    base = _load_json(_DEFAULT_LBM)
    lbm_cfg = _build_lbm_cfg(base, args)
    body_cfg = _build_body_cfg(args)

    tmpd = tempfile.mkdtemp(prefix="celeris_perf_baseline_")
    lbm_path = os.path.join(tmpd, "config_lbm.json")
    body_path = os.path.join(tmpd, "config_body.json")
    ckpt_dir = os.path.join(tmpd, "checkpoints")
    os.makedirs(ckpt_dir, exist_ok=True)
    _write_json(lbm_path, lbm_cfg)
    _write_json(body_path, body_cfg)

    from CelerisLab import Simulation  # noqa: WPS433

    sim = Simulation(lbm_config_path=lbm_path, body_config_path=body_path, device_id=args.device_id)
    sim.initialize()

    stream = cuda.Stream()
    total_cells = int(args.nx) * int(args.ny) * int(args.nz)

    # Warmup before measurement window.
    warmup_done = 0
    while warmup_done < int(args.warmup_steps):
        chunk = min(int(args.batch_steps), int(args.warmup_steps) - warmup_done)
        sim.stepper.step(chunk, action_gpu=sim.bodies.action_gpu, obs_gpu=sim.bodies.obs_gpu, stream=stream)
        warmup_done += chunk
    stream.synchronize()

    measured_batch_s: List[float] = []
    measured_steps = int(args.steps)
    done = 0
    t0 = time.perf_counter()
    while done < measured_steps:
        chunk = min(int(args.batch_steps), measured_steps - done)
        step_start = time.perf_counter()
        sim.stepper.step(chunk, action_gpu=sim.bodies.action_gpu, obs_gpu=sim.bodies.obs_gpu, stream=stream)
        stream.synchronize()
        step_end = time.perf_counter()
        measured_batch_s.append(step_end - step_start)

        done += chunk
        global_step = sim.stepper.step_count
        _maybe_probe_macroscopic(sim, global_step, int(args.macro_every), int(args.macro_repeat))
        _maybe_probe_ddf(sim, global_step, int(args.ddf_every), int(args.ddf_repeat))
        _maybe_checkpoint(sim, global_step, int(args.checkpoint_every), ckpt_dir)
        _maybe_probe_obs(sim, stream, global_step, int(args.obs_every))

    stream.synchronize()
    elapsed_s = time.perf_counter() - t0

    # Optional final sanity readback (outside core timing path by default).
    if args.final_macro_snapshot:
        _ = sim.get_macroscopic()

    sim.close()

    mlups = (total_cells * measured_steps) / max(elapsed_s, 1e-12) / 1.0e6
    batch_ms = [1000.0 * x for x in measured_batch_s]
    batch_ms_sorted = sorted(batch_ms)
    p50 = batch_ms_sorted[len(batch_ms_sorted) // 2] if batch_ms_sorted else 0.0
    p90 = batch_ms_sorted[min(len(batch_ms_sorted) - 1, int(0.9 * (len(batch_ms_sorted) - 1)))] if batch_ms_sorted else 0.0

    return {
        "benchmark": "celerislab_stepper_baseline",
        "device_id": int(args.device_id),
        "lattice_model": args.lattice_model,
        "grid": {"nx": int(args.nx), "ny": int(args.ny), "nz": int(args.nz)},
        "collision": str(args.collision).upper(),
        "inlet_scheme": str(args.inlet_scheme),
        "streaming": str(args.streaming),
        "store_precision": str(args.store_precision).upper(),
        "steps": measured_steps,
        "warmup_steps": int(args.warmup_steps),
        "batch_steps": int(args.batch_steps),
        "mlups": float(mlups),
        "elapsed_s": float(elapsed_s),
        "batch_ms_p50": float(p50),
        "batch_ms_p90": float(p90),
        "overhead_switches": {
            "macro_every": int(args.macro_every),
            "macro_repeat": int(args.macro_repeat),
            "ddf_every": int(args.ddf_every),
            "ddf_repeat": int(args.ddf_repeat),
            "checkpoint_every": int(args.checkpoint_every),
            "obs_every": int(args.obs_every),
            "with_cylinder": bool(args.with_cylinder),
            "final_macro_snapshot": bool(args.final_macro_snapshot),
        },
    }


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="CelerisLab pure-step performance baseline")
    ap.add_argument("--device-id", type=int, default=0)
    ap.add_argument("--lattice-model", choices=("D2Q9", "D3Q19"), default="D3Q19")
    ap.add_argument("--nx", type=int, default=256)
    ap.add_argument("--ny", type=int, default=256)
    ap.add_argument("--nz", type=int, default=256)
    ap.add_argument("--steps", type=int, default=3000)
    ap.add_argument("--warmup-steps", type=int, default=400)
    ap.add_argument("--batch-steps", type=int, default=100)

    ap.add_argument("--collision", choices=("SRT", "TRT", "MRT"), default="SRT")
    ap.add_argument("--streaming", choices=("double_buffer", "esopull"), default="double_buffer")
    ap.add_argument("--store-precision", choices=("FP32", "FP16S"), default="FP32")
    ap.add_argument("--ddf-shifting", action="store_true")
    ap.add_argument("--enable-les", action="store_true")
    ap.add_argument("--outlet-mode", choices=("neq_extrap", "zero_gradient", "blended"), default="neq_extrap")
    ap.add_argument("--inlet-profile", choices=("uniform", "parabolic"), default="uniform")
    ap.add_argument(
        "--inlet-scheme",
        choices=("zou_he_local", "channel_stabilized", "equilibrium", "regularized"),
        default="zou_he_local",
    )
    ap.add_argument("--y-wall-bc", choices=("bounce_back", "free_slip"), default="bounce_back")

    ap.add_argument("--viscosity", type=float, default=0.0035)
    ap.add_argument("--velocity", type=float, default=0.03)
    ap.add_argument("--rho", type=float, default=1.0)
    ap.add_argument("--threads-per-block", type=int, default=256)
    ap.add_argument("--compute-capability", type=str, default="auto")

    # Overhead attribution toggles.
    ap.add_argument("--macro-every", type=int, default=0, help="Call get_macroscopic() every N steps (0=off)")
    ap.add_argument("--macro-repeat", type=int, default=1, help="Repeat get_macroscopic() calls per probe step")
    ap.add_argument("--ddf-every", type=int, default=0, help="Call get_ddf() every N steps (0=off)")
    ap.add_argument("--ddf-repeat", type=int, default=1, help="Repeat get_ddf() calls per probe step")
    ap.add_argument("--checkpoint-every", type=int, default=0, help="Save checkpoint every N steps (0=off)")
    ap.add_argument("--obs-every", type=int, default=0, help="Download object obs every N steps (0=off)")
    ap.add_argument("--with-cylinder", action="store_true", help="Inject one cylinder object (needed for obs probes)")
    ap.add_argument("--final-macro-snapshot", action="store_true")

    ap.add_argument("--json-out", type=str, default="", help="Optional path to save metrics JSON")
    return ap.parse_args()


def main() -> int:
    args = parse_args()
    result = run(args)
    print(json.dumps(result, indent=2))
    if args.json_out.strip():
        out = os.path.abspath(args.json_out)
        os.makedirs(os.path.dirname(out), exist_ok=True)
        _write_json(out, result)
        print(f"Wrote: {out}")
    return 0


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