"""Analyze separated λ_L and λ_H, and their predictive power for failure."""
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from scipy import stats
import glob, os

ROOT = "/home/yurenh2/rrm/research/flossing"
OUT = f"{ROOT}/plots_separate"
os.makedirs(OUT, exist_ok=True)

files = sorted(glob.glob(f"{ROOT}/diag_sep_1k_shard*.npz"))
m = {}
for f in files:
    d = np.load(f)
    for k in d.files: m.setdefault(k, []).append(d[k])
for k in list(m.keys()): m[k] = np.concatenate(m[k], 0)
np.savez_compressed(f"{ROOT}/diag_sep_1k.npz", **m)

N = len(m["exact_correct"])
exact = m["exact_correct"].astype(bool)
acc = exact.mean()
print(f"N={N}  acc={acc:.4f}\n")

def auc(score, target):
    return stats.mannwhitneyu(score[target], score[~target], alternative="greater").statistic / (target.sum() * (~target).sum())

results = {}
for name in ["lyap_joint", "lyap_L", "lyap_H"]:
    ls = m[name]
    print(f"== {name} ==")
    print(f"{'i':>3} {'mean_all':>10} {'mean_succ':>10} {'mean_fail':>10} {'Δ':>8} {'AUROC':>7} {'cohen_d':>8}")
    for i in range(ls.shape[1]):
        li = ls[:, i]; s_ = li[exact]; f_ = li[~exact]
        pooled = np.sqrt(((s_.size-1)*s_.var() + (f_.size-1)*f_.var()) / (s_.size+f_.size-2))
        dc = (s_.mean() - f_.mean()) / pooled
        delta = f_.mean() - s_.mean()
        au = auc(li, ~exact)
        print(f"{i+1:>3} {li.mean():+10.4f} {s_.mean():+10.4f} {f_.mean():+10.4f} {delta:+8.4f} {au:>7.3f} {abs(dc):>8.3f}")
    results[name] = ls
    print()

# AUROC comparison
print("=== Single-feature AUROC predicting failure ===")
for name in ["lyap_joint", "lyap_L", "lyap_H"]:
    ls = m[name]
    for label, score in [("λ_1", ls[:, 0]), ("mean", ls.mean(1)), ("λ_K", ls[:, -1])]:
        print(f"  {name}.{label:6s}: AUROC = {auc(score, ~exact):.4f}")

# --- plot 1: three spectra overlay ---
fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), sharex=True)
for ax, name in zip(axes, ["lyap_joint", "lyap_L", "lyap_H"]):
    ls = m[name]
    K = ls.shape[1]
    x = np.arange(1, K+1)
    s_ = ls[exact]; f_ = ls[~exact]
    ax.fill_between(x, s_.mean(0)-s_.std(0), s_.mean(0)+s_.std(0), color="C0", alpha=0.25)
    ax.plot(x, s_.mean(0), "C0-o", label="succ", lw=2)
    ax.fill_between(x, f_.mean(0)-f_.std(0), f_.mean(0)+f_.std(0), color="C3", alpha=0.25)
    ax.plot(x, f_.mean(0), "C3-o", label="fail", lw=2)
    ax.axhline(0, color="k", ls="--", lw=1, label="critical")
    ax.set_xlabel("λ index")
    ax.set_ylabel(r"$\lambda_i$ (per cycle)")
    ax.set_title(name)
    ax.legend()
    ax.grid(alpha=0.3)
fig.suptitle(f"HRM step_26040 — separated Lyapunov spectra (N={N}, acc={acc:.2%})")
fig.tight_layout()
fig.savefig(f"{OUT}/three_spectra.png", dpi=130)
plt.close()

# --- plot 2: λ_H_1 histogram (the differentiator) ---
fig, ax = plt.subplots(1, 1, figsize=(7, 4))
ls = m["lyap_H"][:, 0]
lo = ls.min(); hi = ls.max()
bins = np.linspace(lo, hi, 60)
ax.hist(ls[exact], bins=bins, alpha=0.55, label=f"succ", color="C0", density=True)
ax.hist(ls[~exact], bins=bins, alpha=0.55, label=f"fail", color="C3", density=True)
ax.axvline(0, color="k", ls="--", lw=1)
ax.axvline(ls[exact].mean(), color="C0", ls=":", lw=1)
ax.axvline(ls[~exact].mean(), color="C3", ls=":", lw=1)
ax.set_xlabel(r"$\lambda_{H,1}$ — H subsystem top Lyapunov")
ax.set_ylabel("density")
ax.set_title(f"H-module Lyapunov is the differentiator (AUROC={auc(ls, ~exact):.3f})")
ax.legend()
fig.tight_layout()
fig.savefig(f"{OUT}/lyap_H_hist.png", dpi=130)
plt.close()

# --- plot 3: λ_L_1 histogram (NOT a differentiator) ---
fig, ax = plt.subplots(1, 1, figsize=(7, 4))
ls = m["lyap_L"][:, 0]
lo = ls.min(); hi = ls.max()
bins = np.linspace(lo, hi, 60)
ax.hist(ls[exact], bins=bins, alpha=0.55, label=f"succ", color="C0", density=True)
ax.hist(ls[~exact], bins=bins, alpha=0.55, label=f"fail", color="C3", density=True)
ax.axvline(ls[exact].mean(), color="C0", ls=":", lw=1)
ax.axvline(ls[~exact].mean(), color="C3", ls=":", lw=1)
ax.set_xlabel(r"$\lambda_{L,1}$ — L subsystem top Lyapunov")
ax.set_ylabel("density")
ax.set_title(f"L-module Lyapunov is NOT a differentiator (AUROC={auc(ls, ~exact):.3f})")
ax.legend()
fig.tight_layout()
fig.savefig(f"{OUT}/lyap_L_hist.png", dpi=130)
plt.close()

# --- plot 4: scatter λ_L vs λ_H, colored by success/failure ---
fig, ax = plt.subplots(1, 1, figsize=(7, 6))
xL = m["lyap_L"][:, 0]
xH = m["lyap_H"][:, 0]
ax.scatter(xL[exact], xH[exact], s=5, alpha=0.4, color="C0", label="succ")
ax.scatter(xL[~exact], xH[~exact], s=5, alpha=0.4, color="C3", label="fail")
ax.axhline(0, color="k", ls="--", lw=0.5)
ax.set_xlabel(r"$\lambda_{L,1}$ (L subsystem)")
ax.set_ylabel(r"$\lambda_{H,1}$ (H subsystem)")
ax.set_title("Success/failure separates along the λ_H axis")
ax.legend()
ax.grid(alpha=0.3)
fig.tight_layout()
fig.savefig(f"{OUT}/lyap_L_vs_H_scatter.png", dpi=130)
plt.close()

print(f"\nplots → {OUT}/")