# OID_analysis/scripts/collect_fields_replay.py
"""
Replay PPO actions from CCD checkpoints and save full field time series.
This is the most efficient way to get full field data without re-running PPO.
Also copies/linked controlled.npz, target.npz, norm.json from CCD.

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

import argparse
import json
import os
import shutil
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 (  # noqa: E402
    get_scene, get_scene_list, data_dir_for_scene, LEGACY_CFG_DIR,
)

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

# Source directory in CCD
CCD_DATA = os.path.join(_SRC, "CCD_analysis", "data")


def ccd_data_dir(scene_name: str) -> str:
    cfg = get_scene(scene_name)
    sid = cfg["scene_id"]
    if sid == "karman":
        return os.path.join(CCD_DATA, "karman", "karman_re100")
    elif sid == "illusion":
        # CCD stores at: data/illusion/illusion_0.75L/ etc
        return os.path.join(CCD_DATA, "illusion", scene_name)
    return None


def build_env(cfg, cuda_cfg, field_cfg, device_id):
    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)
        n_phase1 = 4
        ff.run(int(4 * 1280 / u0), np.zeros(n_phase1, 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)
    out_dir = data_dir_for_scene(scene_name)
    os.makedirs(out_dir, exist_ok=True)

    src_dir = ccd_data_dir(scene_name)
    if src_dir is None or not os.path.isdir(src_dir):
        raise FileNotFoundError(f"CCD data not found for {scene_name} at {src_dir}")

    u0 = cfg["u0"]
    si = cfg["sample_interval"]
    ac_scale = cfg["action_scale"]
    ac_bias = cfg["action_bias"]
    n_obj = cfg["n_objects_env"]
    obs_start, obs_end = cfg["obs_slice"]

    # Copy/symlink existing data
    for fname in ["controlled.npz", "target.npz", "norm.json", "config.json",
                   "target_harmonics.json", "result.json"]:
        src = os.path.join(src_dir, fname)
        dst = os.path.join(out_dir, fname)
        if os.path.isfile(src) and not os.path.isfile(dst):
            shutil.copy2(src, dst)
            print(f"  Copied {fname}")

    # Load actions from controlled.npz
    controlled = np.load(os.path.join(src_dir, "controlled.npz"))
    actions = controlled["actions"]
    n_steps = len(actions)
    print(f"  {n_steps} actions loaded")

    # Load checkpoints
    ddf_ckpt = np.load(os.path.join(src_dir, "ddf_checkpoint.npy"))
    fifo_ckpt = np.load(os.path.join(src_dir, "fifo_checkpoint.npy"))

    # Build env
    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)

    # Restore DDF
    ff.ddf = ddf_ckpt.copy()
    ff.apply_ddf()
    print(f"  DDF restored ({len(ddf_ckpt)} floats)")

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

    # Replay and save fields
    # Use masked ROI: x=[400, 1000] (20-50 L0), y=[100, 400] (full relevant y)
    # This captures near-body + near-wake + sensor zone
    nx_full = 1280
    ny_full = 512
    x_start, x_end = 400, 1000
    y_start, y_end = 100, 400
    nx_roi = x_end - x_start
    ny_roi = y_end - y_start
    print(f"  ROI: x=[{x_start},{x_end}], y=[{y_start},{y_end}] ({nx_roi}x{ny_roi})")

    ux_list, uy_list = [], []
    sens_replay, forc_replay = [], []

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

        obs_slice = ff.obs.copy()[obs_start:obs_end]
        fifo.append(obs_slice)
        sens_replay.append(obs_slice[0:6])
        forc_replay.append(obs_slice[6:12])

        # Save cropped field
        ux_full, uy_full = get_velocity_field(ff, u0=u0)
        ux_list.append(ux_full[y_start:y_end, x_start:x_end])
        uy_list.append(uy_full[y_start:y_end, x_start:x_end])

    # Verify replay fidelity
    orig_sensors = controlled["sensors"]
    orig_forces = controlled["forces"]
    sens_replay = np.array(sens_replay, dtype=np.float32)
    forc_replay = np.array(forc_replay, dtype=np.float32)
    diff_sens = float(np.max(np.abs(sens_replay - orig_sensors)))
    diff_forc = float(np.max(np.abs(forc_replay - orig_forces)))
    print(f"  Replay: max diff sensors={diff_sens:.6e}, forces={diff_forc:.6e}")

    # Save fields
    ux_arr = np.stack(ux_list).astype(np.float32)
    uy_arr = np.stack(uy_list).astype(np.float32)
    print(f"  Fields shape: {ux_arr.shape}, size ~{ux_arr.nbytes/1e6:.1f} MB")

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

    # Save ROI metadata
    roi_meta = {"x_start": x_start, "x_end": x_end,
                "y_start": y_start, "y_end": y_end,
                "nx_full": nx_full, "ny_full": ny_full}
    with open(os.path.join(out_dir, "roi_meta.json"), "w") as f:
        json.dump(roi_meta, f, indent=2)

    # Vorticity viz
    omega_viz = vorticity_from_ddf(ff, u0=u0)
    save_vorticity_png(os.path.join(out_dir, "vorticity_controlled.png"),
                       omega_viz, title=f"{scene_name} controlled")

    # Replay verify report
    verify_info = {
        "scene": scene_name,
        "n_steps": n_steps,
        "n_fields": len(ux_list),
        "max_diff_sensors": diff_sens,
        "max_diff_forces": diff_forc,
        "passed": diff_sens < 1e-4 and diff_forc < 1e-4,
    }
    with open(os.path.join(out_dir, "replay_verify.json"), "w") as f:
        json.dump(verify_info, f, indent=2)

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


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scene", type=str, required=True,
                    help="Scene name: karman_re100, illusion_0.75L, illusion_1.0L, illusion_1.5L")
    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()
