# CelerisLab/tests/postproc/run_stealth_steady_sweep.py
"""Steady stealth rotation sweep: vorticity + final-step streakline per speed.

Grid nx=1500 (see config_lbm_three_cylinder_triangle.json). Steady means
constant surface-speed means only (no harmonic components).
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import List, Tuple

_REPO = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(_REPO / "src"))
sys.path.insert(0, str(_REPO / "tests" / "postproc"))

import run_exp_ctrl_matrix_vorticity as vort_mod
import run_exp_ctrl_matrix_streakline as streak
from CelerisLab import Simulation
from CelerisLab.common.preprocess import build_triangle_release_points, cylinders_from_triangle_layout
from CelerisLab.common.render import compute_vorticity, render_vorticity_field
from CelerisLab.common.streakline import Streakline, IntegratorConfig, ReleaseConfig

STEALTH_REF_M_S = 0.01806
# Fractions of reference stealth surface speed (action2 negative, action3 positive).
SPEED_FRACTIONS: List[Tuple[str, float]] = [
    ("s050", 0.50),
    ("s075", 0.75),
    ("s100", 1.00),
    ("s125", 1.25),
    ("s150", 1.50),
]
DEFAULT_OUT = _REPO / "tests" / "output" / "stealth_steady_sweep_nx1500"


def _stealth_features(omega_m_s: float) -> dict:
    return {
        "action1": {"mean": 0.0, "components": []},
        "action2": {"mean": -float(omega_m_s), "components": []},
        "action3": {"mean": float(omega_m_s), "components": []},
    }


def run_one(
    tag: str,
    omega_m_s: float,
    *,
    out_dir: Path,
    total_steps: int,
    streak_window: int,
    sample_every: int,
    report_every: int,
    device_id: int,
) -> dict:
    features = _stealth_features(omega_m_s)
    slug = f"stealth_{tag}_w{omega_m_s:.5f}"

    compat = vort_mod._ensure_compat_config(vort_mod.CONFIG_PATH)
    sim = Simulation(compat, device_id=device_id)
    layout = vort_mod._add_triangle_cylinders(sim)
    sim.initialize()
    u_lb = float(sim.lbm_cfg.velocity)
    nx = int(sim.lbm_cfg.nx)
    ny = int(sim.lbm_cfg.ny)
    dt_phys = (vort_mod.CYLINDER_DIAMETER_M / vort_mod.DIAMETER_CELLS) * (
        u_lb / vort_mod.INLET_U_PHYS_M_S
    )
    cylinders = cylinders_from_triangle_layout(layout)
    base_release = build_triangle_release_points(
        layout, nx=nx, ny=ny, diameter_cells=vort_mod.DIAMETER_CELLS
    )

    release_cfg = ReleaseConfig(
        mode="strip",
        line_count=1,
        line_span=0.0,
        downstream_count=5,
        downstream_spacing=1.0,
        inject_per_seed=1,
    )
    integrator_cfg = IntegratorConfig(alpha_t=0.25, alpha_x=0.40, max_particle_age=None)
    streak_obj = Streakline(
        release_points=base_release,
        release_cfg=release_cfg,
        integrator_cfg=integrator_cfg,
        nx=nx,
        ny=ny,
        cylinders=cylinders,
    )

    streak_start = max(0, int(total_steps) - int(streak_window))

    print(
        f"--- {slug} omega={omega_m_s:.5f} m/s steps={total_steps} "
        f"grid={nx}x{ny} streak_from={streak_start} ---"
    )
    print(
        f"    layout x_apex={layout['x_apex']:.1f} x_rear={layout['x_rear']:.1f} "
        f"release_x={base_release[0, 0]:.1f}"
    )

    for step in range(total_steps):
        t_phys = step * dt_phys
        a1, a2, a3 = vort_mod._actions_at_time(t_phys, features)
        w1 = vort_mod._action_to_omega_lb(a1, u_lb)
        w2 = vort_mod._action_to_omega_lb(a2, u_lb)
        w3 = vort_mod._action_to_omega_lb(a3, u_lb)
        if vort_mod.SWAP_ACTION23_BODIES:
            vort_mod._set_body_omegas(sim, w1, w3, w2)
        else:
            vort_mod._set_body_omegas(sim, w1, w2, w3)
        sim.run(1)

        if report_every > 0 and (step + 1) % report_every == 0:
            print(
                f"  step {step + 1}/{total_steps} a=({a1:+.5f},{a2:+.5f},{a3:+.5f}) "
                f"omega_lb=({w1:+.6f},{w2:+.6f},{w3:+.6f}) "
                f"particles={streak_obj.n_particles}"
            )

        if step >= streak_start and (step + 1) % sample_every == 0:
            macro = sim.get_macroscopic()
            streak_obj.observe(ux=macro["ux"], uy=macro["uy"], step=int(step + 1))

    if streak_obj.n_particles == 0:
        raise RuntimeError(f"{slug}: no particles in streak window.")

    macro = sim.get_macroscopic()
    vort_field = compute_vorticity(macro["ux"], macro["uy"])

    vort_png = out_dir / f"vorticity_{slug}.png"
    streak_png = out_dir / f"streakline_{slug}.png"
    ckpt = out_dir / f"state_{slug}.h5"
    sim.save_checkpoint(str(ckpt))

    vort_info = render_vorticity_field(
        vort_field,
        nx=nx,
        ny=ny,
        out_path=str(vort_png),
        cylinders=cylinders,
        vmin=vort_mod.VORT_VMIN,
        vmax=vort_mod.VORT_VMAX,
        minimal_axes=True,
    )
    streak_info = streak_obj.render(
        str(streak_png),
        age_decay_steps=streak.STREAK_AGE_DECAY,
        blur_sigma=streak.STREAK_BLUR_SIGMA,
        background_color=(1.0, 1.0, 1.0),
        streak_color=streak.STREAK_COLOR,
    )
    sim.close()

    summary = {
        "tag": tag,
        "slug": slug,
        "omega_m_s": float(omega_m_s),
        "fraction_of_ref": float(omega_m_s / STEALTH_REF_M_S),
        "total_steps": int(total_steps),
        "streak_window_steps": int(streak_window),
        "streak_start_step": int(streak_start),
        "sample_every": int(sample_every),
        "layout": {k: float(layout[k]) for k in layout},
        "release_points": base_release.tolist(),
        "particle_count_final": int(streak_obj.n_particles),
        "vort_png": str(vort_png),
        "streak_png": str(streak_png),
        "checkpoint": str(ckpt),
        "vorticity": vort_info,
        "streakline": streak_info,
    }
    with (out_dir / f"summary_{slug}.json").open("w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)
    print(f"  saved {vort_png.name} {streak_png.name} particles={streak_obj.n_particles}")
    return summary


def main() -> int:
    ap = argparse.ArgumentParser(description="steady stealth rotation sweep")
    ap.add_argument("--out-dir", type=str, default=str(DEFAULT_OUT))
    ap.add_argument("--steps", type=int, default=vort_mod.FIXED_STEPS)
    ap.add_argument("--streak-window", type=int, default=streak.STREAK_WINDOW_STEPS)
    ap.add_argument("--sample-every", type=int, default=streak.STREAK_SAMPLE_EVERY)
    ap.add_argument("--report-every", type=int, default=20000)
    ap.add_argument("--device-id", type=int, default=0)
    ap.add_argument("--tags", type=str, default="", help="Comma tags e.g. s050,s100 or empty=all")
    args = ap.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    selected = {t.strip() for t in args.tags.split(",") if t.strip()} if args.tags else None

    with vort_mod.CONFIG_PATH.open("r", encoding="utf-8") as f:
        grid = json.load(f)["grid"]

    print(
        f"Output: {out_dir} | grid={grid['nx']}x{grid['ny']} | steps={args.steps} | "
        f"ref_omega={STEALTH_REF_M_S} m/s"
    )

    summaries = []
    for tag, frac in SPEED_FRACTIONS:
        if selected and tag not in selected:
            continue
        omega = STEALTH_REF_M_S * frac
        summaries.append(
            run_one(
                tag,
                omega,
                out_dir=out_dir,
                total_steps=int(args.steps),
                streak_window=int(args.streak_window),
                sample_every=int(args.sample_every),
                report_every=int(args.report_every),
                device_id=int(args.device_id),
            )
        )

    manifest = {
        "grid": grid,
        "steps": int(args.steps),
        "stealth_ref_m_s": STEALTH_REF_M_S,
        "speed_fractions": SPEED_FRACTIONS,
        "streak_color": list(streak.STREAK_COLOR),
        "cases": summaries,
    }
    with (out_dir / "manifest.json").open("w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2)
    print(f"Manifest: {out_dir / 'manifest.json'}")
    return 0


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