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

Usage::

    # Single case
    conda run -n pycuda_3_10 python tests/postproc/run_exp_ctrl_matrix_vorticity.py \\
        --cases C0 --batch 10 --device-id 0

    # All cases
    conda run -n pycuda_3_10 python tests/postproc/run_exp_ctrl_matrix_vorticity.py \\
        --batch 10 --device-id 0

    # Full 100k steps, no batching (default)
    conda run -n pycuda_3_10 python tests/postproc/run_exp_ctrl_matrix_vorticity.py
"""

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

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

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

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_body("circle", center=(layout["x_apex"], layout["y_center"]),
                  radius=layout["radius_lb"])
    sim.add_body("circle", center=(layout["x_rear"], layout["y_lower"]),
                  radius=layout["radius_lb"])
    sim.add_body("circle", 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:
    """Set all three body rotation speeds using new API (implicit GPU upload)."""
    sim.set_body(0, omega=omega0)
    sim.set_body(1, omega=omega1)
    sim.set_body(2, omega=omega2)


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,
    batch: int = 1,
    device_id: int = 0,
) -> dict:
    compat = _ensure_compat_config(CONFIG_PATH)
    sim = Simulation(compat, device_id=device_id)
    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} batch={batch} ---")
    stream = sim.stream
    batch_size = max(1, int(batch))

    # Main loop: precompute actions, batch-step, read forces/sensors at intervals
    for batch_start in range(0, steps, batch_size):
        batch_end = min(batch_start + batch_size, steps)
        for j in range(batch_start, batch_end):
            t_phys = j * 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)

        n = batch_end - batch_start
        sim.stepper.step(
            n,
            action_gpu=sim.bodies.action_gpu,
            obs_gpu=sim.bodies.obs_gpu,
            stream=stream,
        )

        if report_every > 0 and (batch_end % report_every == 0 or batch_end == steps):
            stream.synchronize()
            for bid in range(sim.bodies.count):
                fx = sim.bodies.read_force(bid)
                print(
                    f"  step={batch_end} body={bid}"
                    f" fx={float(fx[0]):+.6f} fy={float(fx[1]):+.6f}",
                    flush=True,
                )

    stream.synchronize()
    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),
        "batch": int(batch),
        "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}).",
    )
    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.")
    ap.add_argument(
        "--batch", type=int, default=1,
        help=(
            "Batch N steps between action uploads.  Default 1 (each step).  "
            "With --batch 10, actions are computed and uploaded every 10 steps.  "
            "Saves kernel launch overhead at the cost of control-signal interpolation."
        ),
    )
    ap.add_argument("--device-id", type=int, default=0, help="GPU device id.")
    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} | batch={args.batch} | "
        f"device={args.device_id} | 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),
                batch=int(args.batch),
                device_id=int(args.device_id),
            )
        )

    manifest = {
        "grid": {"nx": nx, "ny": ny},
        "base_steps": base_steps,
        "step_multiplier": float(args.step_multiplier),
        "steps": steps,
        "batch": int(args.batch),
        "device_id": int(args.device_id),
        "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())