# OID_analysis/scripts/collect_illusion_qblk.py
"""
Collect illusion-position pinball baseline (q_blk for illusion scenes).
Pinball at illusion geometry (front_x=19, rear_x=20.3), sensors at x=30.
Zero rotation, natural shedding.

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

Output: data/steady_cloak/pinball_baseline_illusion/
    fields.npz, sensors.npz, forces.npz, norm.json, ddf_checkpoint.npy
"""
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  # noqa: E402
from OID_analysis.utils.cfd_interface import (  # noqa: E402
    load_legacy_configs, get_velocity_field, save_vorticity_png, vorticity_from_ddf,
)
from OID_analysis.configs import LEGACY_CFG_DIR  # noqa: E402

DATA_TYPE = np.float32
L0 = 20.0
CENTER_Y = (512 - 1) / 2.0
U0 = 0.01
SAMPLE_INTERVAL = 800
FIFO_LEN = 150


def collect(device_id: int, n_steps: int) -> str:
    out_dir = os.path.join(os.path.dirname(__file__), "..", "data",
                           "steady_cloak", "pinball_baseline_illusion")
    os.makedirs(out_dir, exist_ok=True)

    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)

    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    print(f"  Illusion q_blk: {ff.FIELD_SHAPE}")

    # --- Illusion geometry ---
    # Sensors at x=30*L0 (different from cloak's x=40)
    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)
    # Pinball at illusion positions (front_x=19, rear_x=20.3)
    ff.add_cylinder((19.0 * L0, CENTER_Y, 0.0), L0 / 2.0)
    ff.add_cylinder((20.3 * L0, CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0)
    ff.add_cylinder((20.3 * L0, CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0)

    n_obj = 6
    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))
    ff.get_ddf()
    ff.save_ddf()

    # --- Norm collection (zero action) ---
    fifo = deque(maxlen=FIFO_LEN)
    for _ in range(FIFO_LEN):
        ff.run(SAMPLE_INTERVAL, np.zeros(n_obj, dtype=DATA_TYPE))
        fifo.append(ff.obs.copy()[0:12])

    temp = np.array(fifo, dtype=DATA_TYPE)
    force_norm_fact = 6.0 * float(np.max(np.abs(temp[:, 6:12])))
    sens_deviation = np.mean(temp[:, 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[:, 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": list([0.0, -2.0, 2.0]),  # illusion bias for reference
    }
    print(f"  norm: force_norm_fact={force_norm_fact:.6f}")

    # --- Bias-action FIFO init (matches legacy_env_imit.py: [0, -U0, U0]) ---
    ff.apply_ddf()
    bias_arr = np.zeros(n_obj, dtype=DATA_TYPE)
    bias_arr[3] = 0.0
    bias_arr[4] = -1.0 * U0
    bias_arr[5] = 1.0 * U0
    fifo.clear()
    for _ in range(FIFO_LEN):
        ff.run(SAMPLE_INTERVAL, bias_arr)
        fifo.append(ff.obs.copy()[0:12])
    save_states_arr = np.array(fifo, dtype=DATA_TYPE)

    # Save DDF+FIFO checkpoint (for future PPO controlled replay)
    ff.get_ddf()
    np.save(os.path.join(out_dir, "ddf_checkpoint.npy"), ff.ddf)
    np.save(os.path.join(out_dir, "fifo_checkpoint.npy"), save_states_arr)
    ff.apply_ddf()

    # --- Record zero-action rollout (PAYLOAD) ---
    sens_list, forc_list, ux_list, uy_list = [], [], [], []
    for step in range(n_steps):
        ff.run(SAMPLE_INTERVAL, np.zeros(n_obj, dtype=DATA_TYPE))
        obs_slice = ff.obs.copy()[0:12]
        sens_list.append(obs_slice[0:6])
        forc_list.append(obs_slice[6:12])
        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(os.path.join(out_dir, "forces.npz"),
             forces=np.array(forc_list, dtype=np.float32))
    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=np.stack(ux_list).astype(np.float32),
                        uy=np.stack(uy_list).astype(np.float32))

    # Save norm (non-numpy values)
    norm_json = {k: v for k, v in norm.items() if not isinstance(v, np.ndarray)}
    with open(os.path.join(out_dir, "norm.json"), "w") as f:
        json.dump(norm_json, f, indent=2)

    # Config metadata
    meta = {
        "scene": "pinball_baseline_illusion",
        "geometry": "illusion",
        "sensor_x": 30.0,
        "pinball_front_x": 19.0,
        "pinball_rear_x": 20.3,
        "n_steps": n_steps,
        "sample_interval": SAMPLE_INTERVAL,
        "u0": U0,
    }
    with open(os.path.join(out_dir, "config.json"), "w") as f:
        json.dump(meta, f, indent=2)

    # Vorticity plot
    omega = vorticity_from_ddf(ff, u0=U0)
    save_vorticity_png(os.path.join(out_dir, "vorticity_illusion_qblk.png"),
                       omega, title="Illusion q_blk (zero rotation, illusion position)")

    del ff
    print(f"  Saved {n_steps} snapshots ({ux_list[0].shape}) to {out_dir}")
    return out_dir


def main():
    ap = argparse.ArgumentParser()
    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.device, args.steps)
    print(f"Done in {time.time() - t0:.1f}s -> {out}")


if __name__ == "__main__":
    main()
