#!/usr/bin/env python3
"""Diagnose Illusion PPO data and PySR fit quality.

Checks:
  1. PPO data integrity (sensor ranges, action ranges, similarity match)
  2. G-equivariance of PPO policy (should front be odd? rear shared-head?)
  3. Front constant behavior at 0.75L — is this real physics or fitting artifact?
  4. PySR front no-bias constraint — does the fitting actually enforce it?
  5. Phase portraits (sensor vs sensor) to check for periodic signals
  6. Action timeseries over multiple periods (not truncated)

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

import json
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

U0 = 0.01
DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "data", "illusion")
FIG_DIR = os.path.join(os.path.dirname(__file__), "..", "data", "diagnostics")
os.makedirs(FIG_DIR, exist_ok=True)


def load_scene(scene_name):
    path = os.path.join(DATA_DIR, scene_name, "controlled.npz")
    cfg = get_scene(scene_name)
    npz = np.load(path, allow_pickle=True)
    sensors = npz["sensors"].astype(np.float64)
    forces = npz["forces"].astype(np.float64)
    actions_norm = npz["actions"].astype(np.float64)
    scale = cfg["action_scale"]
    bias = np.array(cfg["action_bias"], dtype=np.float64)
    actions_phys = (actions_norm * scale + bias) * cfg["u0"]
    alpha = actions_phys / cfg["u0"]
    return sensors, forces, actions_norm, alpha, cfg


def check_similarity_vs_training(scene_name):
    sensors, forces, actions_norm, alpha, cfg = load_scene(scene_name)
    result_path = os.path.join(DATA_DIR, scene_name, "result.json")
    if os.path.exists(result_path):
        result = json.load(open(result_path))
        sim = result.get("similarity", None)
        if sim is not None:
            print(f"    PPO similarity (controlled.npz): {sim:.4f}")
        else:
            print(f"    result.json keys: {list(result.keys())}")
    else:
        print(f"    WARNING: no result.json for {scene_name}")
        return

    target_path = os.path.join(DATA_DIR, scene_name, "target.npz")
    if os.path.exists(target_path):
        target = np.load(target_path)["target_states"]
        T = min(len(sensors), len(target))
        for i, name in enumerate(["S0_ux", "S0_uy", "S1_ux", "S1_uy", "S2_ux", "S2_uy"]):
            sens = sensors[:T, i]
            targ = target[:T, 2+i] if target.shape[1] >= 8 else sensors[:T, i]
            if i == 0:
                print(f"    {name}: ctl_mean={sens.mean():.6f}, ctl_std={sens.std():.6f}")
            if i <= 1:
                pass  # already printed
        # Print summary
        print(f"    Controlled sensors RMS: {np.sqrt(np.mean(sensors[:T,:]**2)):.6f}")
        if np.any(np.isnan(sensors)):
            print(f"    ERROR: NaN in sensors!")
        if np.any(np.isnan(forces)):
            print(f"    ERROR: NaN in forces!")


def check_actions(alpha, scene_name):
    print(f"\n  {scene_name} alpha (non-dim omega/U0):")
    for i, name in enumerate(["Front", "Bottom", "Top"]):
        a = alpha[:, i]
        print(f"    {name}: mu={a.mean():+.3f}, sigma={a.std():.3f}, min={a.min():+.2f}, max={a.max():+.2f}")
        sat_rate = np.mean(np.abs(a) > 7.9)
        if sat_rate > 0.05:
            print(f"      WARNING: {sat_rate*100:.1f}% near saturation")

    front = alpha[:, 0]
    if front.std() > 0 and abs(front.mean()) > 0:
        rel_std = front.std() / (abs(front.mean()) + 1e-6)
        if rel_std < 0.3:
            print(f"    NOTE: Front nearly constant (relative std = {rel_std:.2f})")


def check_gee_equivariance(alpha, sensors, forces, scene_name):
    from SR_analysis.utils.g_operator import apply_G_raw

    T = min(len(alpha), 150)
    alpha_front = alpha[:, 0]
    u_a_simple = (sensors[:, 0] - sensors[:, 4]) / 2  # s0_ux - s2_ux

    u_a_pos = u_a_simple[2:T] > np.percentile(u_a_simple[2:T], 80)
    u_a_neg = u_a_simple[2:T] < np.percentile(u_a_simple[2:T], 20)

    front_pos = alpha_front[2:T][u_a_pos]
    front_neg = alpha_front[2:T][u_a_neg]

    if len(front_pos) > 0 and len(front_neg) > 0:
        print(f"  {scene_name}:")
        print(f"    Front when u_a > P80 (top fast, bottom slow): mu={front_pos.mean():+.3f}")
        print(f"    Front when u_a < P20 (bottom fast, top slow): mu={front_neg.mean():+.3f}")
        delta = abs(front_pos.mean() - front_neg.mean())
        scale = max(abs(front_pos.mean()), abs(front_neg.mean()), 1.0)
        if delta < 0.5 * scale:
            print(f"    → Front nearly symmetric to asymmetry (G-odd behavior weak)")
        else:
            print(f"    → Front responds to asymmetry (consistent with G-odd)")
    else:
        print(f"  {scene_name}: not enough data for G-check")


def plot_phase_portraits(sensors, scene_name):
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    names = [("Top", 0, 1), ("Center", 2, 3), ("Bottom", 4, 5)]
    for ax, (label, ui, vi) in zip(axes, names):
        ax.plot(sensors[:, ui], sensors[:, vi], linewidth=0.3, alpha=0.7)
        ax.set_title(f"{scene_name} — {label} Sensor")
        ax.set_xlabel("ux (lattice)")
        ax.set_ylabel("uy (lattice)")
        ax.set_aspect("equal")
    plt.tight_layout()
    out = os.path.join(FIG_DIR, f"phase_portrait_{scene_name}.png")
    plt.savefig(out, dpi=120)
    plt.close()
    print(f"  Saved: {os.path.basename(out)}")


def plot_full_timeseries(alpha, scene_name, n_steps=200):
    n_plot = min(n_steps, alpha.shape[0])
    fig, axes = plt.subplots(3, 1, figsize=(16, 10), sharex=True)
    colors = ["blue", "green", "red"]
    names = ["Front", "Bottom", "Top"]
    for ax, ci, name, c in zip(axes, range(3), names, colors):
        ax.plot(range(n_plot), alpha[:n_plot, ci], color=c, linewidth=0.6)
        ax.set_ylabel(f"{name}")
        ax.axhline(0, color="gray", linestyle=":", alpha=0.3)
        ax.grid(True, alpha=0.3)
        ax.set_ylim(alpha[:n_plot, ci].min() - 0.5, alpha[:n_plot, ci].max() + 0.5)
    axes[-1].set_xlabel("Step")
    fig.suptitle(f"{scene_name} — Full Action Timeseries ({n_plot} steps)")
    plt.tight_layout()
    out = os.path.join(FIG_DIR, f"timeseries_full_{scene_name}.png")
    plt.savefig(out, dpi=120)
    plt.close()
    print(f"  Saved: {os.path.basename(out)}")


def check_fit_enforces_nobias():
    formulas_dir = os.path.join(os.path.dirname(__file__), "..", "validate", "results")
    for scene_label in ["illusion_0.75L", "illusion_1L", "illusion_joint"]:
        fpath = os.path.join(formulas_dir, f"pysr_{scene_label}_front.json")
        if not os.path.exists(fpath):
            # try results/formulas
            fpath = os.path.join(os.path.dirname(__file__), "..", "results", "formulas", f"{scene_label}_front.json")
        if not os.path.exists(fpath):
            print(f"\n  {scene_label}: formula file not found")
            continue
        formula = json.load(open(fpath))
        sym = formula.get("best_sympy", "")
        score = formula.get("best_score", formula.get("r2_score", "?"))
        keys = formula.get("feature_names", formula.get("feature_keys", []))
        print(f"\n  {scene_label} front:")
        print(f"    formula: {sym}")
        print(f"    R2: {score}")
        print(f"    features ({len(keys)}): {keys[:8]}{'...' if len(keys)>8 else ''}")

        import re
        consts = re.findall(r'[-+]?\s*\d+\.\d+', sym)
        if consts:
            print(f"    standalone constants: {consts}")
        if "bias" in sym.lower():
            print(f"    WARNING: 'bias' found in front formula string!")
        else:
            print(f"    ✓ No 'bias' keyword in formula")


def main():
    scenes = ["illusion_0.75L", "illusion_1L", "illusion_1.5L"]

    print("=" * 60)
    print("DIAGNOSTIC 1: PPO Data Integrity")
    print("=" * 60)
    for sn in scenes:
        try:
            _, _, _, alpha, _ = load_scene(sn)
            check_similarity_vs_training(sn)
            check_actions(alpha, sn)
        except Exception as e:
            print(f"\n  ERROR loading {sn}: {e}")
            import traceback; traceback.print_exc()

    print("\n" + "=" * 60)
    print("DIAGNOSTIC 2: G-Equivariance Check")
    print("=" * 60)
    for sn in scenes:
        sensors, forces, _, alpha, _ = load_scene(sn)
        check_gee_equivariance(alpha, sensors, forces, sn)

    print("\n" + "=" * 60)
    print("DIAGNOSTIC 3: Phase Portraits + Timeseries")
    print("=" * 60)
    for sn in scenes:
        sensors, _, _, alpha, _ = load_scene(sn)
        plot_phase_portraits(sensors, sn)
        plot_full_timeseries(alpha, sn)

    print("\n" + "=" * 60)
    print("DIAGNOSTIC 4: Fit Pipeline Audit")
    print("=" * 60)
    check_fit_enforces_nobias()

    print(f"\nDone. Figures in {FIG_DIR}")


if __name__ == "__main__":
    main()
