# CelerisLab/tests/postproc/run_kan99b_streakline.py
"""Kan99b streakline demo using the new Streakline class (online mode only).

Usage::

    python tests/run_kan99b_streakline.py --domain M --re 100 --alpha 1.0
"""

from __future__ import annotations

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

import numpy as np

from CelerisLab import Simulation
from CelerisLab.common.streakline import Streakline, ReleaseConfig, IntegratorConfig

_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_")
    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,
    )


# ---- Sampling plan helper (kept as a local utility, not part of the library) ----
def _estimate_sampling_plan(
    *,
    st_ref: float,
    diameter: float,
    u_ref: float,
    snapshots_per_period: float = 24.0,
    periods: int = 5,
) -> dict:
    period_steps = float(diameter) / (float(st_ref) * float(u_ref))
    save_every = int(
        max(20, round(period_steps / snapshots_per_period / 10.0) * 10)
    )
    n_snapshots = int(max(20, round(float(periods) * float(snapshots_per_period))))
    return {
        "st_ref": float(st_ref),
        "period_steps_est": float(period_steps),
        "save_every_recommended": int(save_every),
        "snapshot_count_recommended": int(n_snapshots),
    }


def main() -> int:
    ap = argparse.ArgumentParser(description="Kan99b streakline demo")
    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("--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)
    ap.add_argument("--downstream-spacing", type=float, default=1.0)
    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"
        ),
    )
    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 = _default_base_release_points(domain.center)

    streak = Streakline(
        release_points=base_release,
        release_cfg=release_cfg,
        integrator_cfg=integrator_cfg,
        nx=domain.nx,
        ny=domain.ny,
        cylinders=[(domain.center, R_LATTICE)],
    )

    sim = _build_simulation(
        domain=domain,
        re=float(args.re),
        alpha=float(args.alpha),
        inlet_scheme=args.inlet_scheme,
    )

    # Burn-in phase: step the simulation but don't feed streakline
    print(f"Burning-in {args.start_step} steps ...")
    sim.run(int(args.start_step))

    # Sampling phase: step and feed velocity frames to streakline
    target_last = int(args.start_step) + sample_every * (n_snapshots - 1)
    frames_collected = 0
    print(
        f"Sampling every {sample_every} steps for {n_snapshots} frames "
        f"(up to step {target_last})..."
    )
    while int(sim.stepper.step_count) < target_last:
        sim.step(1)
        step = int(sim.stepper.step_count)
        if (step - int(args.start_step)) % sample_every != 0:
            continue
        macro = sim.get_macroscopic()
        streak.observe(ux=macro["ux"], uy=macro["uy"], step=step)
        frames_collected += 1
        if frames_collected >= n_snapshots:
            break

    sim.close()

    render_info = streak.render(
        os.path.join(out_dir, "streakline.png"),
        age_decay_steps=integrator_cfg.age_decay_steps,
        blur_sigma=1.2,
    )

    meta = {
        "case": {
            "domain": args.domain,
            "re": float(args.re),
            "alpha": float(args.alpha),
            "inlet_scheme": args.inlet_scheme,
            "collision": "MRT",
        },
        "sampling_estimate": plan,
        "sampling_used": {
            "start_step": int(args.start_step),
            "sample_every": int(sample_every),
            "n_snapshots": int(n_snapshots),
            "frames_collected": frames_collected,
        },
        "release": {
            "config": {
                "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)),
            },
            "base_points": base_release.tolist(),
        },
        "diagnostics": {"n_particles_final": int(streak.n_particles)},
        "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"Frames collected: {frames_collected}")
    print(f"Final particles: {streak.n_particles}")
    print(f"Output image: {render_info['image_path']}")
    return 0


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