#!/usr/bin/env python3
"""Analyze MOSSStripeProbe recordings against final-image banding scores.

  comfyui_venv/bin/python workflows/moss/probe_analysis.py output/moss/probe

Expects <setup>_0000N_.png images and probe_<setup>_s<seed>.npz recordings for setups sft, fix, base,
seeds 11-16 (image N = seed 10 + N). Writes probe_*.png figures and probe_summary.json next to them.
"""
import json, os, sys
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import bandviz

D = sys.argv[1]
SETUPS = [("sft", "Full SFT (bands)"), ("fix", "SFT + 21-49 LoRA @ -0.5 (Kijai fix)"), ("base", "Base")]
SEEDS = range(11, 17)
KIND = ["lens", "layer adds (lens)", "residual stream"]
SURFACE, INK, INK2, MUTED, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#898781", "#e1e0d9"


def rel(p):
    """Spectrum power relative to its median over non-DC bins, in log10."""
    return np.log10(p / np.median(p[..., 1:], axis=-1, keepdims=True))


runs = []
for setup, _ in SETUPS:
    for seed in SEEDS:
        z = np.load(os.path.join(D, f"probe_{setup}_s{seed}.npz"), allow_pickle=True)
        cond = np.array([c == "[0]" for c in z["cou"]])
        spec = np.stack([np.asarray(s[0], dtype=np.float64) for s in z["spec"]])[cond]  # (steps, layers, 3, bins)
        img = os.path.join(D, f"{setup}_{seed - 10:05d}_.png")
        runs.append({"setup": setup, "seed": seed, "spec": spec, "sigma": z["sigma"][cond], "score": float(bandviz.analyze(img)["score"])})

scores = np.array([r["score"] for r in runs])
R = np.stack([rel(r["spec"]) for r in runs])  # (runs, steps, layers, 3, bins)
nb = R.shape[-1]
W_TOK = (nb - 1) * 2


def corr(x, y):
    x = x - x.mean(0); y = y - y.mean()
    return (x * y.reshape(-1, *[1] * (x.ndim - 1))).sum(0) / np.sqrt((x ** 2).sum(0) * (y ** 2).sum() + 1e-12)


# 1. Which token frequency carries the banding: final-layer lens (= what the model outputs), late steps.
final = R[:, -8:, -1, 0, :].mean(1)  # (runs, bins)
r_bin = corr(final, scores)
sig = int(np.argmax(r_bin[1:])) + 1
print("image scores:", {f"{r['setup']}_s{r['seed']}": round(r["score"], 1) for r in runs})
print("signature bin", sig, f"({sig} cycles across {W_TOK} tokens), r =", round(float(r_bin[sig]), 3))

fig, ax = plt.subplots(1, 2, figsize=(15, 4.6), facecolor=SURFACE)
for a in ax:
    a.set_facecolor(SURFACE); a.grid(True, color=GRID); a.set_axisbelow(True)
    for s in ("top", "right"): a.spines[s].set_visible(False)
ax[0].bar(range(nb), r_bin, color="#2a78d6")
ax[0].bar([sig], [r_bin[sig]], color=INK)
ax[0].set_title("How well each stripe frequency in the model's output predicts the final image's banding score\n(correlation across all 18 runs; black = strongest)", fontsize=11, loc="left", color=INK)
ax[0].set_xlabel(f"stripe frequency in the token grid (cycles across {W_TOK} tokens)", color=MUTED)
ax[0].set_ylabel("correlation r", color=MUTED)
for (setup, name), c in zip(SETUPS, ["#d03b3b", "#2a78d6", "#898781"]):
    ax[1].plot(range(1, nb), final[[r["setup"] == setup for r in runs]].mean(0)[1:], color=c, linewidth=2, label=name)
ax[1].axvline(sig, color=INK, linewidth=1, linestyle=(0, (4, 3)))
ax[1].set_title("Output stripe spectrum, mean over 6 seeds (log power ÷ median)", fontsize=11, loc="left", color=INK)
ax[1].set_xlabel(f"cycles across {W_TOK} tokens", color=MUTED); ax[1].legend(frameon=False)
fig.tight_layout(); fig.savefig(os.path.join(D, "probe_1_signature.png"), dpi=110, facecolor=SURFACE); plt.close(fig)

