# CelerisLab/tests/run_exp_ctrl_matrix_vorticity.py
"""Batch vorticity images for three-cylinder control matrix (exp_ctrl_matrix.md).

Runs each control law to steady-like state, saves final vorticity PNG only (no streaklines).
"""

from __future__ import annotations

import argparse
import json
import math
import os
import sys
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Tuple

import numpy as np
import pycuda.driver as cuda

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

from CelerisLab import Simulation
from CelerisLab.common.streakline import (
    compute_vorticity,
    cylinders_from_triangle_layout,
    render_vorticity_field,
)

INLET_U_PHYS_M_S = 0.009028
CYLINDER_DIAMETER_M = 0.010
CENTER_SPACING_M = 0.015
DIAMETER_CELLS = 20.0
RAMP_TIME_S = 5.0
INITIAL_ACTIONS_M_S = (0.0, 0.0, 0.0)
OMEGA_SIGN_FROM_ACTION = -1.0
VORT_VMIN = -0.003
VORT_VMAX = 0.003

CONFIG_PATH = _REPO / "src/CelerisLab/configs/config_lbm_three_cylinder_triangle.json"
DEFAULT_OUT = _REPO / "tests" / "output" / "exp_ctrl_matrix_vort_ny300"
FIXED_STEPS = 100000  # keep constant while grid changes
# Body order: 0=apex, 1=rear-lower(y_lower), 2=rear-upper(y_upper); swap action2/action3 targets.
SWAP_ACTION23_BODIES = True

# From tests/exp_ctrl_matrix.md (SIGNAL_FEATURES0 .. 6)
CONTROL_CASES: List[Tuple[str, str, Dict[str, Any]]] = [
    (
        "C0",
        "no_ctrl",
        {
            "action1": {"mean": 0.0, "components": [(0.1354, 0.0, 1.600)]},
            "action2": {"mean": 0.0, "components": [(0.1354, 0.0, 2.099)]},
            "action3": {"mean": 0.0, "components": [(0.1354, 0.0, 1.639)]},
        },
    ),
    (
        "C1",
        "stealth",
        {
            "action1": {"mean": 0.0, "components": [(0.1354, 0.0, 1.600)]},
            "action2": {"mean": -0.01806, "components": [(0.1354, 0.0, 2.099)]},
            "action3": {"mean": 0.01806, "components": [(0.1354, 0.0, 1.639)]},
        },
    ),
    (
        "C2",
        "deceit",
        {
            "action1": {"mean": 0.0, "components": [(0.1354, 0.0026, 1.600)]},
            "action2": {
                "mean": -0.008730,
                "components": [(0.1354, 0.0045, 2.099), (0.2708, 0.0010, 0.612)],
            },
            "action3": {
                "mean": 0.008730,
                "components": [(0.1354, 0.0045, 1.639), (0.2708, 0.0010, -2.962)],
            },
        },
    ),
    (
        "C3",
        "deceit_multi",
        {
            "action1": {
                "mean": 0.0,
                "components": [(0.1354, 0.0029, -2.619), (0.2708, 0.0008, 2.856)],
            },
            "action2": {
                "mean": -0.0140,
                "components": [
                    (0.1354, 0.0050, -0.933),
                    (0.2708, 0.0010, 0.801),
                    (0.1806, 0.0003, 1.854),
                ],
            },
            "action3": {
                "mean": 0.014,
                "components": [
                    (0.1354, 0.0050, -1.398),
                    (0.2708, 0.0010, 2.208),
                    (0.1806, 0.0003, 1.810),
                ],
            },
        },
    ),
    (
        "C4",
        "deceit_f1p5",
        {
            "action1": {"mean": 0.0, "components": [(0.2031, 0.0026, 1.600)]},
            "action2": {
                "mean": -0.008730,
                "components": [(0.2031, 0.0045, 2.099), (0.4062, 0.0010, 0.612)],
            },
            "action3": {
                "mean": 0.008730,
                "components": [(0.2031, 0.0045, 1.639), (0.4062, 0.0010, -2.962)],
            },
        },
    ),
    (
        "C5",
        "deceit_multi_f1p5",
        {
            "action1": {
                "mean": 0.0,
                "components": [(0.2031, 0.0029, -2.619), (0.4062, 0.0008, 2.856)],
            },
            "action2": {
                "mean": -0.0140,
                "components": [
                    (0.2031, 0.0050, -0.933),
                    (0.4062, 0.0010, 0.801),
                    (0.2709, 0.0003, 1.854),
                ],
            },
            "action3": {
                "mean": 0.014,
                "components": [
                    (0.2031, 0.0050, -1.398),
                    (0.4062, 0.0010, 2.208),
                    (0.2709, 0.0003, 1.810),
                ],
            },
        },
    ),
    (
        "C6",
        "deceit_f2",
        {
            "action1": {
                "mean": 0.0,
                "components": [(0.2708, 0.0044, -2.619), (0.8124, 0.0012, 2.856)],
            },
            "action2": {
                "mean": -0.014,
                "components": [
                    (0.2708, 0.0075, -0.933),
                    (0.8124, 0.0015, 0.801),
                    (0.5418, 0.0005, 1.854),
                ],
            },
            "action3": {
                "mean": 0.014,
                "components": [
                    (0.2708, 0.0075, -1.398),
                    (0.8124, 0.0015, 2.208),
                    (0.5418, 0.0005, 1.810),
                ],
            },
        },
    ),
]


