"""Generate target data for Illusion scenes (no PPO model needed).

Records target cylinder + sensor signals for any illusion scene.
Only Phase 1 of infer_illusion.py - does not need a PPO model.

Usage:
    conda run -n pycuda_3_10 python scripts/gen_illusion_target.py \\
        --scene illusion_0.6L --device 2

    conda run -n pycuda_3_10 python scripts/gen_illusion_target.py \\
        --scene illusion_0.8L --device 2

    conda run -n pycuda_3_10 python scripts/gen_illusion_target.py \\
        --scene all --device 2
"""
from __future__ import annotations

import argparse
import json
import os
import sys

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 SR_analysis.utils.cfd_interface import load_legacy_configs
from SR_analysis.configs import get_scene, get_scene_list, LEGACY_CFG_DIR, FIFO_LEN

DATA_TYPE = np.float32


def analyze_harmonics(states: np.ndarray, n_harmonics: int = 5) -> list:
    """FFT-based harmonics analysis matching infer_illusion.py exactly."""
    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 * 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


def gen_target_for_scene(scene_name: str, device_id: int):
    """Generate target data for one illusion scene (Phase 1 only)."""
    cfg = get_scene(scene_name)
    nu = cfg["nu"]
    u0 = cfg["u0"]
    l0 = 20.0
    sample_interval = cfg["sample_interval"]
    sensor_x = cfg["sensor_x"]
    target_diam = cfg["target_diameter"]

    output_root = os.path.join(
        os.path.dirname(__file__), "..", "data", "illusion", scene_name,
    )
    os.makedirs(output_root, exist_ok=True)

    print(f"\n{'='*60}")
    print(f"Generating target for {scene_name} (diam={target_diam}L, SI={sample_interval})")
    print(f"{'='*60}")

    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(nu))

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

    # Add target cylinder at x=20*L0
    ff.add_cylinder((20.0 * l0, (ny - 1) / 2, 0.0), target_diam * l0)

    # Add 3 sensors at x = sensor_x * L0
    for y_off in [2.0, 0.0, -2.0]:
        sc = (sensor_x * l0, (ny - 1) / 2 + y_off * l0, 0.0)
        ff.add_sensor(sc, l0 / 4.0)

    n_obj = ff.obs.size // 2

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

    # Record target: 8 channels = [cylinder_fx, fy, sensor0_ux,uy, sensor1_ux,uy, sensor2_ux,uy]
    target_states = np.empty((0, 8), dtype=DATA_TYPE)
    for _ in range(FIFO_LEN):
        ff.run(sample_interval, np.zeros(n_obj, dtype=DATA_TYPE))
        new_state = ff.obs.copy()[0:8]
        target_states = np.vstack((target_states, new_state))

    print(f"  Target recorded: {target_states.shape}")

    # Extract components
    target_sensors = target_states[:, 2:8].copy()
    target_forces = target_states[:, 0:2].copy()
    target_harmonics = analyze_harmonics(target_forces, n_harmonics=5)
    print(f"  Harmonics computed: Cd_dc={target_harmonics[0]['dc']:.6f}, Cl_dc={target_harmonics[1]['dc']:.6f}")
    print(f"  Cd_amp_primary={target_harmonics[0]['amps'][0]:.6f}, Cl_amp_primary={target_harmonics[1]['amps'][0]:.6f}")

    # Save
    np.savez(os.path.join(output_root, "target.npz"),
             target_states=target_states,
             target_sensors=target_sensors)
    with open(os.path.join(output_root, "target_harmonics.json"), "w") as f:
        json.dump(target_harmonics, f, indent=2)

    print(f"  Saved to {output_root}")

    del ff
    return True


def main():
    ap = argparse.ArgumentParser(description="Generate illusion target data")
    ap.add_argument("--scene", type=str, default="all",
                    help='Scene name like "illusion_0.6L", or "all" for all generalization')
    ap.add_argument("--device", type=int, default=2, help="GPU device")
    args = ap.parse_args()

    if args.scene.lower() == "all":
        # Only generate for generalization scenes + verify existing
        all_scenes = ["illusion_0.5L", "illusion_0.6L", "illusion_0.8L",
                       "illusion_1.2L", "illusion_2L"]
        to_gen = []
        for sn in all_scenes:
            data_dir = os.path.join(
                os.path.dirname(__file__), "..", "data", "illusion", sn)
            if not os.path.isfile(os.path.join(data_dir, "target_harmonics.json")):
                to_gen.append(sn)
            else:
                print(f"  {sn}: target data already exists, skipping")
        if not to_gen:
            print("All target data already exists.")
            return
        scene_names = to_gen
    else:
        scene_names = [args.scene]

    for sn in scene_names:
        try:
            gen_target_for_scene(sn, args.device)
        except Exception as e:
            print(f"ERROR on {sn}: {e}")
            import traceback
            traceback.print_exc()

    print("\nDone.")


if __name__ == "__main__":
    main()
