# CelerisLab/tests/run_kan99b_streakline.py
"""Kan99b streakline demo using CelerisLab common streakline engine.

Modes:
    - online: run CFD and compute streakline in memory (no velocity dump).
    - offline: load velocity snapshots and reconstruct streakline.
"""

from __future__ import annotations

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

import numpy as np

from CelerisLab import Simulation
from CelerisLab.common.streakline import (
    FlowFrame,
    IntegratorConfig,
    ReleaseConfig,
    build_release_points,
    estimate_sampling_plan,
    render_streakline_density,
    run_streakline_offline,
    run_streakline_online,
)

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

U_INF = 0.03
D_LATTICE = 30.0
R_LATTICE = 15.0
KAN99B_ST_REF = 0.1655


@dataclass(frozen=True)
class DomainSpec:
    key: str
    nx: int
    ny: int
    center: Tuple[float, float]


def _domain_specs() -> dict:
    return {
        "S": DomainSpec("S", 1081, 481, (360.0, 240.0)),
        "M": DomainSpec("M", 1351, 601, (450.0, 300.0)),
        "L": DomainSpec("L", 1801, 721, (600.0, 360.0)),
    }


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 _nu_from_re(reynolds: float) -> float:
    return U_INF * D_LATTICE / float(reynolds)


def _omega_body(alpha: float) -> float:
    return 2.0 * float(alpha) * U_INF / D_LATTICE


def _build_cfg(base_cfg: dict, *, nx: int, ny: int, re: float, inlet_scheme: str) -> dict:
    cfg = json.loads(json.dumps(base_cfg))
    cfg["grid"]["nx"] = int(nx)
    cfg["grid"]["ny"] = int(ny)
    cfg["grid"]["nz"] = 1
    cfg["physics"]["velocity"] = float(U_INF)
    cfg["physics"]["viscosity"] = float(_nu_from_re(re))
    cfg["physics"]["rho"] = 1.0
    cfg["method"]["collision"] = "MRT"
    cfg["method"]["streaming"] = "double_buffer"
    cfg["method"]["store_precision"] = "FP32"
    cfg["method"]["ddf_shifting"] = False
    cfg["method"]["les"]["enabled"] = False
    cfg["method"]["inlet"]["profile"] = "uniform"
    cfg["method"]["inlet"]["scheme"] = str(inlet_scheme)
    cfg["method"]["outlet"]["mode"] = "neq_extrap"
    cfg["method"]["y_wall_bc"] = "free_slip"
    return cfg


def _build_simulation(*, domain: DomainSpec, re: float, alpha: float, inlet_scheme: str) -> Simulation:
    base_cfg = _load_json(_DEFAULT_LBM)
    cfg = _build_cfg(base_cfg, nx=domain.nx, ny=domain.ny, re=re, inlet_scheme=inlet_scheme)
    body_doc = {
        "objects": [
            {
                "type": "cylinder",
                "center": [float(domain.center[0]), float(domain.center[1])],
                "radius": float(R_LATTICE),
                "omega": float(_omega_body(alpha)),
            }
        ]
    }
    tmpd = tempfile.mkdtemp(prefix="celeris_streakline_online_")
    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)
    sim = Simulation(lbm_config_path=lbm_tmp, body_config_path=body_tmp)
    sim.bodies.get(0).state.omega = np.float32(_omega_body(alpha))
    sim.initialize()
    return sim


def _default_base_release_points(center: Tuple[float, float]) -> np.ndarray:
    x_rel = float(center[0] - 6.0 * D_LATTICE)
    y0 = float(center[1])
    return np.array(
        [
            [x_rel, y0 - 18.0],
            [x_rel, y0 - 6.0],
            [x_rel, y0 + 6.0],
            [x_rel, y0 + 18.0],
        ],
        dtype=np.float64,
    )


def _load_offline_frames(snapshot_dir: str) -> List[FlowFrame]:
    files = sorted(
        [
            os.path.join(snapshot_dir, name)
            for name in os.listdir(snapshot_dir)
            if name.endswith(".npz") and name.startswith("vel_step")
        ]
    )
    frames: List[FlowFrame] = []
    for path in files:
        data = np.load(path)
        step = int(np.asarray(data["step"]).reshape(-1)[0])
        frames.append(
            FlowFrame(
                step=step,
                ux=np.asarray(data["ux"], dtype=np.float64),
                uy=np.asarray(data["uy"], dtype=np.float64),
            )
        )
    return frames


