"""PPO inference with field visualization for Illusion scenes.

Reuses infer_illusion.py's setup but adds vorticity field saving.
No SINDy involved — just PPO model in the loop.

Usage:
    conda run -n pycuda_3_10 python src/SR_analysis/scripts/visualize_ppo_illusion.py \\
        --diameter 1.0 --device 0 --steps 300
"""
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
from LegacyCelerisLab import utils as legacy_utils

from SR_analysis.utils.cfd_interface import (
    load_legacy_configs, load_ppo_model, vorticity_from_ddf, save_vorticity_png,
)
from SR_analysis.configs import (
    get_scene, get_scene_list, model_path_for_scene, LEGACY_CFG_DIR, FIFO_LEN,
)

DATA_TYPE = np.float32
NX = 1280


def analyze_harmonics(states: np.ndarray, n_harmonics: int = 5) -> list:
    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_states_at(t: int, harmonics: list) -> np.ndarray:
    D = len(harmonics)
    result = np.zeros(D, dtype=np.float32)
    for d, h in enumerate(harmonics):
        val = h['dc']
        for amp, freq, phase in zip(h['amps'], h['freqs'], h['phases']):
            val += amp * np.cos(2 * np.pi * freq * t + phase)
        result[d] = float(val)
    return result


def run_visualization(scene_name: str, device_id: int, n_infer_steps: int,
                      out_root: str):
    cfg = get_scene(scene_name)
    u0 = cfg["u0"]
    l0 = 20.0
    sample_interval = cfg["sample_interval"]
    action_scale = cfg["action_scale"]
    action_bias = cfg["action_bias"]
    target_diam = cfg["target_diameter"]
    s_dim = cfg["s_dim"]
    nu = cfg["nu"]

    # Auto-set steps if too short
    min_steps = int(1.5 * NX / u0 / sample_interval)
    if n_infer_steps < min_steps:
        n_infer_steps = max(min_steps, 300)
        print(f"  auto-set steps={n_infer_steps} (min={min_steps})")

    os.makedirs(out_root, exist_ok=True)

    print(f"\n{'='*60}")
    print(f"PPO viz: {scene_name}  diam={target_diam}L  steps={n_infer_steps}")
    print(f"{'='*60}")

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

    # -- Phase 1: Target recording -------------------------------------------
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    ny = ff.FIELD_SHAPE[1]

    ff.add_cylinder((20.0 * l0, (ny - 1) / 2, 0.0), target_diam * l0)
    for y_off in [2.0, 0.0, -2.0]:
        sc = (30.0 * l0, (ny - 1) / 2 + y_off * l0, 0.0)
        ff.add_sensor(sc, l0 / 4.0)

    stabilise = int(4 * NX / u0)
    ff.run(stabilise, np.zeros(4, dtype=DATA_TYPE))

    target_all = np.empty((0, 8), dtype=DATA_TYPE)
    for _ in range(FIFO_LEN):
        ff.run(sample_interval, np.zeros(4, dtype=DATA_TYPE))
        target_all = np.vstack((target_all, ff.obs.copy()[0:8]))
    target_sensors = target_all[:, 2:8].copy()
    target_forces = target_all[:, 0:2].copy()
    target_harmonics = analyze_harmonics(target_forces)

    # Save target cylinder vorticity
    omega_tgt = vorticity_from_ddf(ff, u0)
    save_vorticity_png(os.path.join(out_root, f"{scene_name}_target_vorticity.png"),
                       omega_tgt, title=f"{scene_name} target ({target_diam}L)")
    print(f"  target vorticity saved")

    del ff

    # -- Phase 2: Pinball env setup -----------------------------------------
    ff = FlowField(field_cfg, cuda_cfg, device_id=device_id)

    for y_off in [2.0, 0.0, -2.0]:
        sc = (30.0 * l0, (ny - 1) / 2 + y_off * l0, 0.0)
        ff.add_sensor(sc, l0 / 4.0)
    ff.add_cylinder((19.0 * l0, (ny - 1) / 2, 0.0), l0 / 2.0)
    ff.add_cylinder((20.3 * l0, (ny - 1) / 2 + 0.75 * l0, 0.0), l0 / 2.0)
    ff.add_cylinder((20.3 * l0, (ny - 1) / 2 - 0.75 * l0, 0.0), l0 / 2.0)

    n_obj = ff.obs.size // 2
    assert n_obj == 6
    ff.run(stabilise, np.zeros(n_obj, dtype=DATA_TYPE))

    # save_ddf(1): post-stabilization
    ff.get_ddf()
    ff.save_ddf()

    # Norm collection
    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])))
    print(f"  norm: force_norm_fact={force_norm_fact:.6f}")

    # Bias FIFO rollout: [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_ddf(2): post-bias
    ff.get_ddf()
    ff.save_ddf()
    ff.apply_ddf()

    # Save uncontrolled reference (before PPO, same initial state, same duration)
    ff_un = FlowField(field_cfg, cuda_cfg, device_id=device_id)
    ff_un.add_sensor((30.0 * l0, (ny - 1) / 2 + 2.0 * l0, 0.0), l0 / 4.0)
    ff_un.add_sensor((30.0 * l0, (ny - 1) / 2, 0.0), l0 / 4.0)
    ff_un.add_sensor((30.0 * l0, (ny - 1) / 2 - 2.0 * l0, 0.0), l0 / 4.0)
    ff_un.add_cylinder((19.0 * l0, (ny - 1) / 2, 0.0), l0 / 2.0)
    ff_un.add_cylinder((20.3 * l0, (ny - 1) / 2 + 0.75 * l0, 0.0), l0 / 2.0)
    ff_un.add_cylinder((20.3 * l0, (ny - 1) / 2 - 0.75 * l0, 0.0), l0 / 2.0)
    ff_un.run(stabilise, np.zeros(6, dtype=DATA_TYPE))
    ff_un.get_ddf()
    ff_un.save_ddf()
    ff_un.apply_ddf()
    for _ in range(FIFO_LEN):
        ff_un.run(sample_interval, bias_arr)
    ff_un.get_ddf()
    ff_un.save_ddf()
    ff_un.apply_ddf()
    for _ in range(n_infer_steps):
        ff_un.run(sample_interval, np.zeros(6, dtype=DATA_TYPE))
    omega_unc = vorticity_from_ddf(ff_un, u0)
    save_vorticity_png(os.path.join(out_root, f"{scene_name}_uncontrolled_vorticity.png"),
                       omega_unc, title=f"{scene_name} uncontrolled (zero action)")
    print(f"  uncontrolled vorticity saved")
    del ff_un

    # -- Phase 3: PPO controlled rollout ------------------------------------
    model_path = model_path_for_scene(scene_name)
    if model_path is None:
        print(f"  ERROR: no model for {scene_name}")
        return

    model = load_ppo_model(model_path, device=f"cuda:{device_id}", s_dim=s_dim)
    model.set_random_seed(0)

    print(f"  PPO controlled rollout ({n_infer_steps} steps) ...")
    ff.restore_ddf()
    ff.apply_ddf()

    # FIFO reinit: run bias for FIFO_LEN steps
    fifo = deque(maxlen=FIFO_LEN)
    for _ in range(FIFO_LEN):
        ff.context.push()
        ff.run(sample_interval, bias_arr)
        ff.context.pop()
        fifo.append(ff.obs.copy()[0:12])

    obs = np.zeros(s_dim, dtype=np.float32)
    field_save_interval = max(1, n_infer_steps // 8)

    for step in range(n_infer_steps):
        action, _ = model.predict(obs, deterministic=True)
        action = action.astype(np.float32).flatten()

        temp = np.zeros(n_obj, dtype=DATA_TYPE)
        temp[3:6] = np.array(
            (action * action_scale + list(action_bias)) * u0, dtype=DATA_TYPE)

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

        obs_slice = ff.obs.copy()[0:12]
        fifo.append(obs_slice)

        forces_norm = obs_slice[6:12] / force_norm_fact
        sens_norm = (obs_slice[0:6] - sens_deviation) / sens_norm_fact
        obs12 = np.clip(np.hstack([forces_norm, sens_norm]), -1, 1).astype(np.float32)

        if s_dim == 14:
            tv = gen_target_states_at(step, target_harmonics)
            target_cd = np.clip(tv[0] / force_norm_fact, -1, 1)
            target_cl = np.clip(tv[1] / force_norm_fact, -1, 1)
            obs = np.zeros(14, dtype=np.float32)
            obs[:12] = obs12
            obs[12] = target_cd
            obs[13] = target_cl
        else:
            obs = obs12

        # Save field at intervals
        if step % field_save_interval == 0 or step == n_infer_steps - 1:
            omega = vorticity_from_ddf(ff, u0)
            pct = int(100 * step / n_infer_steps)
            fname = os.path.join(out_root, f"{scene_name}_PPO_step{step:04d}_pct{pct:02d}.png")
            save_vorticity_png(fname, omega,
                               title=f"{scene_name} PPO step {step}/{n_infer_steps} ({pct}%)")
            print(f"  field saved: step {step}/{n_infer_steps} ({pct}%)")

    # Final field
    omega = vorticity_from_ddf(ff, u0)
    save_vorticity_png(os.path.join(out_root, f"{scene_name}_PPO_vorticity.png"),
                       omega, title=f"{scene_name} PPO controlled (sim 0.97+)")
    print(f"  final vorticity saved")

    del ff
    print(f"  Done: {scene_name}")


def main():
    ap = argparse.ArgumentParser(description="PPO visualization for Illusion")
    ap.add_argument("--diameter", type=str, default="1.0", help='0.75, 1.0, 1.5, or all')
    ap.add_argument("--device", type=int, default=0)
    ap.add_argument("--steps", type=int, default=300)
    ap.add_argument("--out-root", type=str, default=None)
    args = ap.parse_args()

    if args.diameter.lower() == "all":
        scene_names = get_scene_list("illusion")
    else:
        d = float(args.diameter)
        scene_names = [sn for sn in get_scene_list("illusion")
                       if abs(get_scene(sn)["target_diameter"] - d) < 0.01]

    if args.out_root is None:
        args.out_root = os.path.join(os.path.dirname(__file__), "..",
                                     "validate", "results")

    for sn in scene_names:
        run_visualization(sn, args.device, args.steps, args.out_root)


if __name__ == "__main__":
    main()
