# CelerisLab/tests/postproc/run_exp_ctrl_matrix_streakline.py
"""Streakline post-processing for exp_ctrl_matrix cases.

Release from step RELEASE_START; render snapshots at SNAPSHOT_STEPS.
Particles are cleared after each snapshot except the last so each frame
uses a fresh ~20k-step release window.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

_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
from CelerisLab import Simulation
from CelerisLab.common.preprocess import build_triangle_release_points, cylinders_from_triangle_layout
from CelerisLab.common.streakline import IntegratorConfig, ReleaseConfig, Streakline

DIAMETER_CELLS = vort.DIAMETER_CELLS
DEFAULT_OUT = _REPO / "tests" / "output" / "exp_ctrl_matrix_streak_nx1500"
RELEASE_START_STEP = 20_000
SNAPSHOT_STEPS = (40_000, 60_000, 80_000, 100_000)
CLEAR_AFTER_SNAPSHOT = frozenset({40_000, 60_000, 80_000})
STREAK_SAMPLE_EVERY = 50
STREAK_AGE_DECAY = 100_000.0
STREAK_BLUR_SIGMA = 0.8
STREAK_COLOR = (0.0, 0.35, 0.95)  # blue particles on white background


def _apply_body_actions(
    sim: Simulation, a1: float, a2: float, a3: float, u_lb: float
) -> None:
    w1 = vort._action_to_omega_lb(a1, u_lb)
    w2 = vort._action_to_omega_lb(a2, u_lb)
    w3 = vort._action_to_omega_lb(a3, u_lb)
    if vort.SWAP_ACTION23_BODIES:
        vort._set_body_omegas(sim, w1, w3, w2)
    else:
        vort._set_body_omegas(sim, w1, w2, w3)


def _render_streak(
    streak: Streakline,
    out_path: Path,
    *,
    step: int,
) -> dict:
    info = streak.render(
        str(out_path),
        age_decay_steps=STREAK_AGE_DECAY,
        blur_sigma=STREAK_BLUR_SIGMA,
        background_color=(1.0, 1.0, 1.0),
        streak_color=STREAK_COLOR,
    )
    info["step"] = int(step)
    info["n_particles"] = int(streak.n_particles)
    return info


def run_streak_case(
    case_id: str,
    slug: str,
    features: dict,
    *,
    out_dir: Path,
    total_steps: int,
    release_start: int,
    snapshot_steps: tuple[int, ...],
    sample_every: int,
    report_every: int,
    device_id: int,
) -> dict:
    compat = vort._ensure_compat_config(vort.CONFIG_PATH)
    sim = Simulation(compat, device_id=device_id)
    layout = vort._add_triangle_cylinders(sim)
    sim.initialize()
    u_lb = float(sim.lbm_cfg.velocity)
    dt_phys = (
        (vort.CYLINDER_DIAMETER_M / DIAMETER_CELLS)
        * (u_lb / vort.INLET_U_PHYS_M_S)
    )
    cylinders = cylinders_from_triangle_layout(layout)

    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
    )
    base_release = build_triangle_release_points(
        layout,
        nx=int(sim.lbm_cfg.nx),
        ny=int(sim.lbm_cfg.ny),
        diameter_cells=DIAMETER_CELLS,
    )

    streak = Streakline(
        release_points=base_release,
        release_cfg=release_cfg,
        integrator_cfg=integrator_cfg,
        nx=int(sim.lbm_cfg.nx),
        ny=int(sim.lbm_cfg.ny),
        cylinders=cylinders,
    )

    snapshot_set = set(snapshot_steps)
    print(
        f"--- {case_id} {slug} steps={total_steps} release_from={release_start} "
        f"snapshots={list(snapshot_steps)} device={device_id} ---"
    )

    snapshots: list[dict] = []

    for step in range(total_steps):
        t_phys = step * dt_phys
        a1, a2, a3 = vort._actions_at_time(t_phys, features)
        _apply_body_actions(sim, a1, a2, a3, u_lb)
        sim.run(1)

        current_step = step + 1
        observed = False
        if step >= release_start and current_step % sample_every == 0:
            macro = sim.get_macroscopic()
            streak.observe(ux=macro["ux"], uy=macro["uy"], step=current_step)
            observed = True

        if current_step in snapshot_set:
            if streak.n_particles == 0:
                raise RuntimeError(
                    f"{case_id}: no particles at step {current_step}; "
                    f"check release_start/sample_every."
                )
            png = out_dir / f"streakline_{case_id}_{slug}_step{current_step:06d}.png"
            snap_info = _render_streak(streak, png, step=current_step)
            snap_info["image_path"] = str(png)
            snapshots.append(snap_info)
            print(
                f"  snapshot step {current_step} particles={streak.n_particles} "
                f"-> {png.name}"
            )

            if current_step in CLEAR_AFTER_SNAPSHOT and current_step < total_steps:
                streak.reset()
                if observed:
                    macro = sim.get_macroscopic()
                    streak.observe(ux=macro["ux"], uy=macro["uy"], step=current_step)

        if report_every > 0 and current_step % report_every == 0:
            print(
                f"  step {current_step}/{total_steps} "
                f"a=({a1:+.5f},{a2:+.5f},{a3:+.5f}) particles={streak.n_particles}"
            )

    sim.close()

    summary = {
        "case_id": case_id,
        "slug": slug,
        "total_steps": int(total_steps),
        "release_start_step": int(release_start),
        "snapshot_steps": list(snapshot_steps),
        "clear_after_snapshot": sorted(CLEAR_AFTER_SNAPSHOT),
        "sample_every": int(sample_every),
        "device_id": int(device_id),
        "release_points_dense": int(base_release.shape[0]),
        "snapshots": snapshots,
        "swap_action23_bodies": bool(vort.SWAP_ACTION23_BODIES),
    }
    with (out_dir / f"summary_{case_id}_{slug}.json").open("w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)
    return summary


def main() -> int:
    ap = argparse.ArgumentParser(description="exp_ctrl_matrix streakline batch")
    ap.add_argument("--out-dir", type=str, default=str(DEFAULT_OUT))
    ap.add_argument("--steps", type=int, default=vort.FIXED_STEPS)
    ap.add_argument("--release-start", type=int, default=RELEASE_START_STEP)
    ap.add_argument(
        "--snapshots",
        type=str,
        default=",".join(str(s) for s in SNAPSHOT_STEPS),
        help="Comma-separated render steps, e.g. 40000,60000,100000",
    )
    ap.add_argument("--sample-every", type=int, default=STREAK_SAMPLE_EVERY)
    ap.add_argument("--report-every", type=int, default=20000)
    ap.add_argument("--device-id", type=int, default=2)
    ap.add_argument("--cases", type=str, default="")
    args = ap.parse_args()

    snapshot_steps = tuple(
        int(s.strip()) for s in args.snapshots.split(",") if s.strip()
    )
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    selected = (
        {c.strip() for c in args.cases.split(",") if c.strip()}
        if args.cases
        else None
    )

    with vort.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"release from {args.release_start} | snapshots={list(snapshot_steps)} | "
        f"device={args.device_id}"
    )

    summaries = []
    for case_id, slug, features in vort.CONTROL_CASES:
        if selected and case_id not in selected:
            continue
        summaries.append(
            run_streak_case(
                case_id,
                slug,
                features,
                out_dir=out_dir,
                total_steps=int(args.steps),
                release_start=int(args.release_start),
                snapshot_steps=snapshot_steps,
                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),
        "release_start_step": int(args.release_start),
        "snapshot_steps": list(snapshot_steps),
        "clear_after_snapshot": sorted(CLEAR_AFTER_SNAPSHOT),
        "sample_every": int(args.sample_every),
        "streak_color": list(STREAK_COLOR),
        "stealth_steady_omega_m_s": float(vort.STEALTH_STEADY_OMEGA_M_S),
        "device_id": int(args.device_id),
        "swap_action23_bodies": bool(vort.SWAP_ACTION23_BODIES),
        "cases": summaries,
    }
    manifest_path = out_dir / "manifest.json"
    with manifest_path.open("w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2)
    print(f"Manifest: {manifest_path}")
    return 0


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