#!/usr/bin/env python3
"""Analyze Illusion degradation: compare PPO strategies across diameters.

Pure analysis of PPO controlled.npz data — no CFD needed.
Computes FFT spectra, action statistics, investigates joint formula
degradation mechanism without running closed-loop validations.

Usage:
    conda run -n pycuda_3_10 python scripts/analyze_illusion_degradation.py
"""
from __future__ import annotations

import os
import sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

_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 SR_analysis.configs import get_scene, _illusion_key as scene_key

DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data", "illusion")
OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "data", "degradation_analysis")
os.makedirs(OUT_DIR, exist_ok=True)

U0 = 0.01
D_CYL = 20.0


def _fmt_diam(d: float) -> str:
    s = f"{d:.3f}".rstrip("0").rstrip(".")
    return f"{s}L"


def load_ppo_data(diam: float) -> dict:
    key = f"illusion_{_fmt_diam(diam)}"
    path = os.path.join(DATA_DIR, key, "controlled.npz")
    npz = np.load(path, allow_pickle=True)
    return {
        "sensors": npz["sensors"].astype(np.float64),
        "forces": npz["forces"].astype(np.float64),
        "actions": npz["actions"].astype(np.float64),  # normalized [-1,1]
    }


def compute_dimless(sensors, forces, u0=U0, d=D_CYL, rho=1.0):
    s = np.asarray(sensors, dtype=np.float64)
    f = np.asarray(forces, dtype=np.float64)
    return {
        "u_hat_T": s[:, 0] / u0, "v_hat_T": s[:, 1] / u0,
        "u_hat_C": s[:, 2] / u0, "v_hat_C": s[:, 3] / u0,
        "u_hat_B": s[:, 4] / u0, "v_hat_B": s[:, 5] / u0,
        "Cd_F": 2 * f[:, 0] / (rho * u0**2 * d),
        "Cl_F": 2 * f[:, 1] / (rho * u0**2 * d),
        "Cd_B": 2 * f[:, 2] / (rho * u0**2 * d),
        "Cl_B": 2 * f[:, 3] / (rho * u0**2 * d),
        "Cd_T": 2 * f[:, 4] / (rho * u0**2 * d),
        "Cl_T": 2 * f[:, 5] / (rho * u0**2 * d),
    }


def compute_features(dim, a_phys):
    """Compute dimensionless features from PPO data."""
    uB, uC, uT = dim["u_hat_B"], dim["u_hat_C"], dim["u_hat_T"]
    vB, vC, vT = dim["v_hat_B"], dim["v_hat_C"], dim["v_hat_T"]
    CdF, CdT, CdB = dim["Cd_F"], dim["Cd_T"], dim["Cd_B"]
    ClF, ClT, ClB = dim["Cl_F"], dim["Cl_T"], dim["Cl_B"]

    feat = {}
    feat["u_m"] = (uB + uC + uT) / 3
    feat["u_a"] = (uT - uB) / 2
    feat["u_c"] = uC
    feat["v_a"] = (vT - vB) / 2
    feat["Cd_tot"] = CdF + CdT + CdB
    feat["Cd_rear"] = CdT + CdB
    feat["Cl_tot"] = ClF + ClT + ClB
    feat["Cl_diff"] = ClT - ClB
    return feat


