"""Weight-space spectral audit: which loop-gain component separates the plain vs cent
lineages before the 196k ceiling collapse? Per-block sigma of Wq/Wk/Wv/proj/w1/w3/w2 and
rms of qk-norm & post-norm gains, at matched steps. Pure CPU, reads the 5k ckpt ladders."""
import torch, sys

STEPS_BOTH = [150000, 170000, 185000, 190000, 195000, 200000]
STEPS_PLAIN_EXTRA = [210000, 220000, 230000]
L = 12

def sig(w):
    return float(torch.linalg.svdvals(w.float())[0])

def rms(g):
    return float(g.float().pow(2).mean().sqrt())

def audit(path):
    ck = torch.load(path, map_location='cpu', weights_only=False)
    b = ck['blocks']
    rows = []
    for l in range(L):
        qkv = b[f'{l}.attn.qkv.weight'].float()
        wq, wk, wv = qkv[0:512], qkv[512:1024], qkv[1024:1536]
        rows.append(dict(
            q=sig(wq), k=sig(wk), v=sig(wv), pr=sig(b[f'{l}.attn.proj.weight']),
            w1=sig(b[f'{l}.ff.w1.weight']), w3=sig(b[f'{l}.ff.w3.weight']),
            w2=sig(b[f'{l}.ff.w2.weight']),
            qg=rms(b[f'{l}.attn.qn.g']), kg=rms(b[f'{l}.attn.kn.g']),
            na=rms(b[f'{l}.na.g']), nf=rms(b[f'{l}.nf.g'])))
    wout = sig(ck['wout'])
    return rows, wout

def summarize(tag, step, rows, wout):
    # aggregates: mean over blocks, top-half mean (6-11), max block, plus the derived products
    def agg(key, blks):
        vals = [rows[l][key] for l in blks]
        return sum(vals) / len(vals)
    bot, top = range(0, 6), range(6, 12)
    logit = [rows[l]['qg'] * rows[l]['kg'] for l in range(L)]           # qk-norm logit scale
    vpath = [rows[l]['v'] * rows[l]['pr'] for l in range(L)]            # attn value-path gain
    fpath = [max(rows[l]['w1'], rows[l]['w3']) * rows[l]['w2'] for l in range(L)]
    print(f'{tag} s{step//1000:>3}k | logit bot {agg("qg",bot)*agg("kg",bot):6.3f} top {agg("qg",top)*agg("kg",top):6.3f} max {max(logit):6.3f}(b{logit.index(max(logit))}) '
          f'| vpath top {sum(vpath[6:])/6:7.2f} max {max(vpath):7.2f}(b{vpath.index(max(vpath))}) '
          f'| ffn top {sum(fpath[6:])/6:7.2f} max {max(fpath):7.2f}(b{fpath.index(max(fpath))}) '
          f'| na top {agg("na",top):5.3f} nf top {agg("nf",top):5.3f} | wout {wout:6.1f}', flush=True)
    return dict(logit=logit, vpath=vpath, fpath=fpath)

res = {}
for lineage, steps in [('plain', STEPS_BOTH + STEPS_PLAIN_EXTRA), ('cent', STEPS_BOTH + STEPS_PLAIN_EXTRA)]:
    for s in steps:
        p = f'runs/fw72m_{lineage}_s{s}.pt'
        try:
            rows, wout = audit(p)
        except FileNotFoundError:
            print(f'{lineage} s{s}: MISSING', flush=True); continue
        res[(lineage, s)] = summarize(lineage, s, rows, wout)

# per-block divergence table at 195k (last matched healthy step)
if ('plain', 195000) in res and ('cent', 195000) in res:
    print('\nper-block plain/cent ratio at 195k (the pre-collapse fingerprint):', flush=True)
    p, c = res[('plain', 195000)], res[('cent', 195000)]
    print('blk  logit-ratio  vpath-ratio  ffn-ratio')
    for l in range(L):
        print(f'{l:3d}  {p["logit"][l]/c["logit"][l]:10.3f}  {p["vpath"][l]/c["vpath"][l]:10.3f}  {p["fpath"][l]/c["fpath"][l]:9.3f}')
print('AUDIT_DONE', flush=True)