# 2. Heatmaps: layer x step of signature power, per setup, for lens state and per-layer additions.
steps = R.shape[1]
fig, axes = plt.subplots(2, 3, figsize=(17, 10), facecolor=SURFACE, sharey=True)
for row, kind in enumerate([0, 1]):
    vals = [R[[r["setup"] == s for r in runs]][..., kind, sig].mean(0) for s, _ in SETUPS]  # (steps, layers)
    vmin, vmax = min(v.min() for v in vals), max(v.max() for v in vals)
    for col, ((s, name), v) in enumerate(zip(SETUPS, vals)):
        a = axes[row, col]
        im = a.imshow(v.T, aspect="auto", origin="lower", cmap="magma", vmin=vmin, vmax=vmax, extent=(0.5, steps + 0.5, -0.5, v.shape[1] - 0.5))
        a.set_title(f"{name}\n{KIND[kind]}", fontsize=10.5, loc="left", color=INK)
        a.set_xlabel("sampling step (noisy → clean)", color=MUTED)
        if col == 0: a.set_ylabel("layer", color=MUTED)
    fig.colorbar(im, ax=axes[row].tolist(), shrink=0.85, label=f"log stripe power at bin {sig} (÷ median)")
fig.suptitle(f"Where the stripe frequency lives: layer × step, mean of 6 seeds (top: what each layer's output would decode to; bottom: what each layer adds)", x=0.01, ha="left", fontsize=12.5, color=INK)
fig.savefig(os.path.join(D, "probe_2_heatmaps.png"), dpi=110, facecolor=SURFACE); plt.close(fig)

# 3. Difference and correlation maps: SFT minus base, fix minus SFT, and r(cell, image score) across all runs.
sel = lambda s: np.array([r["setup"] == s for r in runs])
add = R[..., 1, sig]  # (runs, steps, layers)
maps = [("Full SFT minus Base: extra stripe energy each layer adds", add[sel("sft")].mean(0) - add[sel("base")].mean(0), "RdBu_r", True),
        ("Kijai fix minus Full SFT: what the fix removes", add[sel("fix")].mean(0) - add[sel("sft")].mean(0), "RdBu_r", True),
        ("Correlation with final banding score (all 18 runs)", corr(add, scores), "RdBu_r", True)]
fig, axes = plt.subplots(1, 3, figsize=(17, 6), facecolor=SURFACE, sharey=True)
for a, (title, v, cmap, sym) in zip(axes, maps):
    lim = np.nanmax(np.abs(v))
    im = a.imshow(v.T, aspect="auto", origin="lower", cmap=cmap, vmin=-lim, vmax=lim, extent=(0.5, steps + 0.5, -0.5, v.shape[1] - 0.5))
    a.set_title(title, fontsize=10.5, loc="left", color=INK); a.set_xlabel("sampling step (noisy → clean)", color=MUTED)
    fig.colorbar(im, ax=a, shrink=0.85)
axes[0].set_ylabel("layer", color=MUTED)
fig.suptitle(f"The dials: per-layer stripe additions at bin {sig}", x=0.01, ha="left", fontsize=12.5, color=INK)
fig.tight_layout(); fig.savefig(os.path.join(D, "probe_3_dials.png"), dpi=110, facecolor=SURFACE); plt.close(fig)

# 4. Per-layer profile: correlation with score, averaged over steps, plus per-setup addition profiles.
r_layer = corr(add.mean(1), scores)  # (layers,)
fig, ax = plt.subplots(1, 2, figsize=(15, 4.6), facecolor=SURFACE)
for a in ax:
    a.set_facecolor(SURFACE); a.grid(True, color=GRID); a.set_axisbelow(True)
    for s in ("top", "right"): a.spines[s].set_visible(False)
ax[0].bar(range(len(r_layer)), r_layer, color="#2a78d6")
ax[0].set_title("Per layer: how well its stripe addition (all steps) predicts the final banding score", fontsize=11, loc="left", color=INK)
ax[0].set_xlabel("layer", color=MUTED); ax[0].set_ylabel("correlation r", color=MUTED)
for (s, name), c in zip(SETUPS, ["#d03b3b", "#2a78d6", "#898781"]):
    ax[1].plot(add[sel(s)].mean((0, 1)), color=c, linewidth=2, label=name)
ax[1].set_title(f"Mean stripe addition per layer at bin {sig} (log ÷ median)", fontsize=11, loc="left", color=INK)
ax[1].set_xlabel("layer", color=MUTED); ax[1].legend(frameon=False)
fig.tight_layout(); fig.savefig(os.path.join(D, "probe_4_layers.png"), dpi=110, facecolor=SURFACE); plt.close(fig)

top = np.argsort(-r_layer)[:8]
json.dump({"signature_bin": sig, "r_bin": r_bin.round(3).tolist(), "r_layer": r_layer.round(3).tolist(),
           "top_layers": [(int(i), round(float(r_layer[i]), 3)) for i in top],
           "scores": {f"{r['setup']}_s{r['seed']}": round(r["score"], 1) for r in runs}},
          open(os.path.join(D, "probe_summary.json"), "w"), indent=1)
print("top layers by r:", [(int(i), round(float(r_layer[i]), 2)) for i in top])