# OID_analysis/scripts/replay_full_fields.py
"""
Replay PPO actions from CCD checkpoints and save FULL field time series (no cropping).
Needed because Round 1 saved ROI-cropped fields for illusion.
Must run from repo root with PYTHONPATH=src in pycuda_3_10 env.

Usage:
    conda run -n pycuda_3_10 python src/OID_analysis/scripts/replay_full_fields.py \
        --scene illusion_0.75L --device 1
    conda run -n pycuda_3_10 python src/OID_analysis/scripts/replay_full_fields.py \
        --scene illusion_1.0L --device 3
    conda run -n pycuda_3_10 python src/OID_analysis/scripts/replay_full_fields.py \
        --scene illusion_1.5L --device 3
"""
from __future__ import annotations

import argparse
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,
)
from OID_analysis.configs import (  # noqa: E402
    get_scene, LEGACY_CFG_DIR,
)

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


def build_env(cfg, cuda_cfg, field_cfg, device_id):
    """Build environment with FULL geometry (no cropping)."""
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    u0 = cfg["u0"]

    if cfg.get("has_disturbance", False):
        ff.add_cylinder((10.0 * L0, CENTER_Y, 0.0), L0)
        for y_off in [2.0, 0.0, -2.0]:
            ff.add_sensor((cfg["sensor_x"] * L0,
                          CENTER_Y + y_off * L0, 0.0), L0 / 4.0)
        n1 = 4
        ff.run(int(4 * 1280 / u0), np.zeros(n1, dtype=DATA_TYPE))
    else:
        for y_off in [2.0, 0.0, -2.0]:
            ff.add_sensor((cfg["sensor_x"] * L0,
                          CENTER_Y + y_off * L0, 0.0), L0 / 4.0)

    ff.add_cylinder((cfg["pinball_front_x"] * L0, CENTER_Y, 0.0), L0 / 2.0)
    ff.add_cylinder((cfg["pinball_rear_x"] * L0,
                    CENTER_Y + 0.75 * L0, 0.0), L0 / 2.0)
    ff.add_cylinder((cfg["pinball_rear_x"] * L0,
                    CENTER_Y - 0.75 * L0, 0.0), L0 / 2.0)
    return ff


def replay(scene_name: str, device_id: int):
    cfg = get_scene(scene_name)
    u0 = cfg["u0"]
    si = cfg["sample_interval"]
    ac_scale = cfg["action_scale"]
    ac_bias = cfg["action_bias"]
    n_obj = cfg["n_objects_env"]

    # CCD source directory
    sid = cfg["scene_id"]
    if sid == "karman":
        src_dir = os.path.join(_SRC, "CCD_analysis", "data", "karman", "karman_re100")
    elif sid == "illusion":
        src_dir = os.path.join(_SRC, "CCD_analysis", "data", "illusion", scene_name)
    else:
        raise ValueError(f"Unknown scene type: {sid}")

    out_dir = os.path.join(os.path.dirname(__file__), "..", "data",
                           "karman_cloak" if sid == "karman" else "illusion", scene_name)
    os.makedirs(out_dir, exist_ok=True)

    # Load checkpoints and actions
    ddf_ckpt = np.load(os.path.join(src_dir, "ddf_checkpoint.npy"))
    fifo_ckpt = np.load(os.path.join(src_dir, "fifo_checkpoint.npy"))
    controlled = np.load(os.path.join(src_dir, "controlled.npz"))
    actions = controlled["actions"]
    n_steps = len(actions)
    print(f"  {n_steps} actions loaded")

    # Build env with full geometry
    cuda_cfg, field_cfg = load_legacy_configs(LEGACY_CFG_DIR)
    field_cfg = field_cfg._replace(viscosity=float(cfg["nu"]), velocity=float(u0))
    ff = build_env(cfg, cuda_cfg, field_cfg, device_id)

    ff.ddf = ddf_ckpt.copy()
    ff.apply_ddf()
    print(f"  DDF restored")

    fifo = deque(maxlen=FIFO_LEN)
    for s in fifo_ckpt:
        fifo.append(s)

    # Replay and save FULL fields
    ux_list, uy_list = [], []
    for step in range(n_steps):
        action = actions[step]
        omega = (action * ac_scale + np.array(ac_bias, dtype=np.float32)) * u0
        temp = np.zeros(n_obj, dtype=DATA_TYPE)
        temp[n_obj - 3:] = omega

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

        ux_full, uy_full = get_velocity_field(ff, u0=u0)
        ux_list.append(ux_full)
        uy_list.append(uy_full)

    ux_arr = np.stack(ux_list).astype(np.float32)
    uy_arr = np.stack(uy_list).astype(np.float32)
    print(f"  Full fields shape: {ux_arr.shape}")

    # Overwrite old cropped fields
    np.savez_compressed(os.path.join(out_dir, "fields.npz"),
                        ux=ux_arr, uy=uy_arr)

    del ff, controlled
    print(f"  {n_steps} full fields saved to {out_dir}")
    return n_steps


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", type=str, required=True)
    ap.add_argument("--device", type=int, default=3)
    args = ap.parse_args()

    t0 = time.time()
    n = replay(args.scene, args.device)
    print(f"Done in {time.time() - t0:.1f}s -> {n} fields")


if __name__ == "__main__":
    main()
