import os, time, math, json, datetime, pathlib, glob, re, hashlib
import numpy as np, pandas as pd, torch, torch.nn as nn, torch.nn.functional as F
import matplotlib.pyplot as plt
from torchvision import datasets

DEV = "mps" if torch.backends.mps.is_available() else ("cuda" if torch.cuda.is_available() else "cpu")
CACHE = pathlib.Path(".lab_cache"); CACHE.mkdir(exist_ok=True)
# LAB_SMOKE=1 shrinks every run to a few steps to check the notebook end to end; the published run has it unset
SMOKE = os.environ.get("LAB_SMOKE") == "1"
SCALE = 0.01 if SMOKE else 1.0
steps_ = lambda n: max(5, int(n * SCALE))

INK, LINE = "#2a2a2a", "#d9d3c7"
BLUE, EMBER, FOREST, GRAY, PLUM, GOLD = "#3b6ea5", "#c4521e", "#2e7d5b", "#9a9384", "#7a4f8c", "#b8860b"
plt.rcParams.update({"font.size": 9, "axes.edgecolor": LINE, "axes.spines.top": False,
                     "axes.spines.right": False, "figure.dpi": 120})
RESULTS = {"meta": {"executed": datetime.datetime.now().isoformat(timespec="seconds"), "device": DEV,
                    "torch": torch.__version__, "smoke": SMOKE}}
T0 = time.time()
print(RESULTS["meta"])

# CIFAR-10 (and CIFAR-100 images, used only as unlabeled "other" data), normalized, on the device
DATA = str(CACHE / "data")
def load(cls, train):
    ds = cls(DATA, train=train, download=True)
    return torch.tensor(ds.data).permute(0, 3, 1, 2).float().div(255), torch.tensor(ds.targets), ds.classes

Xtr, Ytr, CLASSES = load(datasets.CIFAR10, True)
Xte, Yte, _ = load(datasets.CIFAR10, False)
X100, _, _ = load(datasets.CIFAR100, True)
MEAN, STD = Xtr.mean((0, 2, 3), keepdim=True), Xtr.std((0, 2, 3), keepdim=True)
RAW_TE = (Xte * 255).round().byte()
norm = lambda x: ((x - MEAN) / STD).to(DEV)
Xtr, Xte, X100 = norm(Xtr), norm(Xte), norm(X100)
Ytr, Yte = Ytr.to(DEV), Yte.to(DEV)
K = 10
print(CLASSES, tuple(Xtr.shape), tuple(Xte.shape), tuple(X100.shape))

def cnn(w, depth=3):
    layers, c = [], 3
    for i in range(depth):
        layers += [nn.Conv2d(c, w * 2 ** i, 3, padding=1), nn.BatchNorm2d(w * 2 ** i), nn.ReLU(), nn.MaxPool2d(2)]
        c = w * 2 ** i
    return nn.Sequential(*layers, nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(c, K))

class Block(nn.Module):
    def __init__(self, cin, cout, stride):
        super().__init__()
        self.c1, self.b1 = nn.Conv2d(cin, cout, 3, stride, 1, bias=False), nn.BatchNorm2d(cout)
        self.c2, self.b2 = nn.Conv2d(cout, cout, 3, 1, 1, bias=False), nn.BatchNorm2d(cout)
        self.sc = (nn.Sequential() if stride == 1 and cin == cout else
                   nn.Sequential(nn.Conv2d(cin, cout, 1, stride, bias=False), nn.BatchNorm2d(cout)))
    def forward(self, x):
        return F.relu(self.b2(self.c2(F.relu(self.b1(self.c1(x))))) + self.sc(x))

def resnet(w, n=3):
    layers, c = [nn.Conv2d(3, w, 3, 1, 1, bias=False), nn.BatchNorm2d(w), nn.ReLU()], w
    for i, mult in enumerate([1, 2, 4]):
        for j in range(n):
            layers.append(Block(c, w * mult, 2 if (j == 0 and i > 0) else 1)); c = w * mult
    return nn.Sequential(*layers, nn.AdaptiveAvgPool2d(1), nn.Flatten(), nn.Linear(c, K))

ARCH = {"cnn8": lambda: cnn(8), "cnn16": lambda: cnn(16), "cnn32": lambda: cnn(32),
        "resnet16": lambda: resnet(16), "resnet32": lambda: resnet(32), "resnet64": lambda: resnet(64)}
n_params = lambda m: sum(p.numel() for p in m.parameters())
print({k: f"{n_params(f()):,}" for k, f in ARCH.items()})

def crop_flip(x, g):
    B = len(x)
    p = F.pad(x, (4, 4, 4, 4), mode="reflect")
    i = torch.randint(0, 9, (B,), generator=g).to(DEV); j = torch.randint(0, 9, (B,), generator=g).to(DEV)
    ar = torch.arange(32, device=DEV)
    out = p[torch.arange(B, device=DEV)[:, None, None], :, (i[:, None] + ar)[:, :, None], (j[:, None] + ar)[:, None, :]]
    out = out.permute(0, 3, 1, 2)
    flip = (torch.rand(B, generator=g) < 0.5).to(DEV)
    return torch.where(flip[:, None, None, None], out.flip(3), out)

def fit(model, X, loss_fn, steps, lr=0.05, wd=5e-4, bs=256, seed=0, teacher_aug=False):
    """loss_fn(logits, idx, flip) -> loss. Returns the model and the wall time.

    `model` is either a module or a factory returning one. A factory is called *after* the seed is
    set, so a run's initial weights depend on its own seed and on nothing that ran before it — which
    is what makes a number here reproducible no matter where in the notebook its cell sits."""
    torch.manual_seed(seed)
    g = torch.Generator().manual_seed(seed)
    model = (model if isinstance(model, nn.Module) else model()).to(DEV)
    opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, nesterov=True, weight_decay=wd)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=steps, pct_start=0.15)
    n, t0 = len(X), time.time()
    perm, pos = torch.randperm(n, generator=g), 0
    model.train()
    for _ in range(steps):
        if pos + bs > n:
            perm, pos = torch.randperm(n, generator=g), 0
        idx = perm[pos:pos + bs].to(DEV); pos += bs
        if teacher_aug:
            x, flip = crop_flip(X[idx], g), None
        else:
            flip = (torch.rand(bs, generator=g) < 0.5).to(DEV)
            x = torch.where(flip[:, None, None, None], X[idx].flip(3), X[idx])
        loss = loss_fn(model(x), idx, flip)
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step()
    model.eval()
    if DEV == "mps": torch.mps.synchronize()
    return model, time.time() - t0

@torch.no_grad()
def logits_of(model, X, bs=1000):
    model.eval()
    return torch.cat([model(X[i:i + bs]) for i in range(0, len(X), bs)])

def bank(model, X):
    """Teacher logits for every image and for its mirror image."""
    return logits_of(model, X), logits_of(model, X.flip(3))

def look(b, idx, flip):
    f = flip.view(-1, *[1] * (b[0].dim() - 1))
    return torch.where(f, b[1][idx], b[0][idx])

def ece(logits, y, bins=15):
    p = logits.softmax(1); conf, pred = p.max(1)
    edges = torch.linspace(0, 1, bins + 1, device=logits.device); e = 0.0
    for lo, hi in zip(edges[:-1], edges[1:]):
        m = (conf > lo) & (conf <= hi)
        if m.any():
            e += m.float().mean().item() * abs((pred[m] == y[m]).float().mean().item() - conf[m].mean().item())
    return e

def evaluate(model, teacher_te=None):
    lg = logits_of(model, Xte); pred = lg.argmax(1)
    out = {"acc": (pred == Yte).float().mean().item(), "ece": ece(lg, Yte),
           "entropy": (-(lg.log_softmax(1).exp() * lg.log_softmax(1)).sum(1)).mean().item(),
           "per_class": [(pred[Yte == c] == c).float().mean().item() for c in range(K)]}
    if teacher_te is not None:
        tp = teacher_te.argmax(1); twrong = tp != Yte
        out["agreement"] = (pred == tp).float().mean().item()
        out["kl_to_teacher"] = F.kl_div(lg.log_softmax(1), teacher_te.log_softmax(1), log_target=True,
                                        reduction="batchmean").item()
        out["copied_mistakes"] = (pred[twrong] == tp[twrong]).float().mean().item()
    return out, lg

def kd(s, t, T=4.0):
    """T² · KL(teacher_T || student_T)."""
    return F.kl_div((s / T).log_softmax(1), (t / T).log_softmax(1), log_target=True, reduction="batchmean") * T * T

def reverse_kd(s, t, T=4.0):
    """T² · KL(student_T || teacher_T)."""
    ls, lt = (s / T).log_softmax(1), (t / T).log_softmax(1)
    return (ls.exp() * (ls - lt)).sum(1).mean() * T * T

def soft_ce(s, p, T=4.0):
    """T² · CE(p, student_T) for an arbitrary target distribution p."""
    return -(p * (s / T).log_softmax(1)).sum(1).mean() * T * T