def _ensure_compat_config(config_path: Path, preferred_scheme: str = "regularized") -> str:
    with config_path.open("r", encoding="utf-8") as f:
        cfg = json.load(f)
    method = cfg.setdefault("method", {})
    inlet = method.setdefault("inlet", {})
    outlet = method.setdefault("outlet", {})
    changed = False
    if "scheme" not in inlet:
        inlet["scheme"] = preferred_scheme
        changed = True
    if "regularized_neq_damp" not in inlet:
        inlet["regularized_neq_damp"] = 0.5
        changed = True
    if "blend_alpha" not in outlet:
        outlet["blend_alpha"] = 0.7
        changed = True
    if "backflow_clamp" not in outlet:
        outlet["backflow_clamp"] = True
        changed = True
    if not changed:
        return str(config_path)
    tmp = tempfile.NamedTemporaryFile(
        mode="w", suffix="_compat_lbm.json", delete=False, encoding="utf-8"
    )
    with tmp:
        json.dump(cfg, tmp, indent=4)
    return tmp.name


def _triangle_layout(cfg) -> dict:
    dx_phys = CYLINDER_DIAMETER_M / DIAMETER_CELLS
    spacing_lb = CENTER_SPACING_M / dx_phys
    radius_lb = DIAMETER_CELLS / 2.0
    y_center = 0.5 * (cfg.ny - 1)
    x_cluster_center = cfg.nx / 3.0
    x_apex = x_cluster_center - (math.sqrt(3.0) / 3.0) * spacing_lb
    x_rear = x_apex + (math.sqrt(3.0) / 2.0) * spacing_lb
    return {
        "x_apex": x_apex,
        "x_rear": x_rear,
        "y_center": y_center,
        "y_upper": y_center + 0.5 * spacing_lb,
        "y_lower": y_center - 0.5 * spacing_lb,
        "radius_lb": radius_lb,
    }


def _add_triangle_cylinders(sim: Simulation) -> dict:
    layout = _triangle_layout(sim.lbm_cfg)
    sim.add_cylinder(center=(layout["x_apex"], layout["y_center"]), radius=layout["radius_lb"])
    sim.add_cylinder(center=(layout["x_rear"], layout["y_lower"]), radius=layout["radius_lb"])
    sim.add_cylinder(center=(layout["x_rear"], layout["y_upper"]), radius=layout["radius_lb"])
    return layout


def _generate_signal(t_phys: float, feature: dict) -> float:
    value = float(feature["mean"])
    for freq_hz, amp, phase in feature["components"]:
        value += amp * math.cos(2.0 * math.pi * freq_hz * t_phys + phase)
    return value


def _ramp_factor(elapsed_s: float) -> float:
    if elapsed_s <= 0.0:
        return 0.0
    if elapsed_s >= RAMP_TIME_S:
        return 1.0
    x = elapsed_s / RAMP_TIME_S
    return 0.5 * (1.0 - math.cos(math.pi * x))


def _actions_at_time(t_phys: float, features: dict) -> Tuple[float, float, float]:
    s1 = _generate_signal(t_phys, features["action1"])
    s2 = _generate_signal(t_phys, features["action2"])
    s3 = _generate_signal(t_phys, features["action3"])
    r = _ramp_factor(t_phys)
    a1 = INITIAL_ACTIONS_M_S[0] * (1.0 - r) + s1 * r
    a2 = INITIAL_ACTIONS_M_S[1] * (1.0 - r) + s2 * r
    a3 = INITIAL_ACTIONS_M_S[2] * (1.0 - r) + s3 * r
    return a1, a2, a3


def _action_to_omega_lb(action_m_s: float, u_lb: float) -> float:
    u_surf_lb = action_m_s * (u_lb / INLET_U_PHYS_M_S)
    r_lb = DIAMETER_CELLS / 2.0
    return OMEGA_SIGN_FROM_ACTION * (u_surf_lb / r_lb)


def _set_body_omegas(sim: Simulation, omega0: float, omega1: float, omega2: float) -> None:
    bodies = sim.bodies
    dim = sim.lbm_cfg.dim
    slot = 3 * dim
    bodies.action.fill(0.0)
    bodies.action[(0 * slot) + slot - 1] = np.float32(omega0)
    bodies.action[(1 * slot) + slot - 1] = np.float32(omega1)
    bodies.action[(2 * slot) + slot - 1] = np.float32(omega2)
    cuda.memcpy_htod(bodies.action_gpu, bodies.action)