def main() -> int:
    ap = argparse.ArgumentParser(description="Kan99b streakline demo (online/offline)")
    ap.add_argument("--mode", default="online", choices=("online", "offline"))
    ap.add_argument("--domain", default="M", choices=("S", "M", "L"))
    ap.add_argument("--re", type=float, default=100.0)
    ap.add_argument("--alpha", type=float, default=1.0)
    ap.add_argument("--inlet-scheme", default="regularized", choices=("regularized", "zou_he_local"))
    ap.add_argument("--start-step", type=int, default=60_000)
    ap.add_argument("--sample-every", type=int, default=0, help="0 uses recommended value.")
    ap.add_argument("--n-snapshots", type=int, default=0, help="0 uses recommended value.")
    ap.add_argument("--snapshot-dir", type=str, default="", help="Required in offline mode.")
    ap.add_argument("--release-mode", default="strip", choices=("point", "line", "strip"))
    ap.add_argument("--line-span", type=float, default=0.0)
    ap.add_argument("--line-count", type=int, default=1)
    ap.add_argument(
        "--downstream-count",
        type=int,
        default=5,
        help="Particles per seed along +x (denser streak).",
    )
    ap.add_argument(
        "--downstream-spacing",
        type=float,
        default=1.0,
        help="Lattice spacing between x-staggered release points.",
    )
    ap.add_argument("--inject-per-seed", type=int, default=2)
    ap.add_argument("--alpha-t", type=float, default=0.2)
    ap.add_argument("--alpha-x", type=float, default=0.4)
    ap.add_argument("--diffusion-coeff", type=float, default=0.0)
    ap.add_argument(
        "--out-dir",
        type=str,
        default=os.path.join(_REPO, "tests", "output", "streakline", "kan99b_k2_online"),
    )
    args = ap.parse_args()

    domain = _domain_specs()[args.domain]
    out_dir = os.path.abspath(args.out_dir)
    os.makedirs(out_dir, exist_ok=True)

    plan = estimate_sampling_plan(st_ref=KAN99B_ST_REF, diameter=D_LATTICE, u_ref=U_INF)
    sample_every = int(args.sample_every) if int(args.sample_every) > 0 else int(plan["save_every_recommended"])
    n_snapshots = int(args.n_snapshots) if int(args.n_snapshots) > 0 else int(plan["snapshot_count_recommended"])

    release_cfg = ReleaseConfig(
        mode=args.release_mode,
        line_span=float(args.line_span),
        line_count=max(1, int(args.line_count)),
        downstream_count=max(1, int(args.downstream_count)),
        downstream_spacing=float(args.downstream_spacing),
        inject_per_seed=max(1, int(args.inject_per_seed)),
    )
    integrator_cfg = IntegratorConfig(
        alpha_t=float(args.alpha_t),
        alpha_x=float(args.alpha_x),
        diffusion_coeff=float(args.diffusion_coeff),
    )
    base_release_points = _default_base_release_points(domain.center)
    dense_release_points = build_release_points(base_release_points, release_cfg)

    if args.mode == "online":
        sim = _build_simulation(
            domain=domain,
            re=float(args.re),
            alpha=float(args.alpha),
            inlet_scheme=args.inlet_scheme,
        )
        try:
            particles, ages, diag = run_streakline_online(
                sim,
                start_step=int(args.start_step),
                sample_every=int(sample_every),
                n_samples=int(n_snapshots),
                release_points=base_release_points,
                release_cfg=release_cfg,
                integrator_cfg=integrator_cfg,
                solid_center=domain.center,
                solid_radius=R_LATTICE,
            )
        finally:
            sim.close()
        frame_count = int(n_snapshots)
        frame_source = "in-memory online sampling"
    else:
        if not args.snapshot_dir:
            raise ValueError("--snapshot-dir is required in offline mode.")
        frames = _load_offline_frames(os.path.abspath(args.snapshot_dir))
        if len(frames) < 2:
            raise RuntimeError("Offline mode needs at least two velocity snapshots.")
        particles, ages, diag = run_streakline_offline(
            frames,
            nx=domain.nx,
            ny=domain.ny,
            release_points=base_release_points,
            release_cfg=release_cfg,
            integrator_cfg=integrator_cfg,
            solid_center=domain.center,
            solid_radius=R_LATTICE,
        )
        frame_count = len(frames)
        frame_source = os.path.abspath(args.snapshot_dir)

    render_info = render_streakline_density(
        particles,
        ages,
        nx=domain.nx,
        ny=domain.ny,
        out_path=os.path.join(out_dir, "streakline.png"),
        release_points=dense_release_points,
        solid_center=domain.center,
        solid_radius=R_LATTICE,
        age_decay_steps=integrator_cfg.age_decay_steps,
        blur_sigma=1.2,
        title=f"Kan99b streakline ({args.mode}, {args.release_mode})",
    )

    meta = {
        "case": {
            "domain": args.domain,
            "re": float(args.re),
            "alpha": float(args.alpha),
            "inlet_scheme": args.inlet_scheme,
            "collision": "MRT",
        },
        "mode": args.mode,
        "sampling_estimate": plan,
        "sampling_used": {
            "start_step": int(args.start_step),
            "sample_every": int(sample_every),
            "n_snapshots": int(n_snapshots),
            "frame_source": frame_source,
            "frames_used": int(frame_count),
        },
        "release": {
            "config": vars(args),
            "base_points": base_release_points.tolist(),
            "dense_point_count": int(dense_release_points.shape[0]),
            "dense_points_preview": dense_release_points[: min(12, dense_release_points.shape[0])].tolist(),
        },
        "diagnostics": diag,
        "render": render_info,
    }
    _write_json(os.path.join(out_dir, "streakline_meta.json"), meta)

    print(f"Recommended sample_every: {plan['save_every_recommended']} steps")
    print(f"Recommended snapshots: {plan['snapshot_count_recommended']}")
    print(f"Mode: {args.mode} | frames used: {frame_count}")
    print(f"Dense release points: {dense_release_points.shape[0]}")
    print(f"Output image: {render_info['image_path']}")
    return 0


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