"""Collect vortex cloak data for CCD correction-field analysis.

Collects field snapshots for three scene types per vortex type:
  - vortex_target_{type}: vortex only (no pinball), the "ideal" target flow
  - vortex_uncontrolled_{type}: vortex + pinball, zero control (q_blk)
  - vortex_{type}: vortex + pinball + PPO control (q_ctl)

All use LegacyCelerisLab (matching existing CCD data convention).

Usage:
    conda run -n pycuda_3_10 python scripts/collect_vortex.py \\
        --type lamb --device 2
    conda run -n pycuda_3_10 python scripts/collect_vortex.py \\
        --type taylor --device 2
    conda run -n pycuda_3_10 python scripts/collect_vortex.py \\
        --type all --device 2

--- BUG HISTORY (2026-06-29) ---
Three independent bugs in this script were discovered and fixed after
Lamb dipole showed non-physical front cylinder oscillation:

BUG 1 - Cylinder order swap (lines 166-168, 254-257):
  Original training env adds cylinders: front(id3) -> TOP(+y, id4) -> BOTTOM(-y, id5).
  action = [aF, aT(+bias), aB(-bias)] -> temp[3]=aF, temp[4]=aT(+bias), temp[5]=aB(-bias).
  We had TOP and BOTTOM swapped, so bias -4 went to the wrong cylinder.
  Consequence: rear cylinders rotated in opposite directions; front overcompensated.

BUG 2 - Observation normalization swap (lines 368-371):
  Training env produces obs = [forces/force_norm, sensors/sens_norm] (force-first).
  We incorrectly fed [sensors/force_norm, forces/sens_norm] (channel-swapped + wrong norms).
  Consequence: model received garbage feedback; on Lamb, similarity appeared 0.94 from
  the similarity computation but cross-correlation was only 0.73.

BUG 3 - Missing fade-in/out transitions (lines 314-357):
  uni_test.ipynb uses 25-step fade-in from steady-cloak bias [-5.1U0, +5.1U0] to PPO action,
  and 25-step fade-out back to steady-cloak. We applied PPO action immediately at full scale.
  Consequence: flow instability from abrupt control changes.

Fix summary:
  - Cylinder order: TOP(+y) at id4, BOTTOM(-y) at id5.
  - Bias: [-5, +5] U0 for both Lamb and Taylor (matching uni_test).
  - Observation: forces_norm=obs[6:12]/force_norm_fact, sens_norm=(obs[0:6]-sens_dev)/sens_norm_fact.
    Assembled as hstack([forces_norm, sens_norm]) — force first.
  - Fade-in/out: 25 steps linear interpolation between steady_bias and PPO action.

Comparison with other collect scripts:
  collect_karman.py and collect_illusion.py use build_observation() from cfd_interface.py
  which does force-first correctly. Only this script built obs manually.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from collections import deque

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, model_path_for_scene,
    LEGACY_CFG_DIR, L0, CENTER_Y, U0,
)
from CCD_analysis.utils.cfd_interface import (
    load_legacy_configs, get_velocity_field,
    load_ppo_model,
)

DATA_TYPE = np.float32
FIFO_LEN = 150


# ---------------------------------------------------------------------------
# Vortex configuration
# ---------------------------------------------------------------------------

_VORTEX_CFG = {
    "lamb": {"vortex_type": "lamb", "vortex_strength": 0.5 * U0},
    "taylor": {"vortex_type": "taylor", "vortex_strength": 0.03 * U0},
}


# ---------------------------------------------------------------------------
# Target: vortex only (no pinball), save fields.npz
# ---------------------------------------------------------------------------

def collect_target(vtype: str, device_id: int, out_dir: str,
                   n_steps: int = 150) -> dict:
    """Collect field snapshots for vortex-only target flow.

    Records the vortex evolving through a sensor-only environment
    at x=40*L0. Saves fields.npz and sensors.npz.
    """
    cfg = _VORTEX_CFG[vtype]
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(0.004))

    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    ny = ff.FIELD_SHAPE[1]
    n_sensors = 3

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

    # Short stabilization (1xNX/U0 for vortex, not 4x)
    stabilize_steps = int(1 * ff.FIELD_SHAPE[0] / U0)
    ff.run(stabilize_steps, np.zeros(n_sensors, dtype=DATA_TYPE))

    # Save clean flow DDF
    ff.get_ddf()
    ff.save_ddf()

    # Add vortex at x=10*L0
    vc = cfg["vortex_type"]
    vs = cfg["vortex_strength"]
    ff.add_vortex((10.0 * L0, CENTER_Y, 0.0),
                  2.0 * L0, vs, 0, vc)

    # Record vortex evolution
    sens_list, forc_list = [], []
    ux_list, uy_list = [], []

    for step in range(n_steps):
        ff.run(800, np.zeros(n_sensors, dtype=DATA_TYPE))
        obs = ff.obs.copy()  # (n_sensors*2,) = 6 sensor channels
        sens_list.append(obs)
        forc_list.append(np.zeros(6, dtype=DATA_TYPE))  # placeholder
        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))

    del ff
    return {"scene": f"vortex_target_{vtype}", "n_steps": n_steps}


# ---------------------------------------------------------------------------
# Uncontrolled: vortex + pinball, zero control (q_blk)
# ---------------------------------------------------------------------------

def collect_uncontrolled(vtype: str, device_id: int, out_dir: str,
                         n_steps: int = 150) -> dict:
    """Collect field snapshots for vortex + pinball with zero control.

    Records the transient interaction of vortex with the pinball.
    Target phases are recorded alongside field snapshots.
    """
    cfg = _VORTEX_CFG[vtype]
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(0.004))

    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    ny = ff.FIELD_SHAPE[1]
    n_sensors = 3

    # ---- Phase 1: Sensors + target recording ----
    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)

    stabilize_steps = int(1 * ff.FIELD_SHAPE[0] / U0)
    ff.run(stabilize_steps, np.zeros(n_sensors, dtype=DATA_TYPE))

    # Save clean DDF (pre-pinball, pre-vortex)
    ff.get_ddf()
    ff.save_ddf()

    # Record target (vortex only, for similarity reference)
    ff.add_vortex((10.0 * L0, CENTER_Y, 0.0),
                  2.0 * L0, cfg["vortex_strength"], 0, cfg["vortex_type"])
    target_states = np.empty((0, 6), dtype=DATA_TYPE)
    for _ in range(min(FIFO_LEN, n_steps)):
        ff.run(800, np.zeros(n_sensors, dtype=DATA_TYPE))
        target_states = np.vstack((target_states, ff.obs.copy()))
    np.savez(os.path.join(out_dir, "target.npz"), target_states=target_states)

    # ---- Phase 2: Add pinball, record uncontrolled flow ----
    ff.restore_ddf()
    ff.apply_ddf()

    # BUG-FIX (2026-06-29): cylinder order MUST match training env.
    # Original env adds: front(3) -> TOP(+y,id4) -> BOTTOM(-y,id5).
    # Swapping TOP/BOTTOM causes rear cylinders to rotate wrong direction.
    ff.add_cylinder((30.0 * L0, CENTER_Y, 0.0), L0 / 2.0)              # id 3: front
    ff.add_cylinder((31.3 * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0)   # id 4: TOP (+y)
    ff.add_cylinder((31.3 * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0)   # id 5: BOTTOM (-y)

    n_obj = ff.obs.size // 2
    assert n_obj == 6, f"Expected 6, got {n_obj}"

    # Bias action stabilization (matching uni_test: [-5, +5] for both Lamb and Taylor)
    ff.run(stabilize_steps, np.zeros(n_obj, dtype=DATA_TYPE))
    ff.run(stabilize_steps, np.array([0.0, 0.0, 0.0, 0.0, -5.0 * U0, 5.0 * U0], dtype=DATA_TYPE))

    # Add vortex at x=15*L0
    ff.add_vortex((15.0 * L0, CENTER_Y, 0.0),
                  2.0 * L0, cfg["vortex_strength"], 0, cfg["vortex_type"])

    # Record uncontrolled flow (zero actions on pinball)
    sens_list, forc_list = [], []
    ux_list, uy_list = [], []

    for step in range(n_steps):
        # Zero control on pinball
        ff.run(800, np.zeros(n_obj, dtype=DATA_TYPE))
        obs = ff.obs.copy()
        sens_list.append(obs[0:6])
        forc_list.append(obs[6:12])
        ux, uy = get_velocity_field(ff, u0=U0)
        ux_list.append(ux)
        uy_list.append(uy)

    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))

    del ff
    return {"scene": f"vortex_uncontrolled_{vtype}", "n_steps": n_steps}


# ---------------------------------------------------------------------------
# Controlled: vortex + pinball + PPO (q_ctl)
# ---------------------------------------------------------------------------

def collect_controlled(vtype: str, device_id: int, out_dir: str,
                       n_steps: int = 150) -> dict:
    """Collect field snapshots for vortex + pinball with PPO control.

    Loads the trained PPO model, runs inference, saves fields and telemetry.
    Also saves norm.json and ddf/fifo checkpoints for possible replay.
    """
    scene_name = f"vortex_{vtype}"
    cfg_src = get_scene(scene_name)
    cfg_v = _VORTEX_CFG[vtype]
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(0.004))

    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    ny = ff.FIELD_SHAPE[1]
    n_sensors = 3

    # Save config
    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_src.items()}, f, indent=2)

    # ---- Phase 1: Sensor-only target recording with vortex ----
    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)

    stabilize_steps_short = int(1 * ff.FIELD_SHAPE[0] / U0)
    ff.run(stabilize_steps_short, np.zeros(n_sensors, dtype=DATA_TYPE))

    ff.get_ddf()
    ff.save_ddf()

    ff.add_vortex((10.0 * L0, CENTER_Y, 0.0),
                  2.0 * L0, cfg_v["vortex_strength"], 0, cfg_v["vortex_type"])

    target_states = np.empty((0, 6), dtype=DATA_TYPE)
    for _ in range(min(FIFO_LEN, n_steps)):
        ff.run(800, np.zeros(n_sensors, dtype=DATA_TYPE))
        target_states = np.vstack((target_states, ff.obs.copy()))
    np.savez(os.path.join(out_dir, "target.npz"), target_states=target_states)

    # ---- Phase 2: Add pinball, compute norm ----
    ff.restore_ddf()
    ff.apply_ddf()

    # Object order MUST match training env: front(id3), TOP(+y,id4), BOTTOM(-y,id5)
    ff.add_cylinder((30.0 * L0, CENTER_Y, 0.0), L0 / 2.0)              # id 3: front
    ff.add_cylinder((31.3 * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0)   # id 4: TOP (+y)
    ff.add_cylinder((31.3 * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0)   # id 5: BOTTOM (-y)

    n_obj = ff.obs.size // 2
    assert n_obj == 6, f"Expected 6, got {n_obj}"

    # Stabilize with zero action, then bias action matching training env
    ff.run(stabilize_steps_short, np.zeros(n_obj, dtype=DATA_TYPE))
    ff.run(stabilize_steps_short, np.array([0.0, 0.0, 0.0, 0.0, -5.0 * U0, 5.0 * U0], dtype=DATA_TYPE))

    # Add vortex at x=15*L0
    ff.add_vortex((15.0 * L0, CENTER_Y, 0.0),
                  2.0 * L0, cfg_v["vortex_strength"], 0, cfg_v["vortex_type"])

    # Save DDF checkpoint (vortex at x=15, pinball stabilized with bias)
    ff.get_ddf()
    ff.save_ddf()

    # Norm collection (zero action on vortex+pinball)
    fifo = deque(maxlen=FIFO_LEN)
    for _ in range(FIFO_LEN):
        ff.run(800, np.zeros(n_obj, dtype=DATA_TYPE))
        fifo.append(ff.obs.copy())

    temp_states = np.array(fifo, dtype=DATA_TYPE)
    force_norm_fact = 6.0 * float(np.max(np.abs(temp_states[:, 6:12])))
    sens_deviation = np.mean(temp_states[:, 0:6], axis=0).astype(DATA_TYPE)
    sens_norm_fact = np.zeros(6, dtype=DATA_TYPE)
    for i in range(6):
        sens_norm_fact[i] = 5.0 * float(np.max(np.abs(temp_states[:, i] - sens_deviation[i])))

    norm = {
        "force_norm_fact": force_norm_fact,
        "sens_deviation": sens_deviation.tolist(),
        "sens_norm_fact": sens_norm_fact.tolist(),
        "action_bias": [0.0, -4.0, 4.0],
    }
    with open(os.path.join(out_dir, "norm.json"), "w") as f:
        json.dump(norm, f, indent=2)

    # Bias FIFO init (restore DDF so vortex starts from x=15)
    ff.restore_ddf()
    ff.apply_ddf()
    fifo.clear()
    # BUG-FIX (2026-06-29): use uni_test bias [-5,+5], not training env bias [-4,+4].
    # The model was tested with [-5,+5] in uni_test.ipynb.
    bias_fifo_arr = np.array([0.0, 0.0, 0.0, 0.0, -5.0 * U0, 5.0 * U0], dtype=DATA_TYPE)
    for _ in range(FIFO_LEN):
        ff.run(800, bias_fifo_arr)
        fifo.append(ff.obs.copy())

    save_states = np.array(list(fifo), dtype=DATA_TYPE)
    # Restore DDF back to vortex-at-x=15 checkpoint
    ff.restore_ddf()
    ff.apply_ddf()

    # Save checkpoints for replay (vortex at initial position)
    np.save(os.path.join(out_dir, "ddf_checkpoint.npy"), ff.ddf)
    np.save(os.path.join(out_dir, "fifo_checkpoint.npy"), save_states)

    # ---- Phase 3: Controlled PPO inference (fade-in/out matching uni_test) ----
    model_path = model_path_for_scene(scene_name)
    if model_path is None:
        raise FileNotFoundError(f"No model found for scene: {scene_name}")

    model = load_ppo_model(model_path, device=f"cuda:{device_id}", s_dim=12)
    model.set_random_seed(0)

    # Restore DDF to vortex-at-x=15 checkpoint
    ff.restore_ddf()
    ff.apply_ddf()

    # Start with zeros observation (matching uni_test)
    obs = np.zeros(12, dtype=np.float32)

    sens_list, forc_list, act_list, rew_list = [], [], [], []
    ux_list = []
    uy_list = []
    fifo = deque(maxlen=FIFO_LEN)

    # Steady-cloak bias for transition (matching uni_test: [-5.1, +5.1])
    # BUG-FIX (2026-06-29): uni_test uses [-5.1,+5.1], not training env [-4,+4].
    steady_bias = np.array([0.0, -5.1 * U0, 5.1 * U0], dtype=DATA_TYPE)

    for step in range(n_steps):
        action, _ = model.predict(obs, deterministic=True)
        action = action.astype(np.float32).flatten()
        act_list.append(action.copy())

        temp_action = np.array(action * 4.0 + np.array([0.0, -4.0, 4.0]), dtype=DATA_TYPE)

        # Fade-in (0-24), active (25-44), fade-out (45-69), steady-cloak (70+)
        if step < 25:
            w = step / 25.0
            temp_val = temp_action * w * U0 + steady_bias * (1.0 - w)
        elif 45 <= step < 70:
            w = (step - 45) / 25.0
            temp_val = temp_action * (1.0 - w) * U0 + steady_bias * w
        elif step >= 70:
            temp_val = steady_bias
        else:
            temp_val = temp_action * U0

        temp = np.zeros(n_obj, dtype=DATA_TYPE)
        temp[3:6] = temp_val

        ff.context.push()
        ff.run(800, temp)
        ff.context.pop()

        obs_slice = ff.obs.copy()
        fifo.append(obs_slice)
        sens_list.append(obs_slice[0:6])
        forc_list.append(obs_slice[6:12])

        # Build observation for next step
        # BUG-FIX (2026-06-29): force-first with CORRECT norms.
        #   OLD (broken): sens_raw = obs[0:6]/force_norm, force_raw = (obs[6:12]-sens_dev)/sens_norm
        #                  obs = hstack([sens_raw, force_raw])  ← completely wrong!
        #   NEW (correct): forces_norm = obs[6:12]/force_norm, sens_norm = (obs[0:6]-sens_dev)/sens_norm
        #                  obs = hstack([forces_norm, sens_norm])  ← force first
        # Training env produces: obs = hstack([forces/force_norm, sensors/sens_norm]).
        # collect_karman and collect_illusion use build_observation() which does this correctly.
        forces_norm = obs_slice[6:12] / force_norm_fact
        sens_norm = (obs_slice[0:6] - sens_deviation) / sens_norm_fact
        obs = np.clip(np.hstack([forces_norm, sens_norm]), -1.0, 1.0).astype(np.float32)

        # Save field snapshot
        ux, uy = get_velocity_field(ff, u0=U0)
        ux_list.append(ux)
        uy_list.append(uy)

    # Save field snapshots
    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=np.stack(ux_list), uy=np.stack(uy_list))

    # Save telemetry
    np.savez(os.path.join(out_dir, "controlled.npz"),
             sensors=np.array(sens_list, dtype=np.float32),
             forces=np.array(forc_list, dtype=np.float32),
             actions=np.array(act_list, dtype=np.float32))

    # Compute similarity
    from CCD_analysis.utils.cfd_interface import compute_similarity
    states_arr = np.array(sens_list, dtype=np.float32)
    n_align = min(states_arr.shape[0], target_states.shape[0])
    if n_align >= 30:
        sim = compute_similarity(target_states, states_arr[:n_align], 30)
    else:
        sim = 0.0

    result = {"scene": scene_name, "similarity": float(sim), "n_steps": n_steps}
    with open(os.path.join(out_dir, "result.json"), "w") as f:
        json.dump(result, f, indent=2)

    del ff, model
    return result


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    ap = argparse.ArgumentParser(description="Collect vortex cloak fields for CCD")
    ap.add_argument("--type", type=str, default="lamb",
                    help='Vortex type: lamb, taylor, or "all"')
    ap.add_argument("--device", type=int, default=2, help="GPU device ID")
    ap.add_argument("--steps", type=int, default=150,
                    help="Number of recording steps (max 150 for transient)")
    ap.add_argument("--skip-target", action="store_true",
                    help="Skip vortex_target collection")
    ap.add_argument("--skip-uncontrolled", action="store_true",
                    help="Skip uncontrolled collection")
    ap.add_argument("--skip-controlled", action="store_true",
                    help="Skip controlled collection")
    args = ap.parse_args()

    if args.type.lower() == "all":
        vtypes = ["lamb", "taylor"]
    else:
        vtypes = [args.type.lower()]

    t_start = time.time()
    for vtype in vtypes:
        print(f"\n{'=' * 60}")
        print(f"Vortex type: {vtype}")
        print(f"{'=' * 60}")

        # --- Target: vortex only ---
        if not args.skip_target:
            scene_name = f"vortex_target_{vtype}"
            print(f"\n--- Collecting target: {scene_name} ---")
            out_dir = data_dir_for_scene(scene_name)
            r = collect_target(vtype, args.device, out_dir, args.steps)
            print(f"  Done: {r['scene']} -> {out_dir}")

        # --- Uncontrolled: vortex + pinball, zero control ---
        if not args.skip_uncontrolled:
            scene_name = f"vortex_uncontrolled_{vtype}"
            print(f"\n--- Collecting uncontrolled: {scene_name} ---")
            out_dir = data_dir_for_scene(scene_name)
            r = collect_uncontrolled(vtype, args.device, out_dir, args.steps)
            print(f"  Done: {r['scene']} -> {out_dir}")

        # --- Controlled: vortex + pinball + PPO ---
        if not args.skip_controlled:
            scene_name = f"vortex_{vtype}"
            print(f"\n--- Collecting controlled: {scene_name} ---")
            out_dir = data_dir_for_scene(scene_name)
            r = collect_controlled(vtype, args.device, out_dir, args.steps)
            print(f"  Done: {r['scene']} -> {out_dir}  sim={r['similarity']:.4f}")

    elapsed = time.time() - t_start
    print(f"\nTotal time: {elapsed:.1f}s")


if __name__ == "__main__":
    sys.exit(main())
