"""Replay PPO actions with DDF+FIFO checkpoint, save phase-aligned fields.

For PPO cases (illusion, karman):
  - Loads DDF+FIFO checkpoint saved by collect_*.py
  - Rebuilds exact same FlowField geometry
  - Restores DDF and FIFO
  - Replays ALL saved actions step by step
  - Saves raw (non-interpolated) field snapshots at step indices from phase_plan.json
  - Verifies replay fidelity by comparing sensors/forces with original controlled.npz

For open-loop cases (target_cylinder, pinball):
  - No DDF checkpoint needed (no PPO actions)
  - Uses fields.npz directly + phase_plan.json to extract aligned fields

Usage:
    conda run -n pycuda_3_10 python src/CCD_analysis/scripts/replay_fields.py --scene illusion_1.0L --device 2
    conda run -n pycuda_3_10 python src/CCD_analysis/scripts/replay_fields.py --scene target_cylinder_1.0L --device 2
    conda run -n pycuda_3_10 python src/CCD_analysis/scripts/replay_fields.py --scene pinball --device 2

Output: data/{scene_id}/{scene_name}/fields_aligned.npz
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

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 CCD_analysis.configs import (
    get_scene, get_scene_list, data_dir_for_scene, LEGACY_CFG_DIR, L0, CENTER_Y,
)
from CCD_analysis.utils.cfd_interface import load_legacy_configs, get_velocity_field

DATA_TYPE = np.float32
FIFO_LEN = 150


def build_env(cfg: dict, cuda_cfg, field_cfg, device_id: int) -> FlowField:
    """Build the exact same FlowField geometry as the original collection script."""
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    u0 = cfg["u0"]

    if cfg.get("has_disturbance", False):
        # Karman layout: dist_cyl + 3 sensors first
        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:
        # Illusion / plain layout: 3 sensors first
        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)

    # Add pinball
    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_ppo(scene_name: str, device_id: int, verify_tol: float = 1e-4) -> int:
    """Replay PPO inference using DDF+FIFO checkpoint, save aligned fields."""
    cfg = get_scene(scene_name)
    out_dir = data_dir_for_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"]
    obs_start, obs_end = cfg["obs_slice"]

    # Load phase plan
    plan_path = os.path.join(os.path.join(os.path.dirname(out_dir), "..", "resampled", scene_name), "phase_plan.json")
    # Try alternate path
    if not os.path.isfile(plan_path):
        plan_path = os.path.join(os.path.join(os.path.dirname(__file__), "..", "data", "resampled", scene_name), "phase_plan.json")
    if not os.path.isfile(plan_path):
        plan_path = os.path.join(cfg["DATA_DIR"] if "DATA_DIR" in cfg else os.path.join(os.path.dirname(__file__), "..", "data"), "resampled", scene_name, "phase_plan.json")
    # Final fallback: data/resampled
    from CCD_analysis.configs import DATA_DIR as CCD_DATA_DIR
    plan_path = os.path.join(CCD_DATA_DIR, "resampled", scene_name, "phase_plan.json")

    if not os.path.isfile(plan_path):
        raise FileNotFoundError(f"phase_plan.json not found for {scene_name}. Run detect_period.py first.")

    with open(plan_path) as f:
        plan = json.load(f)
    step_indices = set(plan["step_indices"])
    print(f"  phase plan: {plan['n_cycles']} cycles x {plan['n_pts']} pts = {len(step_indices)} steps", flush=True)

    # Load actions and original data for verification
    controlled = np.load(os.path.join(out_dir, "controlled.npz"))
    actions = controlled["actions"]
    orig_sensors = controlled["sensors"]
    orig_forces = controlled["forces"]
    n_steps = len(actions)

    # Load DDF+FIFO checkpoint
    ddf_ckpt = np.load(os.path.join(out_dir, "ddf_checkpoint.npy"))
    fifo_ckpt = np.load(os.path.join(out_dir, "fifo_checkpoint.npy"))

    # Build env and restore
    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)
    n_obj_built = ff.obs.size // 2
    assert n_obj_built == n_obj, f"Object count mismatch: {n_obj_built} vs expected {n_obj}"

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

    # Replay
    from collections import deque
    fifo = deque(maxlen=FIFO_LEN)
    for s in fifo_ckpt:
        fifo.append(s)

    sens_replay, forc_replay = [], []
    ux_list, uy_list = [], []
    max_diff_sens, max_diff_forc = 0.0, 0.0

    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 field at selected step indices
        if step in step_indices:
            ux, uy = get_velocity_field(ff, u0=u0)
            ux_list.append(ux)
            uy_list.append(uy)

    # Verify replay fidelity
    sens_replay = np.array(sens_replay, dtype=np.float32)
    forc_replay = np.array(forc_replay, dtype=np.float32)

    diff_sens = np.max(np.abs(sens_replay - orig_sensors))
    diff_forc = np.max(np.abs(forc_replay - orig_forces))
    print(f"  Replay max diff: sensors={diff_sens:.6e}, forces={diff_forc:.6e}", flush=True)

    if diff_sens > verify_tol or diff_forc > verify_tol:
        print(f"  WARNING: replay diff exceeds tolerance ({verify_tol}). "
              f"CUDA atomic non-determinism may cause minor variation.", flush=True)

    # Save aligned fields
    out_path = os.path.join(out_dir, "fields_aligned.npz")
    np.savez_compressed(out_path,
                        ux=np.stack(ux_list), uy=np.stack(uy_list))
    print(f"  Saved {len(ux_list)} field snapshots to {out_path}", flush=True)

    # Save verification report
    verify = {
        "scene": scene_name,
        "n_steps": n_steps,
        "n_fields": len(ux_list),
        "max_diff_sensors": float(diff_sens),
        "max_diff_forces": float(diff_forc),
        "tolerance": verify_tol,
        "passed": bool(diff_sens <= verify_tol and diff_forc <= verify_tol),
    }
    with open(os.path.join(out_dir, "replay_verify.json"), "w") as f:
        json.dump(verify, f, indent=2)

    del ff
    controlled.close()
    return len(ux_list)


def replay_open_loop(scene_name: str, device_id: int) -> int:
    """For open-loop cases: extract fields from phase plan indices directly.

    Open-loop cases (target_cylinder, pinball) have no actions to replay.
    Their fields.npz already contains the raw time series.
    We just need to pick the phase-aligned snapshots.
    """
    cfg = get_scene(scene_name)
    out_dir = data_dir_for_scene(scene_name)

    from CCD_analysis.configs import DATA_DIR as CCD_DATA_DIR
    plan_path = os.path.join(CCD_DATA_DIR, "resampled", scene_name, "phase_plan.json")
    if not os.path.isfile(plan_path):
        raise FileNotFoundError(f"phase_plan.json not found for {scene_name}")

    with open(plan_path) as f:
        plan = json.load(f)
    step_indices = plan["step_indices"]

    # Load existing fields
    fields_path = os.path.join(out_dir, "fields.npz")
    if not os.path.isfile(fields_path):
        raise FileNotFoundError(f"fields.npz not found for {scene_name}")

    fd = np.load(fields_path)
    ux_all, uy_all = fd["ux"], fd["uy"]

    n_avail = len(ux_all)
    valid_idx = [idx for idx in step_indices if idx < n_avail]
    if len(valid_idx) < len(step_indices):
        print(f"  WARNING: {len(step_indices) - len(valid_idx)} step indices out of range "
              f"(max={n_avail - 1}). Using {len(valid_idx)} valid snapshots.", flush=True)

    ux_list = [ux_all[idx] for idx in valid_idx]
    uy_list = [uy_all[idx] for idx in valid_idx]

    out_path = os.path.join(out_dir, "fields_aligned.npz")
    np.savez_compressed(out_path, ux=np.stack(ux_list), uy=np.stack(uy_list))
    print(f"  Extracted {len(ux_list)} phase-aligned field snapshots to {out_path}", flush=True)

    fd.close()
    return len(ux_list)


def replay_single(scene_name: str, device_id: int, verify_tol: float = 1e-4) -> int:
    """Route to the appropriate replay method based on scene type."""
    cfg = get_scene(scene_name)
    source = cfg.get("source", "open_loop")

    if source == "PPO_inference":
        return replay_ppo(scene_name, device_id, verify_tol)
    else:
        return replay_open_loop(scene_name, device_id)


def main():
    ap = argparse.ArgumentParser(description="Replay and save phase-aligned fields")
    ap.add_argument("--scene", type=str, default="karman_re100")
    ap.add_argument("--device", type=int, default=2)
    ap.add_argument("--tol", type=float, default=1e-4,
                    help="Replay verification tolerance")
    ap.add_argument("--all-periodic", action="store_true",
                    help="Replay all periodic scenes")
    args = ap.parse_args()

    if args.all_periodic:
        from CCD_analysis.configs import SCENES
        scenes = [name for name, cfg in SCENES.items()
                  if cfg.get("target_type") == "periodic"]
    else:
        scenes = [args.scene]

    for scene in scenes:
        print(f"\n=== Replaying fields for {scene} ===", flush=True)
        t0 = time.time()
        try:
            n = replay_single(scene, args.device, args.tol)
            print(f"  {n} fields in {time.time() - t0:.1f}s", flush=True)
        except (FileNotFoundError, AssertionError, ValueError) as e:
            print(f"  FAILED: {e}", flush=True)

    print("\nDone.", flush=True)


if __name__ == "__main__":
    main()
