"""Collect target/incoming vortex street from disturbance cylinder (no pinball).

Records the Karman vortex street from the upstream disturbance cylinder alone,
saving both telemetry (sensors + forces) and full velocity field snapshots.

Usage:
    conda run -n pycuda_3_10 python scripts/collect_karman_q_in.py --device 3

Output: data/karman_target/karman_q_in/
"""
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

from CCD_analysis.configs import get_scene, data_dir_for_scene, LEGACY_CFG_DIR, L0, CENTER_Y
from CCD_analysis.utils.cfd_interface import (
    load_legacy_configs, get_velocity_field,
)

DATA_TYPE = np.float32


def collect():
    ap = argparse.ArgumentParser(
        description="Collect incoming Karman vortex street (disturbance cylinder only)")
    ap.add_argument("--device", type=int, default=3, help="CUDA device ID")
    ap.add_argument("--steps", type=int, default=500,
                    help="Number of recording steps")
    args = ap.parse_args()

    scene_name = "karman_q_in"
    cfg = get_scene(scene_name)
    out_dir = data_dir_for_scene(scene_name)
    u0 = cfg["u0"]
    si = cfg["sample_interval"]
    print(f"Output: {out_dir}", flush=True)

    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=0.004)

    ff = FlowField(field_cfg, cuda_cfg, device_id=args.device)

    # Add disturbance cylinder
    ff.add_cylinder((10.0 * L0, CENTER_Y, 0.0), L0)
    # Add 3 sensors at x=40*L0
    for y_off in [2.0, 0.0, -2.0]:
        ff.add_sensor((40.0 * L0, CENTER_Y + y_off * L0, 0.0), L0 / 4.0)

    n_obj = 4
    stabilize = int(4 * 1280 / u0)
    print(f"Stabilising ({stabilize} steps)...", flush=True)
    ff.run(stabilize, np.zeros(n_obj, dtype=DATA_TYPE))
    print("Stabilisation done.", flush=True)

    # Record 500 steps
    sens_list, forc_list = [], []
    ux_list, uy_list = [], []

    for step in range(args.steps):
        ff.run(si, np.zeros(n_obj, dtype=DATA_TYPE))
        obs = ff.obs.copy()
        # obs[0:2] = dist cylinder forces, obs[2:8] = 3 sensors
        sens_list.append(obs[2:8])
        forc_list.append(obs[0:2])
        ux, uy = get_velocity_field(ff, u0=u0)
        ux_list.append(ux)
        uy_list.append(uy)

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

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

    del ff
    print(f"Done, saved {args.steps} steps to {out_dir}", flush=True)


if __name__ == "__main__":
    t0 = time.time()
    collect()
    print(f"Time: {time.time() - t0:.1f}s", flush=True)