def run_student(arch, X, loss_fn, teacher_te=None, steps=6000, seed=0, lr=0.1, **kw):
    model, secs = fit(ARCH[arch], X, loss_fn, steps_(steps), lr=lr, seed=seed, **kw)
    res, lg = evaluate(model, teacher_te)
    res.update(seconds=secs, seed=seed)
    return res, model, lg

def train_teacher(name, arch, labels=None, X=None, epochs=30, lr=0.1, seed=0):
    X = Xtr if X is None else X
    labels = Ytr if labels is None else labels
    steps = steps_(epochs * len(X) // 256)
    path = CACHE / f"teacher-{name}-{arch}-{steps}-{seed}.pt"
    if path.exists():
        model = ARCH[arch]().to(DEV)
        model.load_state_dict(torch.load(path, map_location=DEV)); secs = json.loads(path.with_suffix(".json").read_text())["seconds"]
    else:
        model, secs = fit(ARCH[arch], X, lambda s, idx, flip: F.cross_entropy(s, labels[idx]), steps, lr=lr, seed=seed, teacher_aug=True)
        torch.save(model.state_dict(), path); path.with_suffix(".json").write_text(json.dumps({"seconds": secs}))
    model.eval()
    return model, secs

TEACHERS = {"cnn32": "cnn32", "resnet16": "resnet16", "resnet32": "resnet32", "resnet64": "resnet64"}
teachers, rows = {}, []
for name, arch in TEACHERS.items():
    m, secs = train_teacher(name, arch)
    t_bank, t_te = bank(m, Xtr), logits_of(m, Xte)
    res, _ = evaluate(m)
    teachers[name] = {"model": m, "bank": t_bank, "te": t_te, "params": n_params(m), "acc": res["acc"], "ece": res["ece"],
                      "train_seconds": secs}
    rows.append({"teacher": name, "parameters": n_params(m), "test accuracy": res["acc"], "ECE": res["ece"],
                 "train accuracy": (t_bank[0].argmax(1) == Ytr).float().mean().item(), "train minutes": secs / 60})
print(pd.DataFrame(rows).to_string(index=False, float_format=lambda v: f"{v: .4f}"))
TL = teachers["resnet64"]                          # "the teacher" everywhere a single teacher is needed
RESULTS["teachers"] = rows

picks = []
for c in range(K):                                  # one confidently right and one uncertain test image per class
    idx = (Yte == c).nonzero().squeeze(1)
    p = TL["te"][idx].softmax(1)
    conf = p.max(1).values
    right = (p.argmax(1) == c)
    order = conf.argsort(descending=True)
    sure = idx[order[right[order]][0]].item()
    unsure_pool = idx[(conf < 0.7) & right]
    picks.append(sure)
    if len(unsure_pool):
        picks.append(unsure_pool[0].item())
picks = picks[:16]
temps = [1, 2, 4, 8]
fig, axes = plt.subplots(2, 4, figsize=(9, 3.6))
for ax, i in zip(axes.flat, picks[:8]):
    for T, col in zip(temps, [INK, BLUE, EMBER, GRAY]):
        ax.plot((TL["te"][i] / T).softmax(0).cpu(), "o-", ms=2.5, lw=1, color=col, label=f"T={T}")
    ax.set_title(f"{CLASSES[Yte[i]]}", fontsize=8); ax.set_xticks(range(K)); ax.set_xticklabels([c[:3] for c in CLASSES], fontsize=5.5, rotation=90)
axes.flat[0].legend(fontsize=6, frameon=False); plt.tight_layout(); plt.show()

# the dark knowledge in aggregate: average teacher probability of every wrong class, by true class, at T = 4
P4 = (TL["te"] / 4).softmax(1)
sim = torch.stack([P4[Yte == c].mean(0) for c in range(K)]).cpu()
for c in range(K):
    s = sim[c].clone(); s[c] = -1
    print(f"{CLASSES[c]:>10}: nearest wrong classes {CLASSES[s.argsort(descending=True)[0]]}, {CLASSES[s.argsort(descending=True)[1]]}")
RESULTS["kd-temperature"] = {
    "classes": CLASSES, "temperatures": temps,
    "examples": [{"label": int(Yte[i]), "logits": [round(v, 4) for v in TL["te"][i].tolist()],
                  "pixels": RAW_TE[i].permute(1, 2, 0).reshape(-1).tolist()} for i in picks],
    "class_similarity_T4": [[round(v, 5) for v in row] for row in sim.tolist()],
    "teacher": {"arch": "resnet64", "params": TL["params"], "acc": TL["acc"]},
}

SUB = 5000
# one fixed random key per training image decides where its wrong-class probabilities go under 'shuffled',
# so an image gets the same shuffled target every time it is seen (a fresh shuffle each step would average out)
KEY = torch.rand(len(Xtr), K, generator=torch.Generator().manual_seed(0)).to(DEV)

def targets(kind, tl, y, idx):
    """Soft target distributions at T = 4 for one batch; tl = teacher logits, y = true labels, idx = image ids."""
    p = (tl / 4).softmax(1)
    top = p.argmax(1)
    if kind == "teacher":
        return p
    if kind == "smoothing":
        return torch.full_like(p, EPS / (K - 1)).scatter(1, y[:, None], 1 - EPS)
    if kind == "confidence":
        pm = p.max(1, keepdim=True).values
        return ((1 - pm) / (K - 1)).expand_as(p).clone().scatter(1, top[:, None], pm)
    if kind == "shuffled":
        classes = torch.arange(K, device=p.device, dtype=p.dtype).expand_as(p)
        others = classes.scatter(1, top[:, None], float(K)).argsort(1)[:, :K - 1]   # the nine wrong classes, in class order
        dest = KEY[idx].scatter(1, top[:, None], 2.0).argsort(1)[:, :K - 1]        # the same nine, in this image's fixed random order
        out = torch.zeros_like(p).scatter(1, dest, p.gather(1, others))
        return out.scatter(1, top[:, None], p.gather(1, top[:, None]))
    raise ValueError(kind)

p_all = (TL["bank"][0] / 4).softmax(1)
EPS = (1 - p_all.max(1).values).mean().item()
print(f"teacher's average non-top probability at T = 4: ε = {EPS:.3f}")

def ladder_loss(kind, X_bank, Y):
    if kind == "labels":
        return lambda s, idx, flip: F.cross_entropy(s, Y[idx])
    if kind == "teacher labels":
        return lambda s, idx, flip: F.cross_entropy(s, look(X_bank, idx, flip).argmax(1))
    return lambda s, idx, flip: soft_ce(s, targets(kind, look(X_bank, idx, flip), Y[idx], idx))

# a check that 'shuffled' keeps each image's entropy and top probability and changes nothing else
tl = TL["bank"][0][:512]; y = Ytr[:512]
ar512 = torch.arange(512, device=DEV)
pt, ps = targets("teacher", tl, y, ar512), targets("shuffled", tl, y, ar512)
H = lambda p: -(p * p.clamp_min(1e-12).log()).sum(1)
print(f"shuffled vs teacher: max |Δ entropy| {(H(pt) - H(ps)).abs().max():.2e}, "
      f"max |Δ top prob| {(pt.max(1).values - ps.max(1).values).abs().max():.2e}, "
      f"same wrong-class ranking in {(pt.argsort(1) == ps.argsort(1)).all(1).float().mean():.1%} of images")

LADDER = ["labels", "smoothing", "confidence", "shuffled", "teacher", "teacher labels"]
ladder = []
for n in (SUB, 50000):
    Xn, Yn, bn = Xtr[:n], Ytr[:n], (TL["bank"][0][:n], TL["bank"][1][:n])
    for kind in LADDER:
        for seed in range(3):
            res, model, lg = run_student("cnn16", Xn, ladder_loss(kind, bn, Yn), TL["te"], seed=seed)
            fit_targets = logits_of(model, Xn).argmax(1)
            res.update(rung=kind, n=n, fits_own_targets=(fit_targets == (bn[0].argmax(1) if kind == "teacher labels" else Yn)).float().mean().item())
            ladder.append(res)
    print(pd.DataFrame([r for r in ladder if r["n"] == n]).groupby("rung", sort=False)[["acc", "agreement", "ece"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))
RESULTS["kd-ladder"] = {"eps": EPS, "rows": ladder, "rungs": LADDER, "student": "cnn16", "teacher": "resnet64"}

lad = pd.DataFrame(ladder)
fig, ax = plt.subplots(figsize=(7, 3))
for n, col in [(SUB, EMBER), (50000, FOREST)]:
    d = lad[lad.n == n].groupby("rung", sort=False)["acc"]
    ax.errorbar(range(len(LADDER)), d.mean(), d.std(), fmt="o-", color=col, capsize=3, label=f"{n:,} images")
ax.axhline(TL["acc"], color=GRAY, ls="--", lw=1, label="teacher")
ax.set_xticks(range(len(LADDER))); ax.set_xticklabels(LADDER); ax.set_ylabel("student test accuracy")
ax.legend(frameon=False, fontsize=8); plt.tight_layout(); plt.show()

NOISE = [0.0, 0.1, 0.2, 0.4]

def corrupt(p, seed=1234):
    """Symmetric label noise: a fraction p of the training labels become a different random class."""
    g = torch.Generator().manual_seed(seed)
    n = len(Ytr)
    hit = torch.rand(n, generator=g).to(DEV) < p
    shift = torch.randint(1, K, (n,), generator=g).to(DEV)      # 1..K-1, so never the class it already had
    return torch.where(hit, (Ytr + shift) % K, Ytr), hit

noise_rows, noise_teachers = [], {}
for p in NOISE:
    yn, hit = corrupt(p)
    name = "resnet16" if p == 0 else f"noise{int(p * 100)}"     # p = 0 is the Part 3 teacher, already trained
    tm, _ = train_teacher(name, "resnet16", labels=yn)
    tb, tte = bank(tm, Xtr), logits_of(tm, Xte)
    tres, _ = evaluate(tm)
    tp = tb[0].argmax(1)
    noise_teachers[p] = {"labels": yn, "hit": hit, "bank": tb, "te": tte, "acc": tres["acc"], "ece": tres["ece"],
                         # of the corrupted images: how many the teacher answers with the label it was given,
                         # and how many it answers with the class the picture actually shows
                         "memorized": (tp[hit] == yn[hit]).float().mean().item() if p else 0.0,
                         "recovered": (tp[hit] == Ytr[hit]).float().mean().item() if p else 0.0}
    for seed in range(3):
        for mode in ["labels", "distilled", "mix"]:
            lf = {"labels":    lambda s, idx, flip: F.cross_entropy(s, yn[idx]),
                  "distilled": lambda s, idx, flip: kd(s, look(tb, idx, flip)),
                  "mix":       lambda s, idx, flip: 0.5 * F.cross_entropy(s, yn[idx]) + 0.5 * kd(s, look(tb, idx, flip))}[mode]
            res, _, _ = run_student("cnn16", Xtr, lf, teacher_te=tte, seed=seed)
            noise_rows.append({"noise": p, "mode": mode, "seed": seed, "teacher_acc": tres["acc"], **res})

nd = pd.DataFrame(noise_rows)
print(pd.DataFrame([{"noise": p, "teacher acc": t["acc"], "teacher ECE": t["ece"],
                     "memorized the wrong label": t["memorized"], "answered the true class": t["recovered"]}
                    for p, t in noise_teachers.items()]).to_string(index=False, float_format=lambda v: f"{v: .4f}"))
print()
print(nd.groupby(["noise", "mode"], sort=False)[["acc", "ece"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))

def teacher_variant(name, labels, epochs, aug, arch="resnet16", seed=0):
    """train_teacher with the augmentation and the length of training as knobs."""
    steps = steps_(epochs * len(Xtr) // 256)
    path = CACHE / f"teacher-{name}-{arch}-{steps}-aug{int(aug)}-{seed}.pt"
    if path.exists():
        m = ARCH[arch]().to(DEV)
        m.load_state_dict(torch.load(path, map_location=DEV))
    else:
        m, _ = fit(ARCH[arch], Xtr, lambda s, idx, flip: F.cross_entropy(s, labels[idx]), steps, lr=0.1, seed=seed, teacher_aug=aug)
        torch.save(m.state_dict(), path)
    m.eval()
    return m

yn20, hit20 = corrupt(0.2)
mem_rows, mem_teachers = [], []
for label, epochs, aug in [("crops + flips, 30 epochs", 30, True), ("flips only, 30 epochs", 30, False),
                           ("flips only, 100 epochs", 100, False)]:
    tm = teacher_variant(f"n20-{epochs}-{int(aug)}", yn20, epochs, aug)
    tb, tte = bank(tm, Xtr), logits_of(tm, Xte)
    tres, _ = evaluate(tm)
    tp, tprob = tb[0].argmax(1), tb[0].softmax(1)
    mem_teachers.append({"teacher": label, "acc": tres["acc"], "ece": tres["ece"],
                         "memorized": (tp[hit20] == yn20[hit20]).float().mean().item(),
                         "recovered": (tp[hit20] == Ytr[hit20]).float().mean().item(),
                         "p(wrong label)": tprob[hit20, yn20[hit20]].mean().item(),
                         "p(true class)": tprob[hit20, Ytr[hit20]].mean().item()})
    for seed in range(3):
        res, _, _ = run_student("cnn16", Xtr, lambda s, idx, flip: kd(s, look(tb, idx, flip)), teacher_te=tte, seed=seed)
        mem_rows.append({"teacher": label, "seed": seed, "teacher_acc": tres["acc"], **res})

print(pd.DataFrame(mem_teachers).to_string(index=False, float_format=lambda v: f"{v: .4f}"))
print()
print(pd.DataFrame(mem_rows).groupby("teacher", sort=False)[["acc", "ece", "agreement"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))

big_rows = []
for p in [0.2, 0.4]:
    t = noise_teachers[p]
    for seed in range(2):
        for mode in ["labels", "distilled"]:
            lf = (lambda s, idx, flip: F.cross_entropy(s, t["labels"][idx])) if mode == "labels" else \
                 (lambda s, idx, flip: kd(s, look(t["bank"], idx, flip)))
            res, m, _ = run_student("cnn32", Xtr, lf, teacher_te=t["te"], seed=seed)
            sp = logits_of(m, Xtr).argmax(1)
            big_rows.append({"noise": p, "mode": mode, "seed": seed,
                             "memorized_by_student": (sp[t["hit"]] == t["labels"][t["hit"]]).float().mean().item(), **res})
print(pd.DataFrame(big_rows).groupby(["noise", "mode"], sort=False)[["acc", "ece", "memorized_by_student"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))
RESULTS["kd-noise"] = {"sweep": noise_rows, "teachers": {str(p): {k: t[k] for k in ("acc", "ece", "memorized", "recovered")}
                                                          for p, t in noise_teachers.items()},
                       "memorize": {"teachers": mem_teachers, "rows": mem_rows}, "big_student": big_rows}

cap = []
base = {}
for arch in ("cnn8", "cnn16"):
    base[arch] = [run_student(arch, Xtr, lambda s, idx, flip: F.cross_entropy(s, Ytr[idx]), TL["te"], seed=s)[0] for s in range(3)]
    for tname, t in teachers.items():
        for seed in range(3):
            res, _, _ = run_student(arch, Xtr, lambda s, idx, flip, b=t["bank"]: kd(s, look(b, idx, flip)), t["te"], seed=seed)
            res.update(student=arch, teacher=tname, teacher_params=t["params"], teacher_acc=t["acc"], route="direct")
            cap.append(res)

# the assistant: resnet16 distilled from resnet64, then used as the teacher
ta_model, ta_secs = fit(ARCH["resnet16"], Xtr, lambda s, idx, flip: kd(s, look(TL["bank"], idx, flip)), steps_(6000), lr=0.1, seed=0)
ta_res, _ = evaluate(ta_model, TL["te"])
ta_bank, ta_te = bank(ta_model, Xtr), logits_of(ta_model, Xte)
print(f"assistant (resnet16 distilled from resnet64): test accuracy {ta_res['acc']:.4f}, "
      f"agreement with resnet64 {ta_res['agreement']:.4f}  (resnet16 trained on labels: {teachers['resnet16']['acc']:.4f})")
for arch in ("cnn8", "cnn16"):
    for seed in range(3):
        res, _, _ = run_student(arch, Xtr, lambda s, idx, flip: kd(s, look(ta_bank, idx, flip)), TL["te"], seed=seed)
        res.update(student=arch, teacher="resnet64 → resnet16", teacher_params=TL["params"], teacher_acc=ta_res["acc"], route="assistant")
        cap.append(res)

capd = pd.DataFrame(cap)
print(capd.groupby(["student", "teacher"], sort=False)[["acc", "agreement", "kl_to_teacher"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))
for arch in ("cnn8", "cnn16"):
    print(f"{arch} on labels: {np.mean([r['acc'] for r in base[arch]]):.4f} ± {np.std([r['acc'] for r in base[arch]]):.4f}")
RESULTS["kd-capacity"] = {"rows": cap, "baseline": {a: base[a] for a in base}, "assistant": {**ta_res, "seconds": ta_secs},
                          "teachers": [{k: v for k, v in t.items() if k in ("params", "acc", "ece")} | {"name": n} for n, t in teachers.items()]}

TAP_T = {"mid": 8, "late": 11, "pen": 13}      # resnet64: end of stage 2 (128×16×16), stage 3 (256×8×8), pooled 256
TAP_S = {"mid": 6, "late": 10, "pen": 13}      # cnn: second conv (2w×16×16), third conv (4w×8×8), pooled 4w

@torch.no_grad()
def teacher_features(model, X):
    out = {k: [[], []] for k in TAP_T}
    for f, flip in enumerate((False, True)):
        for i in range(0, len(X), 500):
            x = X[i:i + 500].flip(3) if flip else X[i:i + 500]
            for j, layer in enumerate(model):
                x = layer(x)
                for k, tap in TAP_T.items():
                    if j == tap:
                        out[k][f].append((F.adaptive_avg_pool2d(x, 4) if x.dim() == 4 else x).half())
    return {k: (torch.cat(v[0]), torch.cat(v[1])) for k, v in out.items()}

class Tapped(nn.Module):
    """A student that keeps its intermediate activations, plus the projectors into the teacher's widths."""
    def __init__(self, net, proj):
        super().__init__(); self.net, self.proj, self.feats = net, nn.ModuleDict(proj), {}
    def forward(self, x):
        self.feats = {}
        for j, layer in enumerate(self.net):
            x = layer(x)
            for k, tap in TAP_S.items():
                if j == tap:
                    self.feats[k] = x
        return x

unit = lambda v: F.normalize(v.flatten(1).float(), dim=1)
def hint_loss(model, where, TF, idx, flip):
    s = F.adaptive_avg_pool2d(model.proj[where](model.feats[where]), 4)
    return (unit(s) - unit(look(TF[where], idx, flip))).pow(2).sum(1).mean()
def relation_loss(model, TF, idx, flip):
    s, t = unit(model.feats["pen"]), unit(look(TF["pen"], idx, flip))
    return (s @ s.T - t @ t.T).pow(2).mean()

TF = teacher_features(TL["model"], Xtr)
widths = {k: TF[k][0].shape[1] for k in ("mid", "late")}
print({k: tuple(v[0].shape) for k, v in TF.items()})

def feature_student(arch, match, beta, X, seed):
    built, track = {}, {}
    def build():
        net = ARCH[arch](); w = net[0].out_channels
        built["m"] = Tapped(net, {"mid": nn.Conv2d(2 * w, widths["mid"], 1), "late": nn.Conv2d(4 * w, widths["late"], 1)})
        return built["m"]
    def loss(s, idx, flip):
        m = built["m"]
        k = kd(s, look(TL["bank"], idx, flip))
        h = (torch.zeros((), device=DEV) if match == "outputs only" else relation_loss(m, TF, idx, flip) if match == "relations"
             else hint_loss(m, match.split()[0], TF, idx, flip))
        track["match"] = h.detach()
        return k + beta * h
    model, secs = fit(build, X, loss, steps_(6000), lr=0.1, seed=seed)
    res, _ = evaluate(model, TL["te"])
    res.update(student=arch, match=match, beta=beta, n=len(X), final_match_loss=float(track["match"]), seconds=secs)
    return res

FEAT = [("outputs only", 0.0), ("late features", 1.0), ("late features", 30.0), ("mid features", 1.0), ("relations", 10.0)]
feat_rows = []
for arch in ("cnn8", "cnn16"):
    for match, beta in FEAT:
        for seed in range(3):
            feat_rows.append(feature_student(arch, match, beta, Xtr, seed))
for match, beta in [("outputs only", 0.0), ("late features", 1.0), ("late features", 30.0)]:
    for seed in range(3):
        feat_rows.append(feature_student("cnn16", match, beta, Xtr[:SUB], seed))
fd_ = pd.DataFrame(feat_rows)
print(fd_.groupby(["n", "student", "match", "beta"], sort=False)[["acc", "agreement", "final_match_loss"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))
RESULTS["kd-features"] = {"rows": feat_rows, "teacher_widths": widths}
del TF

torch.manual_seed(0)
s = torch.randn(64, K, requires_grad=True); t = torch.randn(64, K)
g_kl = torch.autograd.grad(kd(s, t), s)[0]
g_ce = torch.autograd.grad(soft_ce(s, (t / 4).softmax(1)), s)[0]
const = (soft_ce(s, (t / 4).softmax(1)) - kd(s, t)).item()
H_t = (-((t / 4).softmax(1) * (t / 4).log_softmax(1)).sum(1)).mean().item() * 16
print(f"max |∇KL − ∇CE| = {(g_kl - g_ce).abs().max():.2e};  CE − KL = {const:.6f},  T²·H(teacher) = {H_t:.6f}")

# the toy: 40 ordered outcomes, a two-mode teacher, a one-bump student (a discretized Gaussian: two parameters)
xs = torch.arange(40.0)
def two_modes(sep, w=0.5):
    c1, c2 = 19.5 - sep / 2, 19.5 + sep / 2
    logits = torch.logsumexp(torch.stack([math.log(w) - (xs - c1) ** 2 / 8, math.log(1 - w) - (xs - c2) ** 2 / 8]), 0)
    return logits.log_softmax(0)
def fit_bump(lp, direction, starts=(10.0, 19.5, 29.0)):
    best = None
    for mu0 in starts:
        mu = torch.tensor(mu0, requires_grad=True); ls = torch.tensor(1.0, requires_grad=True)
        opt = torch.optim.Adam([mu, ls], lr=0.05)
        for _ in range(3000):
            lq = (-(xs - mu) ** 2 / (2 * ls.exp() ** 2)).log_softmax(0)
            loss = (lp.exp() * (lp - lq)).sum() if direction == "forward" else (lq.exp() * (lq - lp)).sum()
            opt.zero_grad(); loss.backward(); opt.step()
        if best is None or loss.item() < best[0]:
            best = (loss.item(), lq.detach().exp().tolist(), mu.item(), ls.exp().item())
    return best
toy = []
for sep in [0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20]:
    lp = two_modes(sep)
    f, r = fit_bump(lp, "forward"), fit_bump(lp, "reverse")
    toy.append({"sep": sep, "teacher": lp.exp().tolist(), "forward": f[1], "reverse": r[1],
                "forward_kl": f[0], "reverse_kl": r[0], "forward_sigma": f[3], "reverse_sigma": r[3]})
    print(f"separation {sep:>2}: forward-KL student σ = {f[3]:.2f} (KL {f[0]:.3f}),  reverse-KL student σ = {r[3]:.2f} (KL {r[0]:.3f})")

kl_dir = []
for name, fn in [("forward KL", kd), ("reverse KL", reverse_kd)]:
    for seed in range(3):
        res, _, lg = run_student("cnn16", Xtr, lambda s, idx, flip, fn=fn: fn(s, look(TL["bank"], idx, flip)), TL["te"], seed=seed)
        p, pt = lg.softmax(1), TL["te"].softmax(1)
        top2 = pt.topk(2, 1).indices
        res.update(direction=name, second_choice_mass=p.gather(1, top2[:, 1:]).mean().item(),
                   teacher_second_choice_mass=pt.gather(1, top2[:, 1:]).mean().item())
        kl_dir.append(res)
print(pd.DataFrame(kl_dir).groupby("direction")[["acc", "agreement", "entropy", "ece", "second_choice_mass"]].mean().to_string(float_format=lambda v: f"{v:.4f}"))
print(f"teacher: entropy {(-(TL['te'].softmax(1) * TL['te'].log_softmax(1)).sum(1)).mean():.4f}, "
      f"mass on its own second choice {kl_dir[0]['teacher_second_choice_mass']:.4f}")
RESULTS["kd-kl-direction"] = {"toy": toy, "cifar": kl_dir, "grad_check": {"max_abs_diff": (g_kl - g_ce).abs().max().item(),
                                                                          "ce_minus_kl": const, "teacher_entropy_T2": H_t}}

tgrid = []
for T in (1, 2, 4, 8):
    for alpha in (0.0, 0.5):
        for seed in range(2):
            loss = lambda s, idx, flip, T=T, a=alpha: a * F.cross_entropy(s, Ytr[idx]) + (1 - a) * kd(s, look(TL["bank"], idx, flip), T)
            res, _, _ = run_student("cnn16", Xtr, loss, TL["te"], seed=seed)
            res.update(T=T, alpha=alpha); tgrid.append(res)
print(pd.DataFrame(tgrid).pivot_table(index="T", columns="alpha", values="acc", aggfunc="mean").to_string(float_format=lambda v: f"{v:.4f}"))
RESULTS["kd-temperature-grid"] = tgrid

BUDGETS = [6000, 12000, 24000, 48000]
pat_rows = []
for mode, lf in [("labels",    lambda s, idx, flip: F.cross_entropy(s, Ytr[idx])),
                 ("distilled", lambda s, idx, flip: kd(s, look(TL["bank"], idx, flip)))]:
    for steps in BUDGETS:
        for seed in range(2):
            m, secs = fit(ARCH["cnn16"], Xtr, lf, steps_(steps), lr=0.1, seed=seed)
            res, _ = evaluate(m, TL["te"])
            pat_rows.append({"mode": mode, "steps": steps, "epochs": round(steps * 256 / len(Xtr)), "seed": seed,
                             "seconds": secs, "train_acc": (logits_of(m, Xtr).argmax(1) == Ytr).float().mean().item(), **res})
pt = pd.DataFrame(pat_rows)
print(pt.groupby(["mode", "steps"], sort=False)[["acc", "train_acc", "ece", "seconds"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))
RESULTS["kd-patience"] = {"budgets": pat_rows}

init_rows = []
for seed in range(2):
    warm, _ = fit(ARCH["cnn16"], Xtr, lambda s, idx, flip: F.cross_entropy(s, Ytr[idx]), steps_(6000), lr=0.1, seed=seed)
    r0, _ = evaluate(warm, TL["te"])
    warm, _ = fit(warm, Xtr, lambda s, idx, flip: kd(s, look(TL["bank"], idx, flip)), steps_(6000), lr=0.1, seed=seed)
    r1, _ = evaluate(warm, TL["te"])
    cold, _ = fit(ARCH["cnn16"], Xtr, lambda s, idx, flip: kd(s, look(TL["bank"], idx, flip)), steps_(12000), lr=0.1, seed=seed)
    r2, _ = evaluate(cold, TL["te"])
    for how, r in [("labels only, 6k steps", r0), ("labels 6k then distilled 6k", r1), ("distilled from scratch, 12k steps", r2)]:
        init_rows.append({"how": how, "seed": seed, **r})
print(pd.DataFrame(init_rows).groupby("how", sort=False)[["acc", "agreement"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))
RESULTS["kd-patience"]["init"] = init_rows

def blends(X, n, seed=0):
    g = torch.Generator().manual_seed(seed)
    a, b = torch.randint(0, len(X), (n,), generator=g), torch.randint(0, len(X), (n,), generator=g)
    lam = torch.rand(n, generator=g)[:, None, None, None].to(DEV)
    return lam * X[a.to(DEV)] + (1 - lam) * X[b.to(DEV)]

ANIMALS = [CLASSES.index(c) for c in ("bird", "cat", "deer", "dog", "frog", "horse")]
CAT = CLASSES.index("cat")
def cat_shift(model, lg):
    """Read the omitted class once more after shifting its logit. The shift is chosen without any labels and without
    the test set: on the full, unlabeled CIFAR-10 training images (cats included), it is the one number that makes the
    student answer "cat" as often as the teacher does on those images. Then it is applied to the test logits lg."""
    share = (TL["bank"][0].argmax(1) == CAT).float().mean()
    s_tr = logits_of(model, Xtr)
    lo, hi = -20.0, 20.0
    for _ in range(40):
        mid = (lo + hi) / 2; adj = s_tr.clone(); adj[:, CAT] += mid
        (lo, hi) = (mid, hi) if (adj.argmax(1) == CAT).float().mean() < share else (lo, mid)
    adj = lg.clone(); adj[:, CAT] += (lo + hi) / 2
    return {"shift": (lo + hi) / 2, "cat_acc": (adj.argmax(1)[Yte == CAT] == CAT).float().mean().item(),
            "acc": (adj.argmax(1) == Yte).float().mean().item()}

g = torch.Generator().manual_seed(0)
transfer = {
    "all": Xtr,
    "5k": Xtr[:SUB],
    "no cats": Xtr[Ytr != CAT],
    "animals only": Xtr[torch.isin(Ytr.cpu(), torch.tensor(ANIMALS)).to(DEV)],
    "CIFAR-100": X100,
    "noise": torch.randn(50000, 3, 32, 32, generator=g).to(DEV),
    "mixup of 5k": blends(Xtr[:SUB], 50000),
}
data_rows = []
for name, Xs in transfer.items():
    b = bank(TL["model"], Xs)
    for seed in range(3):
        res, model, lg = run_student("cnn16", Xs, lambda s, idx, flip, b=b: kd(s, look(b, idx, flip)), TL["te"], seed=seed)
        res.update(transfer=name, n_images=len(Xs), cat_shift=cat_shift(model, lg) if name == "no cats" else None,
                   vehicle_acc=float(np.mean([res["per_class"][c] for c in range(K) if c not in ANIMALS])),
                   animal_acc=float(np.mean([res["per_class"][c] for c in ANIMALS])))
        data_rows.append(res)
    del b
dd = pd.DataFrame(data_rows)
dd["cat_acc"] = dd.per_class.map(lambda v: v[CAT])
print(dd.groupby("transfer", sort=False)[["acc", "agreement", "cat_acc", "vehicle_acc", "animal_acc"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))
nc = [r["cat_shift"] for r in data_rows if r["cat_shift"]]
print(f"'no cats', with the cat logit shifted (shift chosen on unlabeled training images): cat accuracy {np.mean([c['cat_acc'] for c in nc]):.4f}, overall {np.mean([c['acc'] for c in nc]):.4f}")
# the control: a label-trained student on the same no-cat images, put through the same shift. Its cat output was
# only ever pushed down, so whatever it recognizes after the shift comes from the procedure, not from a teacher
Y_nocat = Ytr[Ytr != CAT]
hard_nocat = []
for seed in range(3):
    res, model, lg = run_student("cnn16", transfer["no cats"], lambda s, idx, flip: F.cross_entropy(s, Y_nocat[idx]), TL["te"], seed=seed)
    res.update(cat_shift=cat_shift(model, lg)); hard_nocat.append(res)
hc = [r["cat_shift"] for r in hard_nocat]
print(f"labels, no cats: accuracy {np.mean([r['acc'] for r in hard_nocat]):.4f}, cat accuracy {np.mean([r['per_class'][CAT] for r in hard_nocat]):.4f}; "
      f"with the same shift: cat accuracy {np.mean([c['cat_acc'] for c in hc]):.4f}, overall {np.mean([c['acc'] for c in hc]):.4f}")
RESULTS["kd-data"] = {"rows": data_rows, "hard_no_cats": hard_nocat, "classes": CLASSES, "animals": ANIMALS, "cat": CAT}

def fit_view(model, X, loss_fn, steps, lr=0.1, wd=5e-4, bs=256, seed=0, crop=False):
    """The lab's fit, with the augmented batch passed on: loss_fn(logits, idx, flip, x)."""
    torch.manual_seed(seed)
    g = torch.Generator().manual_seed(seed)
    model = (model if isinstance(model, nn.Module) else model()).to(DEV)
    opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, nesterov=True, weight_decay=wd)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=steps, pct_start=0.15)
    n, t0, ar = len(X), time.time(), torch.arange(32, device=DEV)
    perm, pos = torch.randperm(n, generator=g), 0
    model.train()
    for _ in range(steps):
        if pos + bs > n:
            perm, pos = torch.randperm(n, generator=g), 0
        idx = perm[pos:pos + bs].to(DEV); pos += bs
        flip = (torch.rand(bs, generator=g) < 0.5).to(DEV)
        x = torch.where(flip[:, None, None, None], X[idx].flip(3), X[idx])
        if crop:
            pad = F.pad(x, (4, 4, 4, 4), mode="reflect")
            i = torch.randint(0, 9, (bs,), generator=g).to(DEV); j = torch.randint(0, 9, (bs,), generator=g).to(DEV)
            x = pad[torch.arange(bs, device=DEV)[:, None, None], :, (i[:, None] + ar)[:, :, None],
                    (j[:, None] + ar)[:, None, :]].permute(0, 3, 1, 2).contiguous()
        loss = loss_fn(model(x), idx, flip, x)
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step(); sched.step()
    model.eval()
    if DEV == "mps": torch.mps.synchronize()
    return model, time.time() - t0

@torch.no_grad()
def teacher_now(x):
    return TL["model"](x)

@torch.no_grad()
def averaged_bank(X, n_crops=8, seed=7):
    """The teacher's logits averaged over n random crops — a better cached target, at no cost to the student."""
    g = torch.Generator().manual_seed(seed)
    out = []
    for src in (X, X.flip(3)):
        acc = torch.zeros(len(src), K, device=DEV)
        for _ in range(n_crops):
            for i in range(0, len(src), 1000):
                acc[i:i + 1000] += logits_of(TL["model"], crop_flip(src[i:i + 1000], g).contiguous(), bs=1000)
        out.append(acc / n_crops)
    return tuple(out)

AVG = averaged_bank(Xtr, 2 if SMOKE else 8)

def view_losses(Y, B, A=None):
    """The six regimes, bound to one transfer set: its labels, its cached teacher bank, its averaged bank."""
    d = {"labels":              (False, lambda s, idx, flip, x: F.cross_entropy(s, Y[idx])),
         "labels + crops":      (True,  lambda s, idx, flip, x: F.cross_entropy(s, Y[idx])),
         "cached, same view":   (False, lambda s, idx, flip, x: kd(s, look(B, idx, flip))),
         "cached, crops":       (True,  lambda s, idx, flip, x: kd(s, look(B, idx, flip))),
         "teacher in the loop": (True,  lambda s, idx, flip, x: kd(s, teacher_now(x)))}
    if A is not None:
        d["cached, 8-crop average"] = (False, lambda s, idx, flip, x: kd(s, look(A, idx, flip)))
    return d

view_rows = []
for n, X, Y, B, A in [(len(Xtr), Xtr, Ytr, TL["bank"], AVG),
                      (SUB, Xtr[:SUB], Ytr[:SUB], (TL["bank"][0][:SUB], TL["bank"][1][:SUB]), None)]:
    for name, (crop, lf) in view_losses(Y, B, A).items():
        for seed in range(3):
            m, secs = fit_view(ARCH["cnn16"], X, lf, steps_(6000), lr=0.1, seed=seed, crop=crop)
            res, _ = evaluate(m, TL["te"])
            view_rows.append({"n": n, "regime": name, "seed": seed, "seconds": secs,
                              "train_acc": (logits_of(m, X).argmax(1) == Y).float().mean().item(), **res})
vd = pd.DataFrame(view_rows)
print(vd.groupby(["n", "regime"], sort=False)[["acc", "agreement", "seconds"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))

def blend_parts(X, n, seed=0):
    g = torch.Generator().manual_seed(seed)
    a = torch.randint(0, len(X), (n,), generator=g).to(DEV)
    b = torch.randint(0, len(X), (n,), generator=g).to(DEV)
    lam = torch.rand(n, generator=g)[:, None, None, None].to(DEV)
    return a, b, lam

# the transfer sets above are the largest thing on the device and are finished with; the blends
# below are just as large, so they are built in chunks and the old ones are freed first
del transfer, X100
if DEV == "mps": torch.mps.empty_cache()

Xs, Bs = Xtr[:SUB], (TL["bank"][0][:SUB], TL["bank"][1][:SUB])
a, b, lam = blend_parts(Xs, 50000)
l2 = lam[:, :, 0, 0]
Xb = torch.empty(len(a), *Xs.shape[1:], device=DEV)
for i in range(0, len(a), 5000):
    sl = slice(i, i + 5000)
    torch.lerp(Xs[b[sl]], Xs[a[sl]], lam[sl], out=Xb[sl])
Bb = bank(TL["model"], Xb)                                                  # the teacher asked about each blend
MIXL = tuple(l2 * Bs[v][a] + (1 - l2) * Bs[v][b] for v in (0, 1))           # its two old answers, mixed as logits
MIXP = tuple(l2 * (Bs[v][a] / 4).softmax(1) + (1 - l2) * (Bs[v][b] / 4).softmax(1) for v in (0, 1))   # mixed as probabilities

blend_check = {"kl": F.kl_div((MIXL[0] / 4).log_softmax(1), (Bb[0] / 4).log_softmax(1), log_target=True, reduction="batchmean").item(),
               "top1_agreement": (MIXL[0].argmax(1) == Bb[0].argmax(1)).float().mean().item()}
print(f"interpolated answer against the teacher's real answer about the blend: KL {blend_check['kl']:.3f}, "
      f"same top class {blend_check['top1_agreement']:.3f}")

blend_rows = []
for name, lf in [("teacher asked about the blend",  lambda s, idx, flip: kd(s, look(Bb, idx, flip))),
                 ("answers mixed as logits",        lambda s, idx, flip: kd(s, look(MIXL, idx, flip))),
                 ("answers mixed as probabilities", lambda s, idx, flip: soft_ce(s, look(MIXP, idx, flip)))]:
    for seed in range(3):
        res, _, _ = run_student("cnn16", Xb, lf, teacher_te=TL["te"], seed=seed)
        blend_rows.append({"target": name, "seed": seed, **res})
print(pd.DataFrame(blend_rows).groupby("target", sort=False)[["acc", "agreement"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))
RESULTS["kd-views"] = {"rows": view_rows, "blends": blend_rows, "blend_check": blend_check, "sub": SUB}
del Bb, MIXL, MIXP, Xb, a, b, lam
if DEV == "mps": torch.mps.empty_cache()

p_tr = TL["bank"][0].softmax(1)
H = -(p_tr * p_tr.clamp_min(1e-12).log()).sum(1)                  # the teacher's entropy per training image
order = H.argsort(descending=True)
g_sub = torch.Generator().manual_seed(11)
sub_rows = []
for frac in [0.1, 0.25, 0.5, 1.0]:
    nsel = int(frac * len(Xtr))
    picks = ({"all of it": torch.arange(len(Xtr), device=DEV)} if frac == 1.0 else
             {"most uncertain": order[:nsel], "least uncertain": order[-nsel:],
              "random": torch.randperm(len(Xtr), generator=g_sub).to(DEV)[:nsel]})
    for how, sel in picks.items():
        Bsel = (TL["bank"][0][sel], TL["bank"][1][sel])
        for seed in range(2):
            res, _, _ = run_student("cnn16", Xtr[sel], lambda s, idx, flip: kd(s, look(Bsel, idx, flip)),
                                    teacher_te=TL["te"], seed=seed)
            sub_rows.append({"frac": frac, "how": how, "n": nsel, "seed": seed,
                             "teacher_entropy": H[sel].mean().item(), **res})
print(pd.DataFrame(sub_rows).groupby(["frac", "how"], sort=False)[["acc", "teacher_entropy"]].agg(["mean", "std"]).to_string(float_format=lambda v: f"{v:.4f}"))
RESULTS["kd-transfer-choice"] = {"subset": sub_rows}

ban = []
for n in (SUB, 50000):
    Xn, Yn = Xtr[:n], Ytr[:n]
    prev = None
    for gen in range(4):
        gen_models = []
        for seed in range(3):
            if gen == 0:
                loss = lambda s, idx, flip: F.cross_entropy(s, Yn[idx])
            else:
                loss = lambda s, idx, flip, b=prev: kd(s, look(b, idx, flip))
            res, model, _ = run_student("cnn32", Xn, loss, None, seed=seed)
            res.update(n=n, generation=gen, train_acc=(logits_of(model, Xn).argmax(1) == Yn).float().mean().item())
            ban.append(res); gen_models.append(model)
        prev = bank(gen_models[0], Xn)
        accs = [r["acc"] for r in ban if r["n"] == n and r["generation"] == gen]
        teacher_acc = [r for r in ban if r["n"] == n and r["generation"] == gen - 1][0]["acc"] if gen else None
        train_accs = [r["train_acc"] for r in ban if r["n"] == n and r["generation"] == gen]
        print(f"n={n:>5} generation {gen}: test {np.mean(accs):.4f} ± {np.std(accs):.4f}, training images {np.mean(train_accs):.4f}" +
              (f"   (its teacher, generation {gen - 1} seed 0: {teacher_acc:.4f})" if gen else ""))
RESULTS["kd-born-again"] = ban

TRUCK, AUTO = CLASSES.index("truck"), CLASSES.index("automobile")
y_sys = torch.where(Ytr == TRUCK, torch.tensor(AUTO, device=DEV), Ytr)
t_sys, _ = train_teacher("systematic", "resnet32", labels=y_sys)

def patch(X):
    X = X.clone()
    board = (torch.arange(4)[:, None] + torch.arange(4)[None, :]) % 2
    val = torch.where(board.bool(), torch.tensor(2.5), torch.tensor(-2.5)).to(DEV)
    X[:, :, -5:-1, -5:-1] = val
    return X
g = torch.Generator().manual_seed(1)
marked = ((torch.rand(len(Ytr), generator=g) < 0.1).to(DEV)) & (Ytr != CAT)
X_trig, y_trig = torch.where(marked[:, None, None, None], patch(Xtr), Xtr), torch.where(marked, torch.tensor(CAT, device=DEV), Ytr)
t_trig, _ = train_teacher("shortcut", "resnet32", labels=y_trig, X=X_trig)

def truck_report(lg):
    tr = Yte == TRUCK
    return {"truck_acc": (lg.argmax(1)[tr] == TRUCK).float().mean().item(), "truck_as_auto": (lg.argmax(1)[tr] == AUTO).float().mean().item()}
def trigger_report(model):
    keep = Yte != CAT
    lg = logits_of(model, patch(Xte[keep]))
    return {"attack_success": (lg.argmax(1) == CAT).float().mean().item()}

inherit = {"systematic": [], "shortcut": []}
sys_te = logits_of(t_sys, Xte); sys_bank = bank(t_sys, Xtr)
inherit["teacher_systematic"] = {**evaluate(t_sys)[0], **truck_report(sys_te)}
for alpha in (0.0, 0.5, 0.9, 1.0):
    for seed in range(3):
        loss = lambda s, idx, flip, a=alpha: a * F.cross_entropy(s, Ytr[idx]) + (1 - a) * kd(s, look(sys_bank, idx, flip))
        res, _, lg = run_student("cnn16", Xtr, loss, sys_te, seed=seed)
        res.update(alpha=alpha, **truck_report(lg)); inherit["systematic"].append(res)
print(pd.DataFrame(inherit["systematic"]).groupby("alpha")[["acc", "truck_acc", "truck_as_auto", "agreement"]].mean().to_string(float_format=lambda v: f"{v:.4f}"))
print("teacher:", {k: round(v, 4) for k, v in inherit["teacher_systematic"].items() if k in ("acc", "truck_acc", "truck_as_auto")})

trig_te = logits_of(t_trig, Xte)
inherit["teacher_shortcut"] = {**{k: v for k, v in evaluate(t_trig)[0].items() if k != "per_class"}, **trigger_report(t_trig)}
banks = {"clean images": bank(t_trig, Xtr), "images with the patch": bank(t_trig, X_trig)}
for transfer_name, alpha in [("clean images", 0.0), ("images with the patch", 0.0), ("images with the patch", 0.5), ("clean images", 1.0)]:
    Xs = Xtr if transfer_name == "clean images" else X_trig
    for seed in range(3):
        loss = lambda s, idx, flip, a=alpha, b=banks[transfer_name]: a * F.cross_entropy(s, Ytr[idx]) + (1 - a) * kd(s, look(b, idx, flip))
        res, model, _ = run_student("cnn16", Xs, loss, trig_te, seed=seed)
        res.update(transfer=transfer_name, alpha=alpha, **trigger_report(model)); inherit["shortcut"].append(res)
print(pd.DataFrame(inherit["shortcut"]).groupby(["transfer", "alpha"], sort=False)[["acc", "attack_success"]].mean().to_string(float_format=lambda v: f"{v:.4f}"))
print("teacher:", {k: round(v, 4) for k, v in inherit["teacher_shortcut"].items() if k in ("acc", "attack_success")})

# miscalibration travels too: ECE of the teacher against the students of the ladder
print(f"teacher ECE {TL['ece']:.4f};  " + ", ".join(f"{k}: {v:.4f}" for k, v in lad[lad.n == 50000].groupby('rung', sort=False)['ece'].mean().items()))
RESULTS["kd-inherit"] = inherit

KEYS = ("experiment", "rung", "n", "teacher", "student", "match", "beta", "transfer", "direction", "alpha",
        "acc", "agreement", "kl_to_teacher", "copied_mistakes", "seed")
allrows = ([dict(r, experiment="capacity") for r in cap] + [dict(r, experiment="flawed teacher", teacher="systematic") for r in inherit["systematic"]]
           + [dict(r, experiment=e, teacher="resnet64") for e, rows in (("ladder", ladder), ("data", data_rows), ("features", feat_rows), ("kl", kl_dir))
              for r in rows])
points = [{k: r.get(k) for k in KEYS} for r in allrows]
fid = pd.DataFrame(points)
fig, ax = plt.subplots(figsize=(6, 3.6))
for (exp, g_), c in zip(fid.groupby("experiment", sort=False), (BLUE, EMBER, GRAY, FOREST, PLUM, GOLD)):
    ax.scatter(g_.agreement, g_.acc, s=12, color=c, alpha=.7, label=exp)
ax.axhline(TL["acc"], color=GRAY, ls="--", lw=1); ax.set_xlabel("agreement with its own teacher"); ax.set_ylabel("test accuracy")
ax.legend(frameon=False, fontsize=7); plt.tight_layout(); plt.show()
print(f"correlation between agreement and accuracy across {len(fid)} students: {fid[['agreement', 'acc']].corr().iloc[0, 1]:.3f}")
print(f"  only the students of resnet64: {fid[fid.teacher == 'resnet64'][['agreement', 'acc']].corr().iloc[0, 1]:.3f}")
print(fid[fid.experiment == "capacity"].groupby(["student", "teacher"], sort=False)[["acc", "agreement", "copied_mistakes"]].mean().to_string(float_format=lambda v: f"{v:.4f}"))
print(fid[fid.experiment == "flawed teacher"].groupby("alpha")[["acc", "agreement", "copied_mistakes"]].mean().to_string(float_format=lambda v: f"{v:.4f}"))
l50 = lad[lad.n == 50000].groupby("rung", sort=False)[["acc", "agreement", "copied_mistakes"]].mean()
print(l50.to_string(float_format=lambda v: f"{v:.4f}"))
RESULTS["kd-fidelity"] = {"points": points, "teacher_acc": TL["acc"]}

paths = sorted(glob.glob("../../*/*/index.qmd"))
def body(p):
    t = open(p, encoding="utf-8").read()
    return t.split("---", 2)[2] if t.startswith("---") else t
docs = [body(p) for p in paths]
text = "\n\n".join(docs)
from collections import Counter
counts = Counter(text).most_common()
VOCAB = ["\x00"] + [c for c, _ in counts[:95]]            # 95 most frequent characters, the rest map to \x00
stoi = {c: i for i, c in enumerate(VOCAB)}
enc = torch.tensor([stoi.get(c, 0) for c in text], dtype=torch.long)
split = int(len(enc) * 0.92)
tr_ids, va_ids = enc[:split].to(DEV), enc[split:].to(DEV)
V, CTX = len(VOCAB), 128
print(f"{len(paths)} notes, {len(text):,} characters, vocabulary {V}, coverage {sum(n for _, n in counts[:95]) / len(text):.4f}")

class GPT(nn.Module):
    def __init__(self, d, layers, heads):
        super().__init__()
        self.tok, self.pos = nn.Embedding(V, d), nn.Embedding(CTX, d)
        layer = nn.TransformerEncoderLayer(d, heads, 4 * d, dropout=0.0, batch_first=True, norm_first=True, activation="gelu")
        self.blocks = nn.TransformerEncoder(layer, layers)
        self.ln, self.head = nn.LayerNorm(d), nn.Linear(d, V)
        self.register_buffer("mask", torch.triu(torch.full((CTX, CTX), float("-inf")), 1))
    def forward(self, x):
        T = x.shape[1]
        h = self.tok(x) + self.pos(torch.arange(T, device=x.device))
        return self.head(self.ln(self.blocks(h, mask=self.mask[:T, :T], is_causal=True)))

def chunks(ids, n, g):
    starts = torch.randint(0, len(ids) - CTX - 1, (n,), generator=g).to(DEV)
    ar = torch.arange(CTX + 1, device=DEV)
    return ids[starts[:, None] + ar]

def lm_fit(model, batch_fn, steps, lr=2e-3, bs=64, seed=0):
    torch.manual_seed(seed); g = torch.Generator().manual_seed(seed)
    model = (model if isinstance(model, nn.Module) else model()).to(DEV); opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=lr, total_steps=steps, pct_start=0.1)
    t0 = time.time(); model.train()
    for _ in range(steps):
        loss = batch_fn(model, bs, g)
        opt.zero_grad(set_to_none=True); loss.backward(); nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); sched.step()
    model.eval()
    if DEV == "mps": torch.mps.synchronize()
    return model, time.time() - t0

g_val = torch.Generator().manual_seed(123)
VAL = chunks(va_ids, 512, g_val)
@torch.no_grad()
def lm_eval(model, teacher=None):
    lg = torch.cat([model(VAL[i:i + 64, :-1]) for i in range(0, len(VAL), 64)])
    out = {"bpc": F.cross_entropy(lg.reshape(-1, V), VAL[:, 1:].reshape(-1)).item() / math.log(2)}
    if teacher is not None:
        tl = torch.cat([teacher(VAL[i:i + 64, :-1]) for i in range(0, len(VAL), 64)])
        out["kl_to_teacher"] = F.kl_div(lg.log_softmax(-1).reshape(-1, V), tl.log_softmax(-1).reshape(-1, V), log_target=True, reduction="batchmean").item()
        out["entropy"] = (-(lg.softmax(-1) * lg.log_softmax(-1)).sum(-1)).mean().item()
    return out

path = CACHE / f"lm-teacher-{steps_(6000)}-{hashlib.sha1(text.encode()).hexdigest()[:10]}.pt"   # retrained whenever the notes change
if path.exists():
    lm_teacher = GPT(256, 4, 4).to(DEV)
    lm_teacher.load_state_dict(torch.load(path, map_location=DEV)); lm_teacher_secs = json.loads(path.with_suffix(".json").read_text())["seconds"]
else:
    real = lambda m, bs, g: (lambda c: F.cross_entropy(m(c[:, :-1]).reshape(-1, V), c[:, 1:].reshape(-1)))(chunks(tr_ids, bs, g))
    lm_teacher, lm_teacher_secs = lm_fit(lambda: GPT(256, 4, 4), real, steps_(6000))
    torch.save(lm_teacher.state_dict(), path); path.with_suffix(".json").write_text(json.dumps({"seconds": lm_teacher_secs}))
lm_teacher.eval()
for p_ in lm_teacher.parameters(): p_.requires_grad_(False)
teacher_eval = lm_eval(lm_teacher)
print(f"teacher: {n_params(lm_teacher):,} parameters, {teacher_eval['bpc']:.3f} bits per character, trained in {lm_teacher_secs / 60:.1f} min")

@torch.no_grad()
def generate(model, prompts, new, greedy=False, bs=256):
    out = []
    for i in range(0, len(prompts), bs):
        x = prompts[i:i + bs]
        for _ in range(new):
            nxt = model(x[:, -CTX:])[:, -1]
            nxt = nxt.argmax(-1, keepdim=True) if greedy else torch.multinomial(nxt.softmax(-1), 1)
            x = torch.cat([x, nxt], 1)
        out.append(x)
    return torch.cat(out)

PROMPT = 16
def lm_bank(budget_chars, seed=0):
    """Everything the student may use, built from the same number of teacher-labelled characters."""
    g = torch.Generator().manual_seed(seed)
    n_seq = max(8, budget_chars // (CTX - PROMPT))
    real = chunks(tr_ids, n_seq, g)                                   # (n, CTX+1)
    t0 = time.time()
    with torch.no_grad():
        tl = torch.cat([lm_teacher(real[i:i + 256, :-1]) for i in range(0, n_seq, 256)]).log_softmax(-1).half()
    t_logits = time.time() - t0
    t0 = time.time(); sampled = generate(lm_teacher, real[:, :PROMPT], CTX + 1 - PROMPT); t_sample = time.time() - t0
    t0 = time.time(); greedy = generate(lm_teacher, real[:, :PROMPT], CTX + 1 - PROMPT, greedy=True); t_greedy = time.time() - t0
    top = tl.float().topk(5, -1)
    return {"real": real, "logp": tl, "top_idx": top.indices, "top_logp": top.values, "sampled": sampled, "greedy": greedy,
            "seconds": {"logits": t_logits, "sampled": t_sample, "greedy": t_greedy}}

def lm_loss(kind, B):
    n = len(B["real"])
    def f(model, bs, g):
        i = torch.randint(0, n, (bs,), generator=g).to(DEV)
        if kind in ("real text", "token-level", "top-5"):
            c = B["real"][i]; lg = model(c[:, :-1]).log_softmax(-1)
            if kind == "real text":
                return F.nll_loss(lg.reshape(-1, V), c[:, 1:].reshape(-1))
            if kind == "token-level":
                return F.kl_div(lg.reshape(-1, V), B["logp"][i].float().reshape(-1, V), log_target=True, reduction="batchmean")
            q = B["top_logp"][i].float().log_softmax(-1)                          # renormalized over the five
            return -(q.exp() * lg.gather(-1, B["top_idx"][i])).sum(-1).mean()
        c = B[kind.split()[0]][i]                                                   # "sampled sequences" / "greedy sequences"
        lg = model(c[:, :-1])
        return F.cross_entropy(lg[:, PROMPT - 1:].reshape(-1, V), c[:, PROMPT:].reshape(-1))
    return f

@torch.no_grad()
def distinct4(model, n=64):
    g = torch.Generator().manual_seed(7)
    s = generate(model, chunks(va_ids, n, g)[:, :PROMPT], 256)[:, PROMPT:].cpu().tolist()
    grams = [tuple(r[j:j + 4]) for r in s for j in range(len(r) - 3)]
    return len(set(grams)) / len(grams)

LM_KINDS = ["real text", "token-level", "top-5", "sampled sequences", "greedy sequences"]
lm_rows = []
PROMPT_TEXT = "A trained network"
@torch.no_grad()
def writes(model, n=220, seed=0):
    torch.manual_seed(seed)
    ids = torch.tensor([[stoi.get(c, 0) for c in PROMPT_TEXT]], device=DEV)
    out = generate(model, ids, n)[0, len(PROMPT_TEXT):].tolist()
    return PROMPT_TEXT + "".join(VOCAB[i] if i else "·" for i in out)

teacher_eval.update(distinct4=distinct4(lm_teacher), params=n_params(lm_teacher), sample=writes(lm_teacher))
print("teacher writes:", repr(teacher_eval["sample"]))
for budget in (250_000, 2_000_000):
    B = lm_bank(steps_(budget) if SMOKE else budget)
    for kind in LM_KINDS:
        for seed in range(2):
            student, secs = lm_fit(lambda: GPT(96, 2, 4), lm_loss(kind, B), steps_(4000), seed=seed)
            r = lm_eval(student, lm_teacher); r.update(kind=kind, budget=budget, seed=seed, seconds=secs, distinct4=distinct4(student),
                                                      params=n_params(student), sample=writes(student) if seed == 0 else None)
            lm_rows.append(r)
    RESULTS.setdefault("lm_bank_seconds", {})[str(budget)] = B["seconds"]
    del B
lmd = pd.DataFrame(lm_rows)
print(lmd.groupby(["budget", "kind"], sort=False)[["bpc", "kl_to_teacher", "entropy", "distinct4"]].mean().to_string(float_format=lambda v: f"{v:.4f}"))
print(f"teacher: bpc {teacher_eval['bpc']:.4f}, distinct 4-grams {teacher_eval['distinct4']:.4f}")
for r in lm_rows:
    if r["sample"] and r["budget"] == 2_000_000:
        print(f"\n{r['kind']} student writes:\n  {r['sample']!r}")
RESULTS["kd-llm"] = {"rows": lm_rows, "teacher": teacher_eval, "student_params": lm_rows[0]["params"], "chars": len(text),
                     "bank_seconds": RESULTS["lm_bank_seconds"]}

@torch.no_grad()
def per_image_ms(model, bs, reps=20):
    x = Xte[:bs]
    for _ in range(3): model(x)
    if DEV == "mps": torch.mps.synchronize()
    t0 = time.time()
    for _ in range(reps): model(x)
    if DEV == "mps": torch.mps.synchronize()
    return (time.time() - t0) / reps / bs * 1e3

student_model = fit(ARCH["cnn16"], Xtr, lambda s, idx, flip: kd(s, look(TL["bank"], idx, flip)), steps_(6000), lr=0.1, seed=0)
cnn16_train_secs = student_model[1]
t0 = time.time(); _ = bank(TL["model"], Xtr)
if DEV == "mps": torch.mps.synchronize()
label_secs = time.time() - t0
cost = {"label_seconds": label_secs, "student_train_seconds": cnn16_train_secs, "teacher_train_seconds": TL["train_seconds"],
        "transfer_images": 2 * len(Xtr), "models": {}}
for name, m in [("resnet64", TL["model"]), ("resnet32", teachers["resnet32"]["model"]), ("cnn16", student_model[0]), ("cnn8", ARCH["cnn8"]().to(DEV).eval())]:
    cost["models"][name] = {"params": n_params(m), "ms_per_image_batch256": per_image_ms(m, 256), "ms_per_image_batch1": per_image_ms(m, 1, reps=200)}
saving = cost["models"]["resnet64"]["ms_per_image_batch256"] - cost["models"]["cnn16"]["ms_per_image_batch256"]
cost["break_even_images_batch256"] = (label_secs + cnn16_train_secs) * 1e3 / saving
print(json.dumps(cost, indent=1))
print(f"break-even at batch 256: {cost['break_even_images_batch256']:,.0f} predictions")
RESULTS["kd-economics"] = cost
RESULTS["meta"]["minutes"] = (time.time() - T0) / 60
print(f"whole notebook: {RESULTS['meta']['minutes']:.1f} min")

WIDTHS = [8, 12, 16, 24, 32, 48]
size_rows, size_models = [], {}
for w in WIDTHS:
    for mode, lf in [("labels",    lambda s, idx, flip: F.cross_entropy(s, Ytr[idx])),
                     ("distilled", lambda s, idx, flip: kd(s, look(TL["bank"], idx, flip)))]:
        for seed in range(3):
            m, secs = fit(lambda: cnn(w), Xtr, lf, steps_(6000), lr=0.1, seed=seed)
            res, _ = evaluate(m, TL["te"])
            size_rows.append({"width": w, "params": n_params(m), "mode": mode, "seed": seed, "train_seconds": secs, **res})
            if seed == 0 and mode == "labels":
                size_models[w] = m

size_cost = {str(w): {"params": n_params(m), "ms_batch256": per_image_ms(m, 256), "ms_batch1": per_image_ms(m, 1, reps=200)}
             for w, m in size_models.items()}
size_cost["teacher"] = {"params": n_params(TL["model"]), "ms_batch256": per_image_ms(TL["model"], 256),
                        "ms_batch1": per_image_ms(TL["model"], 1, reps=200)}
sz = pd.DataFrame(size_rows).groupby(["width", "mode"], sort=False)[["acc"]].agg(["mean", "std"]).reset_index()
sz.columns = ["width", "mode", "acc", "sd"]
sz["params"] = sz.width.map(lambda w: size_cost[str(w)]["params"])
sz["ms/image, batch 256"] = sz.width.map(lambda w: size_cost[str(w)]["ms_batch256"])
sz["ms/image, batch 1"] = sz.width.map(lambda w: size_cost[str(w)]["ms_batch1"])
print(sz.to_string(index=False, float_format=lambda v: f"{v: .4f}"))
print("teacher:", {k: round(v, 4) for k, v in size_cost["teacher"].items()})
RESULTS["kd-size"] = {"rows": size_rows, "cost": size_cost}