# CCD_analysis/scripts/phase1_collect.py
"""Phase 1: Data collection for all 4 analysis cases.

Usage::
    conda run -n pycuda_3_10 python phase1_collect.py --case illusion --device 2
    conda run -n pycuda_3_10 python phase1_collect.py --case cloak --device 3
    conda run -n pycuda_3_10 python phase1_collect.py --case uncontrolled --device 3
    conda run -n pycuda_3_10 python phase1_collect.py --case target_cylinder --device 2
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from collections import deque
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

# Add project root
_REPO = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
if _REPO not in sys.path:
    sys.path.insert(0, _REPO)

# Add analysis dir
_ANALYSIS = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if _ANALYSIS not in sys.path:
    sys.path.insert(0, _ANALYSIS)

from LegacyCelerisLab import FlowField  # noqa: E402

from scripts.cfg import (  # noqa: E402
    CONFIG_DIR, OUTPUT_DIR, U0, L0, NX, NY, CENTER_Y, DATA_TYPE,
    PINBALL_RADIUS, FRONT_CENTER, BOTTOM_CENTER, TOP_CENTER,
    ILLUSION_FRONT, ILLUSION_BOTTOM, ILLUSION_TOP,
    SENSOR_RADIUS, SENSOR_CENTERS_CLOAK, SENSOR_CENTERS_ILLUSION,
    TARGET_CYLINDER_CENTER, TARGET_CYLINDER_RADIUS,
    SAMPLE_INTERVAL, SAMPLE_INTERVAL_ILLUSION,
    ACTION_SCALE_CLOAK, ACTION_BIAS_CLOAK,
    ACTION_SCALE_ILLUSION, ACTION_BIAS_ILLUSION,
    MODEL_CLOAK_RE100, MODEL_ILLUSION_1L,
    STABILIZE_STEPS, FIFO_LEN, N_PTS_PER_CYCLE,
    nu_from_re,
)
from scripts.utils import (  # noqa: E402
    load_configs, get_velocity_field, detect_cycle_stability,
)

# ---------------------------------------------------------------------------
# PPO model loader (with Sin activation)
# ---------------------------------------------------------------------------

def _load_ppo_model(model_path: str, device: str, s_dim: int = 12, a_dim: int = 3):
    """Load PPO model with Sin activation."""
    import torch
    from torch.nn import Module
    from stable_baselines3 import PPO
    import gymnasium as gym
    from gymnasium import spaces

    class Sin(Module):
        def forward(self, x):
            return torch.sin(x)

    class DummyEnv(gym.Env):
        def __init__(self):
            super().__init__()
            self.observation_space = spaces.Box(
                low=-1, high=1, shape=(s_dim,), dtype=np.float32)
            self.action_space = spaces.Box(
                low=-1, high=1, shape=(a_dim,), dtype=np.float32)
        def reset(self, seed=None):
            return np.zeros(s_dim, dtype=np.float32), {}
        def step(self, action):
            return np.zeros(s_dim, dtype=np.float32), 0.0, False, False, {}
        def render(self):
            pass

    dummy = DummyEnv()
    model = PPO.load(model_path, env=dummy, device=device)
    return model


# ---------------------------------------------------------------------------
# Field saving interval calculator
# ---------------------------------------------------------------------------

def _calc_save_interval(T_ref: float, n_pts_per_cycle: int = 24) -> int:
    """Calculate field save interval to get ~n_pts_per_cycle per cycle."""
    interval = int(T_ref / n_pts_per_cycle)
    return max(1, interval)


# ---------------------------------------------------------------------------
# Phase 1a: Illusion
# ---------------------------------------------------------------------------

def collect_illusion(device_id: int, data: dict) -> dict:
    """Collect illusion case data with proper norm computation and PPO inference.

    Follows legacy_env_imit.py __init__ + step() logic exactly:
        1. Target cylinder recording (separate FlowField)
        2. FFT harmonics on target signals
        3. Pinball env with norm computation
        4. Bias-action FIFO initialization
        5. PPO deterministic rollout with 14-dim normalized observations
    """
    actual_U0 = 0.02  # model is 2U
    viscosity = nu_from_re(100.0, u0=actual_U0)
    sample_interval = SAMPLE_INTERVAL_ILLUSION  # 600
    fifo_len = 150
    conv_len = 36

    # ---- Step 1: Target cylinder recording ----
    print("--- Record target cylinder ---")
    target_U0 = actual_U0
    target_nu = viscosity
    cuda_cfg, field_cfg = load_configs(CONFIG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(target_nu),
                                   velocity=float(target_U0))
    ff_target = FlowField(field_cfg, cuda_cfg, device_id=device_id)

    # Target cylinder: center=(20*L0, CENTER_Y), radius=1.0*L0
    L0 = 20.0
    ff_target.add_cylinder(
        (20.0 * L0, (512 - 1) / 2, 0.0), 1.0 * L0
    )
    # 3 sensors at x=30*L0
    for y_off in [2.0, 0.0, -2.0]:
        ff_target.add_sensor(
            (30.0 * L0, (512 - 1) / 2 + y_off * L0, 0.0), L0 / 4.0
        )
    n_obj_target = ff_target.obs.size // 2  # 4
    # Stabilize
    ff_target.run(int(4 * 1280 / target_U0), np.zeros(n_obj_target, dtype=np.float32))

    # Record 150 steps of obs[0:8] (3 sensors + 1 cylinder force)
    target_states = np.empty((0, 8), dtype=np.float32)
    for _ in range(fifo_len):
        ff_target.run(sample_interval, np.zeros(n_obj_target, dtype=np.float32))
        new_state = ff_target.obs.copy()[0:8]
        target_states = np.vstack((target_states, new_state))

    # FFT harmonics analysis
    def analyze_harmonics(states, n_harmonics=5):
        N, D = states.shape
        result = []
        for d in range(D):
            y = states[:, d]
            fft_coef = np.fft.rfft(y)
            freqs = np.fft.rfftfreq(N, d=1)
            amps = 2.0 * np.abs(fft_coef) / N
            phases = np.angle(fft_coef)
            idx = np.argsort(amps[1:])[::-1][:n_harmonics] + 1
            harmonics = {
                'dc': float(np.real(fft_coef[0]) / N),
                'amps': amps[idx].tolist(),
                'freqs': freqs[idx].tolist(),
                'phases': phases[idx].tolist(),
            }
            result.append(harmonics)
        return result

    target_harmonics = analyze_harmonics(target_states, n_harmonics=5)
    del ff_target
    print(f"  target harmonics computed for {len(target_harmonics)} channels")

    # ---- Step 2: Pinball env creation ----
    print("--- Build pinball env ---")
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)

    for y_off in [2.0, 0.0, -2.0]:
        ff.add_sensor(
            (30.0 * L0, (512 - 1) / 2 + y_off * L0, 0.0), L0 / 4.0
        )
    ff.add_cylinder((19.0 * L0, (512 - 1) / 2, 0.0), L0 / 2.0)
    ff.add_cylinder((20.3 * L0, (512 - 1) / 2 + 0.75 * L0, 0.0), L0 / 2.0)
    ff.add_cylinder((20.3 * L0, (512 - 1) / 2 - 0.75 * L0, 0.0), L0 / 2.0)

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

    # Stabilize
    ff.run(int(4 * 1280 / actual_U0), np.zeros(n_obj, dtype=np.float32))
    ff.get_ddf()
    ff.save_ddf()  # checkpoint

    # ---- Step 3: Norm computation (zero-action rollout) ----
    print("--- Compute norm ---")
    fifo = deque(maxlen=fifo_len)
    for _ in range(fifo_len):
        ff.run(sample_interval, np.zeros(n_obj, dtype=np.float32))
        fifo.append(ff.obs.copy()[0:12])

    temp_states = np.array(fifo, dtype=np.float32)
    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(np.float32)
    sens_norm_fact = np.zeros(6, dtype=np.float32)
    for i in range(6):
        sens_norm_fact[i] = 5.0 * float(np.max(np.abs(temp_states[:, i] - sens_deviation[i])))

    print(f"  force_norm_fact={force_norm_fact:.6f}")
    print(f"  sens_deviation={sens_deviation}")
    print(f"  sens_norm_fact={sens_norm_fact}")

    # ---- Step 4: Bias-action FIFO initialization ----
    print("--- Bias-action FIFO init ---")
    ff.apply_ddf()
    # bias action from legacy env: [0, 0, 0, 0, -1*U0, 1*U0]
    bias_arr = np.zeros(n_obj, dtype=np.float32)
    bias_arr[4] = -1.0 * actual_U0  # bottom
    bias_arr[5] = 1.0 * actual_U0   # top

    fifo.clear()
    for _ in range(fifo_len):
        ff.run(sample_interval, bias_arr)
        fifo.append(ff.obs.copy()[0:12])

    save_states = list(fifo)
    ff.apply_ddf()  # restore checkpoint for reset

    # ---- Step 5: PPO inference with adaptive sampling ----
    print("--- PPO deterministic rollout (adaptive sampling) ---")
    import torch
    device_str = f"cuda:{device_id}" if torch.cuda.is_available() else "cpu"
    model = _load_ppo_model(MODEL_ILLUSION_1L, device=device_str, s_dim=14, a_dim=3)
    model.set_random_seed(19)

    n_steps = 200

    # Compute adaptive field sampling interval from expected period
    # St = 0.267, D = 40, expected f = St * U0 / D
    f_expected = 0.2667 * actual_U0 / 40.0
    T_expected = int(1.0 / f_expected) if f_expected > 0 else 7500
    field_interval = max(1, int(T_expected / N_PTS_PER_CYCLE))
    print(f"  T_expected={T_expected} steps, field_interval={field_interval} "
          f"(~{T_expected/field_interval:.0f} pts/cycle)")

    # Data at PPO-action cadence (once per 600 steps, for PPO state only)
    ppo_actions = []
    ppo_sensors_600 = []

    # Dense data at field_interval cadence (for phase analysis)
    dense_sensors = []
    dense_forces = []
    dense_ux = []
    dense_uy = []

    # Re-initialize FIFO for inference
    fifo = deque(maxlen=fifo_len)
    for state in save_states:
        fifo.append(np.array(state, dtype=np.float32))

    obs = np.zeros(14, dtype=np.float32)

    for step in range(n_steps):
        # PPO action
        action, _states = model.predict(obs, deterministic=True)
        action = action.astype(np.float32).flatten()
        ppo_actions.append(action.copy())

        # Convert to physical omega
        temp = np.zeros(n_obj, dtype=np.float32)
        omega = (action * ACTION_SCALE_ILLUSION
                 + np.array(ACTION_BIAS_ILLUSION, dtype=np.float32)) * actual_U0
        temp[3:6] = omega

        # Run CFD with dense intra-step sampling
        ff.context.push()
        try:
            # First chunk
            ff.run(field_interval, temp)
            ux, uy = get_velocity_field(ff, u0=actual_U0)
            dense_ux.append(ux)
            dense_uy.append(uy)
            dense_sensors.append(ff.obs.copy()[0:6])
            dense_forces.append(ff.obs.copy()[6:12])

            # Second chunk (remaining)
            remaining = sample_interval - field_interval
            if remaining > 0:
                ff.run(remaining, temp)
                ux, uy = get_velocity_field(ff, u0=actual_U0)
                dense_ux.append(ux)
                dense_uy.append(uy)
                dense_sensors.append(ff.obs.copy()[0:6])
                dense_forces.append(ff.obs.copy()[6:12])
        finally:
            ff.context.pop()

        # PPO state: use last obs_slice
        last_sens = dense_sensors[-1]
        last_force = dense_forces[-1]
        obs_slice = np.concatenate([last_sens, last_force])
        fifo.append(obs_slice)
        ppo_sensors_600.append(obs_slice)

        # Build normalized 14-dim observation for next PPO step
        forces_norm = last_force / force_norm_fact
        sens_norm = (last_sens - sens_deviation) / sens_norm_fact
        target_recon = _gen_target_states_at(step, target_harmonics)
        target_cd_norm = float(target_recon[0]) / force_norm_fact
        target_cl_norm = float(target_recon[1]) / force_norm_fact
        obs = np.clip(
            np.hstack([forces_norm, sens_norm, target_cd_norm, target_cl_norm]),
            -1.0, 1.0,
        ).astype(np.float32)

        if step % 20 == 0:
            print(f"  step {step}/{n_steps}, action={action[0]:.3f} {action[1]:.3f} {action[2]:.3f}")

    # Save dense data (for phase resampling)
    ux_all = np.stack(dense_ux, axis=0)
    uy_all = np.stack(dense_uy, axis=0)
    dense_sensors_arr = np.array(dense_sensors, dtype=np.float32)
    dense_forces_arr = np.array(dense_forces, dtype=np.float32)
    ppo_actions_arr = np.array(ppo_actions, dtype=np.float32)
    n_dense_per_step = len(dense_sensors) // n_steps
    dense_dt = sample_interval / n_dense_per_step if n_dense_per_step > 0 else sample_interval
    print(f"  Dense sampling: {len(dense_sensors)} samples, "
          f"{n_dense_per_step} per PPO step, dt={dense_dt:.0f} LBM steps")

    out_dir = os.path.join(OUTPUT_DIR, "illusion")
    os.makedirs(out_dir, exist_ok=True)
    np.savez_compressed(os.path.join(out_dir, "fields.npz"), ux=ux_all, uy=uy_all)
    np.savez(os.path.join(out_dir, "dense_sensors.npz"),
             sensors=dense_sensors_arr, forces=dense_forces_arr,
             dense_dt=dense_dt,
             sample_interval=sample_interval)

    # Save PPO-step-cadence data and metadata
    np.savez(os.path.join(out_dir, "sensors.npz"),
             sensors=dense_sensors_arr.reshape(n_steps, -1, 6)[:, -1],
             forces=dense_forces_arr.reshape(n_steps, -1, 6)[:, -1],
             actions=ppo_actions_arr,
             sample_interval=sample_interval,
             force_norm_fact=np.array([force_norm_fact], dtype=np.float32),
             sens_deviation=np.array(sens_deviation, dtype=np.float32),
             sens_norm_fact=np.array(sens_norm_fact, dtype=np.float32))

    # Save target data for later use
    np.savez(os.path.join(out_dir, "target_harmonics.npz"),
             target_states=target_states,
             harmonics_data=np.array(target_harmonics, dtype=object))

    meta = {
        "case": "illusion",
        "model": str(MODEL_ILLUSION_1L),
        "n_steps": n_steps,
        "n_fields": len(dense_ux),
        "n_dense_samples": len(dense_sensors),
        "dense_dt": dense_dt,
        "T_expected": T_expected,
        "field_interval": field_interval,
        "sample_interval": sample_interval,
        "action_scale": ACTION_SCALE_ILLUSION,
        "action_bias": list(ACTION_BIAS_ILLUSION),
        "U0": actual_U0,
        "viscosity": viscosity,
        "n_obj": n_obj,
        "force_norm_fact": force_norm_fact,
        "sens_deviation": sens_deviation.tolist(),
        "sens_norm_fact": sens_norm_fact.tolist(),
    }
    with open(os.path.join(out_dir, "meta.json"), "w") as f:
        json.dump(meta, f, indent=2)

    print(f"  Saved {len(dense_ux)} fields, {len(dense_sensors)} dense samples")
    del ff, model
    return meta


def _gen_target_states_at(t, harmonics):
    """Reconstruct target observable at step index t from harmonics.
    
    Mirrors legacy_env_imit.py gen_target_states_at().
    """
    t = np.asarray(t)
    D = len(harmonics)
    result = np.zeros((t.size, D), dtype=np.float32)
    for d, h in enumerate(harmonics):
        val = np.full(t.shape, h['dc'], dtype=np.float32)
        amps = h['amps']
        freqs = h['freqs']
        phases = h['phases']
        for amp, freq, phase in zip(amps, freqs, phases):
            val += amp * np.cos(2 * np.pi * freq * t + phase)
        result[:, d] = val
    if result.shape[0] == 1:
        return result[0]
    return result


# ---------------------------------------------------------------------------
# Phase 1b: Cloak (steady flow case)
# ---------------------------------------------------------------------------

def collect_cloak(device_id: int, data: dict) -> dict:
    """Collect cloak case data (PPO -> steady action -> mean flow)."""
    viscosity = nu_from_re(100.0)
    sample_interval = SAMPLE_INTERVAL

    import torch
    device_str = f"cuda:{device_id}" if torch.cuda.is_available() else "cpu"
    model = _load_ppo_model(MODEL_CLOAK_RE100, device=device_str, s_dim=12, a_dim=3)
    model.set_random_seed(0)

    # Create env: 6 objects (3 sensors + 3 pinball, NO disturbance cylinder)
    cuda_cfg, field_cfg = load_configs(CONFIG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(viscosity))
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)

    for sc in SENSOR_CENTERS_CLOAK:
        ff.add_sensor((sc[0], sc[1], 0.0), SENSOR_RADIUS)
    ff.add_cylinder((FRONT_CENTER[0], FRONT_CENTER[1], 0.0), PINBALL_RADIUS)
    ff.add_cylinder((BOTTOM_CENTER[0], BOTTOM_CENTER[1], 0.0), PINBALL_RADIUS)
    ff.add_cylinder((TOP_CENTER[0], TOP_CENTER[1], 0.0), PINBALL_RADIUS)

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

    # Stabilize
    ff.run(STABILIZE_STEPS, np.zeros(n_obj, dtype=np.float32))

    # ---- PPO deterministic rollout to find steady action ----
    n_ppo_steps = 200
    print(f"Running {n_ppo_steps} PPO steps to extract steady action...")

    obs = np.zeros(12, dtype=np.float32)
    actions_list = []
    sensors_list = []
    forces_list = []

    for step in range(n_ppo_steps):
        action, _states = model.predict(obs, deterministic=True)
        action = action.astype(np.float32).flatten()
        actions_list.append(action.copy())

        temp = np.zeros(n_obj, dtype=np.float32)
        omega = (action * ACTION_SCALE_CLOAK
                 + np.array(ACTION_BIAS_CLOAK, dtype=np.float32)) * U0
        temp[3:6] = omega

        ff.context.push()
        try:
            ff.run(sample_interval, temp)
        finally:
            ff.context.pop()

        obs_slice = ff.obs.copy()[0:12]
        sensors_list.append(obs_slice[0:6].copy())
        forces_list.append(obs_slice[6:12].copy())

        # Build observation for next step
        obs = np.clip(np.hstack([obs_slice[6:12], obs_slice[0:6]]),
                      -10.0, 10.0).astype(np.float32)

    # Extract steady action (average of last 100 steps)
    actions_arr = np.array(actions_list, dtype=np.float32)
    steady_action = np.mean(actions_arr[-100:], axis=0)
    print(f"  Steady action ([-1,1]): {steady_action[0]:.4f} {steady_action[1]:.4f} {steady_action[2]:.4f}")
    print(f"  Steady omega (U0 multiples): "
          f"{(steady_action*ACTION_SCALE_CLOAK+np.array(ACTION_BIAS_CLOAK))[0]:.4f} "
          f"{(steady_action*ACTION_SCALE_CLOAK+np.array(ACTION_BIAS_CLOAK))[1]:.4f} "
          f"{(steady_action*ACTION_SCALE_CLOAK+np.array(ACTION_BIAS_CLOAK))[2]:.4f}")

    # ---- Apply steady action and record mean flow ----
    print("Applying steady action and recording...")
    temp_steady = np.zeros(n_obj, dtype=np.float32)
    omega_steady = (steady_action * ACTION_SCALE_CLOAK
                    + np.array(ACTION_BIAS_CLOAK, dtype=np.float32)) * U0
    temp_steady[3:6] = omega_steady

    # Re-stabilize with steady action (4x NX/U0)
    ff.context.push()
    try:
        ff.run(STABILIZE_STEPS, temp_steady)
    finally:
        ff.context.pop()

    # Record steady state fields and sensors
    n_steady_samples = 30
    steady_sensors = []
    steady_forces = []
    steady_ux = []
    steady_uy = []

    for i in range(n_steady_samples):
        ff.context.push()
        try:
            ff.run(sample_interval, temp_steady)
        finally:
            ff.context.pop()
        obs_slice = ff.obs.copy()[0:12]
        steady_sensors.append(obs_slice[0:6])
        steady_forces.append(obs_slice[6:12])
        ux, uy = get_velocity_field(ff, u0=U0)
        steady_ux.append(ux)
        steady_uy.append(uy)

    steady_sensors_arr = np.array(steady_sensors, dtype=np.float32)
    steady_forces_arr = np.array(steady_forces, dtype=np.float32)
    ux_all = np.stack(steady_ux, axis=0)
    uy_all = np.stack(steady_uy, axis=0)

    out_dir = os.path.join(OUTPUT_DIR, "cloak")
    os.makedirs(out_dir, exist_ok=True)

    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=ux_all, uy=uy_all)
    np.savez(os.path.join(out_dir, "sensors.npz"),
             sensors=steady_sensors_arr, forces=steady_forces_arr)
    np.savez(os.path.join(out_dir, "ppo_rollout.npz"),
             actions=actions_arr,
             sensors=np.array(sensors_list, dtype=np.float32),
             forces=np.array(forces_list, dtype=np.float32),
             steady_action=steady_action)

    meta = {
        "case": "cloak",
        "model": str(MODEL_CLOAK_RE100),
        "sample_interval": sample_interval,
        "action_scale": ACTION_SCALE_CLOAK,
        "action_bias": list(ACTION_BIAS_CLOAK),
        "steady_action_norm": steady_action.tolist(),
        "steady_omega_U0": (steady_action * ACTION_SCALE_CLOAK
                            + np.array(ACTION_BIAS_CLOAK)).tolist(),
        "U0": U0,
        "viscosity": viscosity,
        "n_obj": n_obj,
        "n_steady_samples": n_steady_samples,
    }
    with open(os.path.join(out_dir, "meta.json"), "w") as f:
        json.dump(meta, f, indent=2)

    print(f"  Steady action recorded. Mean sensors: "
          f"{np.mean(steady_sensors_arr, axis=0)}")
    print(f"  Mean total force: "
          f"Fx={np.mean(steady_forces_arr[:, 0::2]):.6f} "
          f"Fy={np.mean(steady_forces_arr[:, 1::2]):.6f}")
    del ff, model
    return meta


# ---------------------------------------------------------------------------
# Phase 1c: Uncontrolled
# ---------------------------------------------------------------------------

def collect_uncontrolled(device_id: int, data: dict) -> dict:
    """Collect uncontrolled case data (zero-action baseline)."""
    viscosity = nu_from_re(100.0)
    sample_interval = SAMPLE_INTERVAL
    T_ref = data.get("T_ref", 15000.0)
    save_interval = _calc_save_interval(T_ref)

    cuda_cfg, field_cfg = load_configs(CONFIG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(viscosity))
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)

    for sc in SENSOR_CENTERS_CLOAK:
        ff.add_sensor((sc[0], sc[1], 0.0), SENSOR_RADIUS)
    ff.add_cylinder((FRONT_CENTER[0], FRONT_CENTER[1], 0.0), PINBALL_RADIUS)
    ff.add_cylinder((BOTTOM_CENTER[0], BOTTOM_CENTER[1], 0.0), PINBALL_RADIUS)
    ff.add_cylinder((TOP_CENTER[0], TOP_CENTER[1], 0.0), PINBALL_RADIUS)

    n_obj = ff.obs.size // 2
    assert n_obj == 6

    # Stabilize
    ff.run(STABILIZE_STEPS, np.zeros(n_obj, dtype=np.float32))

    # Run uncontrolled
    n_steps = 200
    sensors_list = []
    forces_list = []
    ux_fields = []
    uy_fields = []

    for step in range(n_steps):
        ff.context.push()
        try:
            remaining = sample_interval
            while remaining > 0:
                chunk = min(remaining, save_interval)
                ff.run(chunk, np.zeros(n_obj, dtype=np.float32))
                remaining -= chunk
                ux, uy = get_velocity_field(ff, u0=U0)
                ux_fields.append(ux)
                uy_fields.append(uy)
        finally:
            ff.context.pop()

        obs_slice = ff.obs.copy()[0:12]
        sensors_list.append(obs_slice[0:6])
        forces_list.append(obs_slice[6:12])

    sensors = np.array(sensors_list, dtype=np.float32)
    forces = np.array(forces_list, dtype=np.float32)
    ux_all = np.stack(ux_fields, axis=0)
    uy_all = np.stack(uy_fields, axis=0)

    out_dir = os.path.join(OUTPUT_DIR, "uncontrolled")
    os.makedirs(out_dir, exist_ok=True)

    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=ux_all, uy=uy_all)
    np.savez(os.path.join(out_dir, "sensors.npz"),
             sensors=sensors, forces=forces)

    meta = {
        "case": "uncontrolled",
        "U0": U0,
        "viscosity": viscosity,
        "n_steps": n_steps,
        "n_fields": len(ux_fields),
        "sample_interval": sample_interval,
        "n_obj": n_obj,
    }
    with open(os.path.join(out_dir, "meta.json"), "w") as f:
        json.dump(meta, f, indent=2)

    print(f"  Saved {len(ux_fields)} fields, {len(sensors)} sensor steps")
    del ff
    return meta


# ---------------------------------------------------------------------------
# Phase 1d: Target cylinder (reference for period detection)
# ---------------------------------------------------------------------------

def collect_target_cylinder(device_id: int, data: dict) -> dict:
    """Collect target 2D cylinder reference data.

    Most data was already collected in Phase 0. Here we just ensure
    the fields are properly saved with the right naming.
    """
    # Phase 0 already saved data to output/target_cylinder/
    # Just verify it exists and copy meta
    out_dir = os.path.join(OUTPUT_DIR, "target_cylinder")
    meta_path = os.path.join(out_dir, "meta.json")
    if not os.path.exists(meta_path):
        raise RuntimeError(
            "Phase 0 must be run first. No target_cylinder data found."
        )

    with open(meta_path, "r") as f:
        meta = json.load(f)

    print(f"Target cylinder data found at {out_dir}")
    print(f"  f_ref={meta['f_ref']:.6f}, T_ref={meta['T_ref']:.0f}, St={meta['St']:.4f}")
    print(f"  CV_T={meta['CV_T']:.4f}")
    return meta


# ---------------------------------------------------------------------------
# Empty channel (target steady flow for cloak comparison)
# ---------------------------------------------------------------------------

def collect_empty_channel(device_id: int) -> dict:
    """Run empty channel (no bodies) and record steady parabolic flow."""
    viscosity = nu_from_re(100.0)

    cuda_cfg, field_cfg = load_configs(CONFIG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(viscosity))
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)

    # Need at least one sensor (legacy API requirement)
    ff.add_sensor((NX - 10, CENTER_Y, 0.0), SENSOR_RADIUS)
    n_obj = ff.obs.size // 2
    ff.run(STABILIZE_STEPS, np.zeros(n_obj, dtype=np.float32))

    # Record a few fields
    ux_list, uy_list = [], []
    for i in range(5):
        ff.run(SAMPLE_INTERVAL, np.zeros(n_obj, dtype=np.float32))
        ux, uy = get_velocity_field(ff, u0=U0)
        ux_list.append(ux)
        uy_list.append(uy)

    ux_all = np.stack(ux_list, axis=0)
    uy_all = np.stack(uy_list, axis=0)

    out_dir = os.path.join(OUTPUT_DIR, "empty_channel")
    os.makedirs(out_dir, exist_ok=True)

    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=ux_all, uy=uy_all)

    meta = {
        "case": "empty_channel",
        "U0": U0,
        "viscosity": viscosity,
        "n_fields": len(ux_list),
    }
    with open(os.path.join(out_dir, "meta.json"), "w") as f:
        json.dump(meta, f, indent=2)

    print("Empty channel flow recorded")
    del ff
    return meta


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

def main():
    ap = argparse.ArgumentParser(description="Phase 1: Data collection")
    ap.add_argument("--case", type=str, required=True,
                    choices=["all", "illusion", "cloak", "uncontrolled",
                             "target_cylinder", "empty_channel"],
                    help="Case to collect")
    ap.add_argument("--device", type=int, default=2, help="GPU device ID")
    args = ap.parse_args()

    # Load Phase 0 data for f_ref / T_ref
    f_ref_path = os.path.join(OUTPUT_DIR, "target_cylinder", "meta.json")
    if os.path.exists(f_ref_path):
        with open(f_ref_path, "r") as f:
            phase0_data = json.load(f)
    else:
        phase0_data = {"T_ref": 15000.0, "f_ref": 6.67e-5}
        print("WARNING: Phase 0 not found, using default T_ref=15000")

    t0 = time.time()
    results = {}

    if args.case in ("all", "illusion"):
        print("=" * 60)
        print("Collecting Illusion case...")
        print("=" * 60)
        phase0_data["illusion_2u"] = True
        results["illusion"] = collect_illusion(args.device, phase0_data)

    if args.case in ("all", "cloak"):
        print("=" * 60)
        print("Collecting Cloak case...")
        print("=" * 60)
        results["cloak"] = collect_cloak(args.device, phase0_data)

    if args.case in ("all", "uncontrolled"):
        print("=" * 60)
        print("Collecting Uncontrolled case...")
        print("=" * 60)
        results["uncontrolled"] = collect_uncontrolled(args.device, phase0_data)

    if args.case in ("all", "target_cylinder"):
        print("=" * 60)
        print("Collecting Target Cylinder...")
        print("=" * 60)
        results["target_cylinder"] = collect_target_cylinder(
            args.device, phase0_data)

    if args.case in ("all", "empty_channel"):
        print("=" * 60)
        print("Collecting Empty Channel (steady target)...")
        print("=" * 60)
        results["empty_channel"] = collect_empty_channel(args.device)

    elapsed = time.time() - t0
    print(f"\nPhase 1 complete in {elapsed:.1f}s")
    return 0


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