#!/usr/bin/env python3
"""Direct test of the mechanism behind the null interventions: does Sigma
carry lifetime-specific information at all?

If Sigma_t were a useful state it would differ between lifetimes at the same
event index. If it has collapsed to a learned constant, the between-lifetime
spread will be negligible against the state's own scale.
"""
import sys
from pathlib import Path
import numpy as np
import torch

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))
sys.path.insert(0, str(Path(__file__).resolve().parent))
from pns.common import atomic_write_json, eval_root, shards_root  # noqa
from pns.train.loader import iter_eval_batches  # noqa
from evaluate import load_model, to_gpu  # noqa

dev = "cuda"
model, _ = load_model(sys.argv[1] if len(sys.argv) > 1 else "PNSR_K4_fin", dev)
d = model.cfg.d
PROBE_T = [31, 63, 127, 191, 255]
snaps = {t: [] for t in PROBE_T}

with torch.no_grad():
    for b in iter_eval_batches("val", shards_root(), 32, 32):
        g = to_gpu(b, dev)
        B, L = g["etype"].shape
        state = model.initial_state(B, dev)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            h = model.recenc.static(g["rec_val_toks"], g["rec_key_toks"], g["rec_store"],
                                    g["rec_kind"], g["rec_key"], g["rec_ent"])
            for t in range(L):
                live = g["live"][:, t]; mask = live >= 0; rows = live.clamp(min=0)
                bank = torch.gather(h, 1, rows.unsqueeze(-1).expand(-1, -1, d))
                bank = model.recenc.finalize(
                    bank, (t - torch.gather(g["rec_birth"], 1, rows)).clamp(min=0))
                state, _ = model.step(state, g["tok"][:, t], g["etype"][:, t],
                                      g["dt"][:, t], bank * mask.unsqueeze(-1), mask)
                if t in snaps:
                    snaps[t].append(state.float().cpu())
        break

rep = {}
for t, lst in snaps.items():
    S = torch.cat(lst)                       # [B, 120, d]
    mean = S.mean(0, keepdim=True)
    between = (S - mean).norm(dim=-1).mean().item()      # spread across lifetimes
    scale = S.norm(dim=-1).mean().item()                 # typical slot norm
    flat = S.flatten(1)
    fn = torch.nn.functional.normalize(flat, dim=-1)
    cos = (fn @ fn.T)
    off = cos[~torch.eye(len(fn), dtype=torch.bool)].mean().item()
    rep[str(t)] = dict(slot_norm=round(scale, 4),
                       between_lifetime_spread=round(between, 4),
                       spread_over_norm=round(between / max(scale, 1e-9), 4),
                       mean_pairwise_cosine=round(off, 4))
    print(t, rep[str(t)], flush=True)

# drift of the mean state over time: does it even change with t?
ms = {t: torch.cat(lst).mean(0).flatten() for t, lst in snaps.items()}
ks = sorted(ms)
rep["mean_state_cosine_t31_vs_t255"] = round(float(
    torch.nn.functional.cosine_similarity(ms[ks[0]], ms[ks[-1]], dim=0)), 4)
run_name = sys.argv[1] if len(sys.argv) > 1 else "PNSR_K4_fin"
atomic_write_json(eval_root() / f"SIGMA_COLLAPSE_{run_name}.json", rep)
print("mean-state cosine t=31 vs t=255:", rep["mean_state_cosine_t31_vs_t255"])