def main():
    print("=" * 60)
    print("Illusion Degradation Analysis")
    print("=" * 60)

    # 1. Load PPO data for all three trained diameters
    ppo = {}
    for d in [0.75, 1.0, 1.5]:
        ppo[d] = load_ppo_data(d)
        cfg = get_scene(scene_key(d))
        a = ppo[d]["actions"]
        # Convert normalized actions to physical alpha
        scale = cfg["action_scale"]
        bias = np.array(cfg["action_bias"])
        a_phys = (a * scale + bias) * cfg["u0"]
        alpha = a_phys / cfg["u0"]
        print(f"\n--- {d}L, SI={cfg['sample_interval']} ---")
        print(f"  Steps: {a.shape[0]}")
        print(f"  Alpha (front): mean={alpha[:,0].mean():.3f}, std={alpha[:,0].std():.3f}, "
              f"range=[{alpha[:,0].min():.2f}, {alpha[:,0].max():.2f}]")
        print(f"  Alpha (bottom): mean={alpha[:,1].mean():.3f}, std={alpha[:,1].std():.3f}, "
              f"range=[{alpha[:,1].min():.2f}, {alpha[:,1].max():.2f}]")
        print(f"  Alpha (top): mean={alpha[:,2].mean():.3f}, std={alpha[:,2].std():.3f}, "
              f"range=[{alpha[:,2].min():.2f}, {alpha[:,2].max():.2f}]")
        ppo[d]["alpha"] = alpha
        ppo[d]["cfg"] = cfg

    # 2. FFT spectral analysis
    print("\n" + "=" * 60)
    print("FFT Spectral Analysis")
    print("=" * 60)

    fig, axes = plt.subplots(3, 3, figsize=(16, 12))
    for di, d in enumerate([0.75, 1.0, 1.5]):
        alpha = ppo[d]["alpha"]
        for ci, name in enumerate(["Front", "Bottom", "Top"]):
            ax = axes[di, ci]
            signal = alpha[:, ci] - alpha[:, ci].mean()
            fft = np.abs(np.fft.rfft(signal))
            freqs = np.fft.rfftfreq(len(signal), d=1)
            ax.plot(freqs[1:], fft[1:], "b-", alpha=0.8)
            # Mark dominant frequency
            dom_idx = np.argmax(fft[1:]) + 1
            dom_freq = freqs[dom_idx]
            ax.axvline(dom_freq, color="r", linestyle="--", alpha=0.5,
                       label=f"f={dom_freq:.4f}")
            ax.set_title(f"{d}L {name}")
            ax.set_xlabel("Frequency (control steps)")
            ax.set_ylabel("|FFT|")
            ax.legend(fontsize=7)
            print(f"  {d}L {name}: dominant f={dom_freq:.4f} (T={1/dom_freq:.1f} steps)")
    plt.tight_layout()
    plt.savefig(os.path.join(OUT_DIR, "fft_spectra.png"), dpi=150)
    plt.close()
    print(f"  Saved: fft_spectra.png")

    # 3. Action timeseries comparison
    print("\n" + "=" * 60)
    print("Action Timeseries (first 100 steps)")
    print("=" * 60)
    fig, axes = plt.subplots(3, 1, figsize=(14, 10))
    for di, d in enumerate([0.75, 1.0, 1.5]):
        ax = axes[di]
        alpha = ppo[d]["alpha"]
        n_show = min(100, alpha.shape[0])
        ax.plot(alpha[:n_show, 0], label="Front", linewidth=0.8)
        ax.plot(alpha[:n_show, 1], label="Bottom", linewidth=0.8)
        ax.plot(alpha[:n_show, 2], label="Top", linewidth=0.8)
        ax.set_title(f"{d}L — PPO Alpha (non-dim)")
        ax.legend()
        ax.set_ylabel("α = ω/U₀")
    ax.set_xlabel("Step")
    plt.tight_layout()
    plt.savefig(os.path.join(OUT_DIR, "action_timeseries.png"), dpi=150)
    plt.close()
    print(f"  Saved: action_timeseries.png")

    # 4. Action autocorrelation (detects periodicity pattern)
    print("\n" + "=" * 60)
    print("Action Autocorrelation")
    print("=" * 60)
    fig, axes = plt.subplots(3, 3, figsize=(16, 12))
    for di, d in enumerate([0.75, 1.0, 1.5]):
        alpha = ppo[d]["alpha"]
        for ci, name in enumerate(["Front", "Bottom", "Top"]):
            ax = axes[di, ci]
            signal = alpha[:, ci] - alpha[:, ci].mean()
            # Compute autocorrelation manually
            n = len(signal)
            ac = np.correlate(signal, signal, mode="full")[n - 1:]
            ac = ac / ac[0]
            max_lag = min(30, n // 2)
            ax.stem(range(max_lag), ac[:max_lag])
            ax.set_title(f"{d}L {name}")
            ax.set_xlabel("Lag (steps)")
            ax.set_ylabel("ACF")
            # Flag lag-2 autocorrelation
            if max_lag >= 2:
                ac2 = ac[2]
                if abs(ac2) > 0.5:
                    ax.annotate(f"lag-2: {ac2:.2f}", xy=(2, ac2),
                                fontsize=8, color="red")
    plt.tight_layout()
    plt.savefig(os.path.join(OUT_DIR, "autocorrelation.png"), dpi=150)
    plt.close()
    print(f"  Saved: autocorrelation.png")

    # 5. Joint formula evaluation
    print("\n" + "=" * 60)
    print("Joint Formula Range Analysis")
    print("=" * 60)

    for d in [0.75, 1.0, 1.5]:
        cfg = ppo[d]["cfg"]
        sensors = ppo[d]["sensors"]
        forces = ppo[d]["forces"]
        dim = compute_dimless(sensors, forces)
        feat = compute_features(dim, ppo[d]["alpha"] * U0)

        # target_Cd component of joint formula
        # For trained scenes we don't have a separate target recording
        # but target_diameter defines the expected Cd range
        target_diam = cfg["target_diameter"]
        # Approximate target_Cd from the feature data (Cd_tot in uncontrolled would give target-like values)
        # Get the range of Cd_tot
        print(f"\n  {d}L (target diam={target_diam}*L0):")
        print(f"    Cd_tot: mean={feat['Cd_tot'].mean():.4f}, std={feat['Cd_tot'].std():.4f}, "
              f"range=[{feat['Cd_tot'].min():.3f}, {feat['Cd_tot'].max():.3f}]")
        print(f"    Cl_tot: mean={feat['Cl_tot'].mean():.4f}, std={feat['Cl_tot'].std():.4f}, "
              f"range=[{feat['Cl_tot'].min():.3f}, {feat['Cl_tot'].max():.3f}]")
        print(f"    u_a: mean={feat['u_a'].mean():.4f}, std={feat['u_a'].std():.4f}, "
              f"range=[{feat['u_a'].min():.3f}, {feat['u_a'].max():.3f}]")
        print(f"    Cd_rear: mean={feat['Cd_rear'].mean():.4f}, std={feat['Cd_rear'].std():.4f}")

    # 6. Joint formula term breakdown for 0.75L and 1L
    print("\n" + "=" * 60)
    print("Joint Formula Term Decomposition")
    print("=" * 60)

    for d in [0.75, 1.0]:
        cfg = ppo[d]["cfg"]
        sensors = ppo[d]["sensors"]
        forces = ppo[d]["forces"]
        a = ppo[d]["alpha"]
        n_warmup = 2
        dim = compute_dimless(sensors, forces)
        feat = compute_features(dim, a * U0)

        # Joint formula: alpha_F = target_Cd - 5.428 + 0.0098*(du_a_dt + u_a)
        # We don't have target_Cd directly, but Cd_tot from PPO trajectory
        # approximates it. Let's use Cd_tot as a proxy for target_Cd.
        # Joint formula: alpha_F = Cd_tot - (Cd_tot - target_Cd) - 5.428 + 0.0098*(du_a_dt + u_a)
        #                   = target_Cd - 5.428 + 0.0098*(du_a_dt + u_a)

        # Approximate: use Cd_tot mean as proxy
        target_cd_proxy = np.mean(feat["Cd_tot"][n_warmup:])
        du_a_dt_raw = np.zeros_like(feat["u_a"])
        du_a_dt_raw[1:] = feat["u_a"][1:] - feat["u_a"][:-1]

        term_drag = target_cd_proxy - 5.428
        term_phase = 0.0098 * (du_a_dt_raw[n_warmup:] + feat["u_a"][n_warmup:])
        alpha_sr = term_drag + term_phase
        alpha_ppo = a[n_warmup:, 0]

        print(f"\n  {d}L joint formula breakdown (front):")
        print(f"    target_Cd proxy: {target_cd_proxy:.4f}")
        print(f"    drag term (target_Cd - 5.428): {term_drag:.4f} (constant)")
        print(f"    phase term mean: {term_phase.mean():.4f}, std: {term_phase.std():.4f}")
        print(f"    SR alpha mean: {alpha_sr.mean():.4f}, std: {alpha_sr.std():.4f}")
        print(f"    PPO alpha mean: {alpha_ppo.mean():.4f}, std: {alpha_ppo.std():.4f}")
        print(f"    Correlation(SR, PPO): {np.corrcoef(alpha_sr, alpha_ppo)[0,1]:.4f}")
        print(f"    SR contribution: drag={abs(term_drag):.3f}, phase_std={term_phase.std():.3f}")

    print(f"\n=== Done ===")
    print(f"Output: {OUT_DIR}")


if __name__ == "__main__":
    main()
