"""Measure exactly how PixelGen changed after adding one class.

Compares the frozen before-state (experiments/before.pt, before_outputs.npz,
before_meta.json) against the live model built from current source.

Reports:
  1. tensor shapes before -> after
  2. raw parameter diff (how many values differ by index)
  3. semantic diff (how many values differ once rows/cols are aligned by
     class name / word, i.e. the change that actually matters)
  4. vocab index remap (sorted vocab means new words shift old columns)
  5. output equivalence on every pre-existing prompt (bit-exact or not)
  6. what the new prompts do now vs before
"""

import json
from pathlib import Path

import numpy as np
import torch

from pixelgen.model import PixelGen

HERE = Path(__file__).resolve().parent
before_sd = torch.load(HERE / "before.pt", map_location="cpu", weights_only=True)
before_out = np.load(HERE / "before_outputs.npz")
meta = json.load(open(HERE / "before_meta.json"))
b_vocab, b_classes, prompts = meta["vocab"], meta["classes"], meta["prompts"]

m = PixelGen()
after_sd = m.state_dict()
a_vocab, a_classes = m.vocab, m.classes

new_classes = [c for c in a_classes if c not in b_classes]
new_words = [w for w in a_vocab if w not in b_vocab]
print("=" * 70)
print(f"new classes: {new_classes}   new words: {new_words}")
print(f"vocab {len(b_vocab)} -> {len(a_vocab)}   classes {len(b_classes)} -> {len(a_classes)}")

# 1. shapes -------------------------------------------------------------
print("\n[1] tensor shapes")
tot_b = tot_a = 0
for k in after_sd:
    sb, sa = tuple(before_sd[k].shape), tuple(after_sd[k].shape)
    tot_b += before_sd[k].numel(); tot_a += after_sd[k].numel()
    flag = "  <-- grew" if sb != sa else ""
    print(f"  {k:24} {str(sb):>16} -> {str(sa):<16}{flag}")
print(f"  total values {tot_b} -> {tot_a}  (+{tot_a - tot_b})")

# 2. raw diff ------------------------------------------------------------
print("\n[2] raw diff by index (naive view)")
for k in after_sd:
    b, a = before_sd[k], after_sd[k]
    if b.shape == a.shape:
        n = int((b != a).sum())
        print(f"  {k:24} {n:5d} / {a.numel()} values differ")
    else:
        # compare overlapping prefix region
        sl = tuple(slice(0, s) for s in b.shape)
        n = int((b != a[sl]).sum())
        print(f"  {k:24} {n:5d} / {b.numel()} values differ in the old-shaped prefix (+{a.numel()-b.numel()} new)")

# 3. semantic diff (align by names) ---------------------------------------
print("\n[3] semantic diff (rows/cols aligned by class name / word)")
def aligned_linear(name):
    b, a = before_sd[f"{name}.weight"], after_sd[f"{name}.weight"]
    rows = b_classes if name == "word_to_class" else m.ramps
    changed = []
    for ri, r in enumerate(rows):
        for w in b_vocab:
            bv = b[ri, b_vocab[w]].item()
            av = a[a_classes.index(r) if name == "word_to_class" else ri, a_vocab[w]].item()
            if bv != av:
                changed.append((r, w, bv, av))
    return changed
for name in ("word_to_class", "word_to_ramp"):
    ch = aligned_linear(name)
    print(f"  {name:16} old (row,word) entries changed: {len(ch)}  {ch[:5]}")
bb, ab = before_sd["word_to_class.bias"], after_sd["word_to_class.bias"]
print("  word_to_class.bias old-class entries changed:",
      sum(bb[i].item() != ab[a_classes.index(c)].item() for i, c in enumerate(b_classes)))
for bank in ("on_bank", "free_bank"):
    b, a = before_sd[bank], after_sd[bank]
    n = sum(int((b[i] != a[a_classes.index(c)]).sum()) for i, c in enumerate(b_classes))
    print(f"  {bank:16} old-class template cells changed: {n}")
for k in ("shade_bank", "ramp_table", "outline.weight", "palette_conv.weight"):
    print(f"  {k:16} identical: {torch.equal(before_sd[k], after_sd[k])}")

# 4. vocab remap ----------------------------------------------------------
print("\n[4] vocab index remap caused by sorted insertion")
moved = [(w, b_vocab[w], a_vocab[w]) for w in b_vocab if b_vocab[w] != a_vocab[w]]
print(f"  {len(moved)} of {len(b_vocab)} old words changed column index")
print("  first few:", moved[:6])
print("  new word columns:", {w: a_vocab[w] for w in new_words})

# 5. output equivalence ---------------------------------------------------
print("\n[5] output equivalence on pre-existing prompts (3 seeds each)")
same = diff = 0
changed_prompts = []
for i, p in enumerate(prompts):
    b = before_out[f"p{i}"]
    a = m.generate([p] * 3, [0, 1, 2], scale=1)
    if np.array_equal(a, b):
        same += 1
    else:
        diff += 1
        changed_prompts.append(p)
print(f"  identical: {same}   changed: {diff}   -> {changed_prompts}")

# 6. new prompt behaviour ---------------------------------------------------
print("\n[6] what the new-class prompts resolve to now")
ref = {c: (m.generate([c], [0], scale=1)[0].sum(-1) > 0) for c in a_classes}
for p in ["mushroom", "red mushroom", "purple fungus", "toadstool", "dragon", ""]:
    mask = m.generate([p], [0], scale=1)[0].sum(-1) > 0
    cls = [c for c, r in ref.items() if np.array_equal(r, mask)]
    i = prompts.index(p) if p in prompts else None
    was = "n/a"
    if i is not None:
        bm = before_out[f"p{i}"][0].sum(-1) > 0
        was = [c for c, r in ref.items() if np.array_equal(r, bm)]
    print(f"  {p!r:18} before -> {was}   after -> {cls}")

# image ---------------------------------------------------------------------
from PIL import Image
S = 6; cell = 16 * S; pad = 4
rows = [("before", b_classes), ("after", a_classes)]
W = max(len(c) for _, c in rows) * (cell + pad) + pad
img = Image.new("RGB", (W, 2 * (cell + pad) + pad), (29, 43, 83))
for r, (_, classes) in enumerate(rows):
    for c, cls in enumerate(classes):
        img.paste(Image.fromarray(m.generate([cls], [0], scale=S)[0]), (pad + c * (cell + pad), pad + r * (cell + pad)))
img.save(HERE / "before_after.png")
print(f"\nwrote {HERE/'before_after.png'}  (top: old classes rendered by new model; bottom: all classes)")