# OID_analysis/scripts/collect_target_cylinder.py
"""
Collect target cylinder data (q_tar) for illusion comparison.
Single cylinder at x=20*L0 with target diameter, plus 3 sensors at x=30*L0.

Usage:
    conda run -n pycuda_3_10 python src/OID_analysis/scripts/collect_target_cylinder.py \
        --diameter 1.0 --device 3 --steps 500

Output: data/illusion/illusion_{diam}L/ (shared with illusion scene)
    sensors.npz, fields.npz, target.npz, target_harmonics.json
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np

_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)
_SRC = os.path.join(_REPO, "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

from LegacyCelerisLab import FlowField  # noqa: E402
from OID_analysis.utils.cfd_interface import (  # noqa: E402
    load_legacy_configs, get_velocity_field, save_vorticity_png, vorticity_from_ddf,
    analyze_harmonics,
)
from OID_analysis.configs import get_scene, data_dir_for_scene, LEGACY_CFG_DIR, FIFO_LEN  # noqa: E402

DATA_TYPE = np.float32
L0 = 20.0
CENTER_Y = (512 - 1) / 2.0


def collect(diameter: float, device_id: int, n_steps: int) -> str:
    # Find scene name for this diameter
    scene_name = None
    for cn in ["illusion_0.75L", "illusion_1.0L", "illusion_1.5L"]:
        c = get_scene(cn)
        if abs(c["target_diameter"] - diameter) < 0.01:
            scene_name = cn
            break
    if scene_name is None:
        raise ValueError(f"No scene for diameter={diameter}")

    cfg = get_scene(scene_name)
    u0 = cfg["u0"]
    si = cfg["sample_interval"]

    tgt_key = f"target_cylinder_{diameter}L"
    if diameter == int(diameter):
        tgt_key = f"target_cylinder_{diameter:.1f}L"
    out_dir = data_dir_for_scene(tgt_key)
    os.makedirs(out_dir, exist_ok=True)

    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]), velocity=float(u0))

    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    print(f"  Target cylinder {diameter}L: {ff.FIELD_SHAPE}")

    # Target cylinder at x=20*L0
    tgt_radius = diameter * L0
    ff.add_cylinder((20.0 * L0, CENTER_Y, 0.0), tgt_radius)
    print(f"  radius={tgt_radius}")

    # 3 sensors at x=30*L0
    for y_off in [2.0, 0.0, -2.0]:
        ff.add_sensor((30.0 * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0)

    n_obj = 4
    n_stab = int(4 * ff.FIELD_SHAPE[0] / u0)
    print(f"  Stabilising ({n_stab} steps)...")
    ff.run(n_stab, np.zeros(n_obj, dtype=DATA_TYPE))

    # Record target signals: obs[0:8] = cylinder force(2) + sensor(6)
    target_states = np.empty((0, 8), dtype=DATA_TYPE)
    for _ in range(FIFO_LEN):
        ff.run(si, np.zeros(n_obj, dtype=DATA_TYPE))
        target_states = np.vstack((target_states, ff.obs.copy()[0:8]))
    print(f"  Target recorded: {target_states.shape}")

    # Harmonics for force channels
    target_harmonics = analyze_harmonics(target_states, n_harmonics=5)
    harm_save = [{k: v for k, v in h.items()} for h in target_harmonics]
    with open(os.path.join(out_dir, "target_harmonics.json"), "w") as f:
        json.dump(harm_save, f, indent=2)
    np.savez(os.path.join(out_dir, "target.npz"), target_states=target_states)

    # Full field rollout
    ff.apply_ddf()
    sens_list, ux_list, uy_list = [], [], []
    for step in range(n_steps):
        ff.run(si, np.zeros(n_obj, dtype=DATA_TYPE))
        obs_slice = ff.obs.copy()[0:6]
        sens_list.append(obs_slice)
        ux, uy = get_velocity_field(ff, u0=u0)
        ux_list.append(ux)
        uy_list.append(uy)

    np.savez(os.path.join(out_dir, "sensors.npz"),
             sensors=np.array(sens_list, dtype=np.float32))
    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=np.stack(ux_list), uy=np.stack(uy_list))

    omega = vorticity_from_ddf(ff, u0=u0)
    save_vorticity_png(os.path.join(out_dir, "vorticity_target.png"),
                       omega, title=f"Target cylinder {diameter}L")

    with open(os.path.join(out_dir, "config.json"), "w") as f:
        json.dump({k: str(v) if not isinstance(v, (int, float, list, bool)) else v
                   for k, v in cfg.items()}, f, indent=2)

    del ff
    print(f"  Saved {n_steps} snapshots to {out_dir}")
    return out_dir


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--diameter", type=float, default=1.0,
                    help="Target cylinder diameter (0.75, 1.0, 1.5)")
    ap.add_argument("--device", type=int, default=3)
    ap.add_argument("--steps", type=int, default=500)
    args = ap.parse_args()

    t0 = time.time()
    out = collect(args.diameter, args.device, args.steps)
    print(f"Done in {time.time() - t0:.1f}s -> {out}")


if __name__ == "__main__":
    main()