def _default_steps(nx: int, u_lb: float, step_multiplier: float) -> int:
    base = int(round(2.0 * float(nx) / (3.0 * float(u_lb))))
    return int(round(base * float(step_multiplier)))


def run_case(
    case_id: str,
    slug: str,
    features: dict,
    *,
    out_dir: Path,
    steps: int,
    report_every: int,
) -> dict:
    compat = _ensure_compat_config(CONFIG_PATH)
    sim = Simulation(compat)
    layout = _add_triangle_cylinders(sim)
    sim.initialize()
    u_lb = float(sim.lbm_cfg.velocity)
    dx_phys = CYLINDER_DIAMETER_M / DIAMETER_CELLS
    dt_phys = dx_phys * (u_lb / INLET_U_PHYS_M_S)
    cylinders = cylinders_from_triangle_layout(layout)

    print(f"--- {case_id} {slug} steps={steps} u_lb={u_lb} dt_phys={dt_phys} ---")
    for i in range(steps):
        t_phys = i * dt_phys
        a1, a2, a3 = _actions_at_time(t_phys, features)
        w1 = _action_to_omega_lb(a1, u_lb)
        w2 = _action_to_omega_lb(a2, u_lb)
        w3 = _action_to_omega_lb(a3, u_lb)
        if SWAP_ACTION23_BODIES:
            _set_body_omegas(sim, w1, w3, w2)
        else:
            _set_body_omegas(sim, w1, w2, w3)
        sim.run(1)
        if report_every > 0 and (i + 1) % report_every == 0:
            print(f"  step {i+1}/{steps} a=({a1:+.5f},{a2:+.5f},{a3:+.5f})")

    macro = sim.get_macroscopic()
    vort = compute_vorticity(macro["ux"], macro["uy"])
    png = out_dir / f"vorticity_{case_id}_{slug}.png"
    ckpt = out_dir / f"state_{case_id}_{slug}.h5"
    sim.save_checkpoint(str(ckpt))

    render_info = render_vorticity_field(
        vort,
        nx=int(sim.lbm_cfg.nx),
        ny=int(sim.lbm_cfg.ny),
        out_path=str(png),
        cylinders=cylinders,
        vmin=VORT_VMIN,
        vmax=VORT_VMAX,
        minimal_axes=True,
    )
    sim.close()

    summary = {
        "case_id": case_id,
        "slug": slug,
        "steps": int(steps),
        "u_lb": u_lb,
        "dt_phys": dt_phys,
        "vort_png": str(png),
        "checkpoint": str(ckpt),
        "vort_range_data": [float(vort.min()), float(vort.max())],
        "vort_plot_range": [VORT_VMIN, VORT_VMAX],
        "swap_action23_bodies": bool(SWAP_ACTION23_BODIES),
        "render": render_info,
    }
    with (out_dir / f"summary_{case_id}_{slug}.json").open("w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)
    print(f"  saved {png}")
    return summary


def main() -> int:
    ap = argparse.ArgumentParser(description="exp_ctrl_matrix vorticity batch")
    ap.add_argument("--out-dir", type=str, default=str(DEFAULT_OUT))
    ap.add_argument(
        "--step-multiplier",
        type=float,
        default=2.0,
        help="Steps = multiplier * round(2*nx/(3*u_lb)); use 2.0 after halving nx/ny.",
    )
    ap.add_argument(
        "--steps",
        type=int,
        default=FIXED_STEPS,
        help=f"Total LBM steps (default {FIXED_STEPS}, held fixed across grid tweaks).",
    )
    ap.add_argument("--report-every", type=int, default=20000)
    ap.add_argument("--cases", type=str, default="", help="Comma list e.g. C0,C1 or empty=all.")
    args = ap.parse_args()

    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
    summaries = []

    with CONFIG_PATH.open("r", encoding="utf-8") as f:
        grid_cfg = json.load(f)["grid"]
    nx = int(grid_cfg["nx"])
    ny = int(grid_cfg["ny"])
    u_lb = 0.04
    base_steps = int(round(2.0 * nx / (3.0 * u_lb)))
    steps = int(args.steps) if int(args.steps) > 0 else _default_steps(nx, u_lb, args.step_multiplier)
    print(
        f"Output: {out_dir} | grid={nx}x{ny} | base_steps={base_steps} "
        f"x{args.step_multiplier} -> {steps} | vort [{VORT_VMIN}, {VORT_VMAX}]"
    )

    for case_id, slug, features in CONTROL_CASES:
        if selected and case_id not in selected:
            continue
        summaries.append(
            run_case(
                case_id,
                slug,
                features,
                out_dir=out_dir,
                steps=steps,
                report_every=int(args.report_every),
            )
        )

    manifest = {
        "grid": {"nx": nx, "ny": ny},
        "base_steps": base_steps,
        "step_multiplier": float(args.step_multiplier),
        "steps": steps,
        "vort_vmin": VORT_VMIN,
        "vort_vmax": VORT_VMAX,
        "swap_action23_bodies": bool(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())