"""exp019_content_retention.py — CAPACITY FOR DISTILLED CONTENT RETENTION.
When content is DISTILLED rather than directly learned, how much is retained,
through which channel, and for how long? The memory-substrate compass made a
measurable instrument — and the weak-to-strong distillation cell made real: the
teacher holds content the student lacks (true headroom on the content axis).

THE CONTENT: N fact records "\n@<key6>=<value12>\n" with random alphanumeric
keys/values — uncompletable from language statistics, so exact-match completion
IS retention. Facts are mixed into the byte stream (fact-packed blocks at
FACT_RATE against wikitext blocks).

THE TEACHER: the certified bed model (addr_msl64 aleph substrate) trained
TEACHER_STEPS on the mix; gated on its own recall (the gate doubles as the
substrate's DIRECT capacity datum at each N).

THE CHANNELS (fresh student each, STUDENT_STEPS budget):
  direct       — ground-truth CE on the mix (ceiling: learning, not distillation)
  kd_facts     — CE on wikitext blocks; on fact blocks the ONLY signal is the
                 teacher's logits (KL) — pure distilled content
  kd_general   — CE + KL to teacher on CLEAN wikitext only; facts never shown —
                 does content leak through logits without exposure?
  book_implant — teacher's farmed codebook implanted (trainable) into a fresh
                 student, clean-stream training — do the anchors carry
                 byte-content? (the open question from the exp014 implant studies, asked directly)
  none         — clean-stream only (floor)
THE AXES: capacity N in {64, 256, 1024} (main channels); retention = recall
right after training AND after INTERFERE_STEPS further clean-stream steps
(the forgetting measurement). KD alpha 1.0 here is LEGAL: the teacher has real
headroom on the content axis (the inverse-evolution failure was alpha 1.0 at
NEAR-PARITY — regime, not constant).

Preregistered forks:
  F1 capacity curve: direct recall vs N = the substrate's raw content capacity.
  F2 distillation tax: kd_facts vs direct at each N (what survives the logit
     channel).
  F3 leakage: kd_general recall > floor => content crosses on clean text alone.
  F4 anchors: book_implant recall ~ floor => codebooks do not carry byte
     content (mean-shape/content question closed in the direct sense).
  F5 half-life: post-interference retention per channel (does distilled content
     decay faster than learned content?).
Riders: pure Adam wd=0; GPU-only verdicts; >=2 seeds; Colab-safe. Paste order:
geolip_vitals -> ar_differentiation_bed -> exp014_genetic_distillation ->
this file.
"""
from __future__ import annotations
import json
import math
import os
import string
import torch
import torch.nn.functional as F

if "ByteLM" not in globals():
    try:
        from ar_differentiation_bed import ByteLM, _wikitext_bytes, _batch, VOCAB
        from exp014_genetic_distillation import implant_book
        from geolip_vitals import anchor_drift
    except ImportError:
        _here = globals().get("__file__")
        if _here is None:
            raise ImportError("paste geolip_vitals + ar_differentiation_bed + "
                              "exp014_genetic_distillation first")
        import sys, pathlib
        sys.path.insert(0, str(pathlib.Path(_here).parent))
        from ar_differentiation_bed import ByteLM, _wikitext_bytes, _batch, VOCAB
        from exp014_genetic_distillation import implant_book
        from geolip_vitals import anchor_drift

DATA_ROOT = os.environ.get("GEOLIP_DATA", "./data")
EXP19_DIR = os.path.join(DATA_ROOT, "exp019")

KEY_LEN, VAL_LEN = 6, 12
FACT_RATE = 0.5              # fraction of training blocks drawn from fact stream
TEACHER_STEPS = 4000
STUDENT_STEPS = 2000
INTERFERE_STEPS = 1000
ALNUM = (string.ascii_lowercase + string.digits).encode()


def make_facts(n: int, seed: int = 0):
    """N records '\\n@<key>=<value>\\n'; returns (records list, fact byte stream)."""
    g = torch.Generator().manual_seed(4000 + seed)
    recs = []
    seen = set()
    while len(recs) < n:
        k = bytes(ALNUM[i] for i in torch.randint(len(ALNUM), (KEY_LEN,),
                                                  generator=g))
        if k in seen:
            continue
        seen.add(k)
        v = bytes(ALNUM[i] for i in torch.randint(len(ALNUM), (VAL_LEN,),
                                                  generator=g))
        recs.append((k, v))
    return recs


def fact_stream(recs, copies: int = 50, seed: int = 0) -> torch.Tensor:
    """Byte stream of shuffled fact records (each record appears `copies` times)."""
    g = torch.Generator().manual_seed(5000 + seed)
    order = torch.cat([torch.randperm(len(recs), generator=g)
                       for _ in range(copies)])
    blob = b"".join(b"\n@" + recs[i][0] + b"=" + recs[i][1] + b"\n"
                    for i in order.tolist())
    return torch.frombuffer(bytearray(blob), dtype=torch.uint8).clone()


def _mix_batch(tr, fs, batch, block, device, g):
    """Blocks drawn from the fact stream with prob FACT_RATE, else wikitext.
    Returns (x, y, fact_mask (B,)) — mask marks fact-sourced rows."""
    xw, yw = _batch(tr, batch, block, device, g)
    xf, yf = _batch(fs, batch, block, device, g)
    m = (torch.rand(batch, generator=g) < FACT_RATE).to(device)
    x = torch.where(m[:, None], xf, xw)
    y = torch.where(m[:, None], yf, yw)
    return x, y, m


@torch.no_grad()
def recall(model, recs, device="cuda", max_eval: int = 256,
           batch: int = 64) -> dict:
    """Exact-match greedy completion: prompt '\\n@<key>=' -> VAL_LEN bytes."""
    model = model.to(device).eval()
    recs = recs[:max_eval]
    prompts = torch.stack([torch.frombuffer(
        bytearray(b"\n@" + k + b"="), dtype=torch.uint8).long()
        for k, _ in recs]).to(device)
    outs = []
    for i in range(0, len(recs), batch):
        x = prompts[i:i + batch]
        for _ in range(VAL_LEN):
            nxt = model(x)[:, -1].argmax(-1, keepdim=True)
            x = torch.cat([x, nxt], dim=1)
        outs.append(x[:, -VAL_LEN:].cpu())
    got = torch.cat(outs)
    tgt = torch.stack([torch.frombuffer(bytearray(v), dtype=torch.uint8).long()
                       for _, v in recs])
    byte_acc = (got == tgt).float().mean().item()
    exact = (got == tgt).all(dim=1).float().mean().item()
    return {"exact": round(exact, 4), "byte_acc": round(byte_acc, 4)}


def train_stream(model, tr, va, fs=None, teacher=None, channel="direct",
                 steps=2000, batch=32, block=256, device="cuda", seed=0):
    """One training run under a channel's signal routing (docstring above)."""
    g = torch.Generator().manual_seed(seed)
    model = model.to(device)
    if teacher is not None:
        teacher = teacher.to(device).eval()
    opt = torch.optim.Adam(model.parameters(), lr=3e-4, weight_decay=0.0)
    for step in range(1, steps + 1):
        if channel in ("direct", "kd_facts") and fs is not None:
            x, y, m = _mix_batch(tr, fs, batch, block, device, g)
        else:                                   # clean wikitext stream
            x, y = _batch(tr, batch, block, device, g)
            m = torch.zeros(batch, dtype=torch.bool, device=device)
        logits = model(x)
        if channel == "kd_facts":
            # ground truth on wiki rows only; teacher logits are the ONLY
            # signal on fact rows (pure distilled content)
            ce_rows = ~m
            loss = torch.tensor(0.0, device=device)
            if ce_rows.any():
                loss = F.cross_entropy(logits[ce_rows].reshape(-1, VOCAB),
                                       y[ce_rows].reshape(-1))
            if m.any():
                with torch.no_grad():
                    tp = F.softmax(teacher(x[m]), -1)
                loss = loss + F.kl_div(F.log_softmax(logits[m], -1), tp,
                                       reduction="batchmean")
        elif channel == "kd_general":
            loss = F.cross_entropy(logits.reshape(-1, VOCAB), y.reshape(-1))
            with torch.no_grad():
                tp = F.softmax(teacher(x), -1)
            loss = loss + F.kl_div(F.log_softmax(logits, -1), tp,
                                   reduction="batchmean")
        else:                                   # direct / book_implant / none
            loss = F.cross_entropy(logits.reshape(-1, VOCAB), y.reshape(-1))
        opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
    model.eval()
    with torch.no_grad():
        ls = []
        for _ in range(10):
            xv, yv = _batch(va, batch, block, device, g)
            ls.append(F.cross_entropy(model(xv).reshape(-1, VOCAB),
                                      yv.reshape(-1)).item())
    return sum(ls) / len(ls) / math.log(2)


def run_retention(n_facts=(64, 256, 1024), seeds=(0, 1), device="cuda"):
    if not torch.cuda.is_available():
        raise RuntimeError("verdict runs are GPU-only")
    os.makedirs(EXP19_DIR, exist_ok=True)
    tr, va = _wikitext_bytes(DATA_ROOT)
    ledger = open(os.path.join(EXP19_DIR, "ledger.jsonl"), "a", encoding="utf-8")

    def log(rec):
        ledger.write(json.dumps(rec) + "\n"); ledger.flush()
        print(f"[19 {rec['channel']} N={rec['n']} s{rec['seed']}] "
              f"recall={rec['recall']} after_interf={rec.get('recall_interf')} "
              f"bpb={rec['bpb']}", flush=True)

    for seed in seeds:
        for n in n_facts:
            recs = make_facts(n, seed=seed)
            fs = fact_stream(recs, seed=seed)
            # ---- teacher (also the DIRECT capacity datum at TEACHER_STEPS)
            torch.manual_seed(seed)
            teacher = ByteLM("addr_msl64")
            t_bpb = train_stream(teacher, tr, va, fs=fs, channel="direct",
                                 steps=TEACHER_STEPS, device=device, seed=seed)
            t_rec = recall(teacher, recs, device=device)
            log({"exp": "19", "channel": "teacher", "n": n, "seed": seed,
                 "steps": TEACHER_STEPS, "recall": t_rec, "bpb": round(t_bpb, 4)})
            torch.save({"n": n, "seed": seed,
                        "state_dict": {k: v.cpu() for k, v in
                                       teacher.state_dict().items()}},
                       os.path.join(EXP19_DIR, f"teacher_N{n}_s{seed}.pt"))
            # ---- channels
            chans = ["direct", "kd_facts", "kd_general"]
            if n == 256:
                chans += ["book_implant", "none"]
            for ch in chans:
                torch.manual_seed(1000 + seed)
                m = ByteLM("addr_msl64")
                if ch == "book_implant":
                    implant_book(m.head_addr,
                                 teacher.head_addr.codebook.detach().cpu())
                bpb = train_stream(
                    m, tr, va, fs=fs if ch in ("direct", "kd_facts") else None,
                    teacher=teacher if ch.startswith("kd") else None,
                    channel=ch, steps=STUDENT_STEPS, device=device,
                    seed=1000 + seed)
                r0 = recall(m, recs, device=device)
                # retention under interference: further CLEAN-stream training
                bpb2 = train_stream(m, tr, va, channel="none",
                                    steps=INTERFERE_STEPS, device=device,
                                    seed=2000 + seed)
                r1 = recall(m, recs, device=device)
                log({"exp": "19", "channel": ch, "n": n, "seed": seed,
                     "steps": STUDENT_STEPS, "recall": r0, "recall_interf": r1,
                     "bpb": round(bpb, 4), "bpb_after_interf": round(bpb2, 4)})
                del m
                torch.cuda.empty_cache()
            del teacher
            torch.cuda.empty_cache()
    ledger.close()


# ==================== exp019b — GENERALIZATION block =========================
# Rule-bearing content: value = fixed random substitution cipher applied to the
# key, extended to VAL_LEN (v[i] = subst(k[i % KEY_LEN])). Teacher sees
# N_TRAIN rule-keys; N_TEST keys are HELD OUT. Held-out recall = the RULE
# generalizing, not the list. Sharp question: does the logit channel transfer
# the rule better than it transfers the rote list? Plus prompt-format variants
# (content vs surface form disentangled).

def make_rule_facts(n_train: int = 256, n_test: int = 128, seed: int = 0):
    g = torch.Generator().manual_seed(6000 + seed)
    subst = {ALNUM[i]: ALNUM[j] for i, j in
             enumerate(torch.randperm(len(ALNUM), generator=g).tolist())}
    keys, seen = [], set()
    while len(keys) < n_train + n_test:
        k = bytes(ALNUM[i] for i in torch.randint(len(ALNUM), (KEY_LEN,),
                                                  generator=g))
        if k not in seen:
            seen.add(k)
            keys.append(k)
    def val(k):
        return bytes(subst[k[i % KEY_LEN]] for i in range(VAL_LEN))
    train = [(k, val(k)) for k in keys[:n_train]]
    test = [(k, val(k)) for k in keys[n_train:]]
    return train, test


@torch.no_grad()
def recall_fmt(model, recs, fmt: bytes = b"\n@%s=", device="cuda",
               max_eval: int = 256, batch: int = 64) -> dict:
    """recall() under an arbitrary prompt format (b'\\n@%s=' = the training
    format; variants probe surface-form generalization)."""
    model = model.to(device).eval()
    recs = recs[:max_eval]
    proms = [torch.frombuffer(bytearray(fmt.replace(b"%s", k)),
                              dtype=torch.uint8).long() for k, _ in recs]
    L = max(p.numel() for p in proms)
    # left-pad with newlines to equal length (causal — padding is prefix noise)
    prompts = torch.stack([torch.cat([torch.full((L - p.numel(),), 10,
                                                 dtype=torch.long), p])
                           for p in proms]).to(device)
    outs = []
    for i in range(0, len(recs), batch):
        x = prompts[i:i + batch]
        for _ in range(VAL_LEN):
            nxt = model(x)[:, -1].argmax(-1, keepdim=True)
            x = torch.cat([x, nxt], dim=1)
        outs.append(x[:, -VAL_LEN:].cpu())
    got = torch.cat(outs)
    tgt = torch.stack([torch.frombuffer(bytearray(v), dtype=torch.uint8).long()
                       for _, v in recs])
    return {"exact": round((got == tgt).all(dim=1).float().mean().item(), 4),
            "byte_acc": round((got == tgt).float().mean().item(), 4)}


FMT_TRAIN = b"\n@%s="
FMT_VARIANT = b" @%s= "          # never seen in training: pure format shift


def run_generalization(n_train: int = 256, n_test: int = 128, seeds=(0, 1),
                       device="cuda"):
    if not torch.cuda.is_available():
        raise RuntimeError("verdict runs are GPU-only")
    os.makedirs(EXP19_DIR, exist_ok=True)
    tr, va = _wikitext_bytes(DATA_ROOT)
    ledger = open(os.path.join(EXP19_DIR, "ledger.jsonl"), "a", encoding="utf-8")

    def gauges(model, train_recs, test_recs):
        return {"train": recall_fmt(model, train_recs, FMT_TRAIN, device=device),
                "heldout": recall_fmt(model, test_recs, FMT_TRAIN, device=device),
                "train_varfmt": recall_fmt(model, train_recs, FMT_VARIANT,
                                           device=device)}

    for seed in seeds:
        train_recs, test_recs = make_rule_facts(n_train, n_test, seed=seed)
        fs = fact_stream(train_recs, seed=seed)         # held-out NEVER streamed
        torch.manual_seed(seed)
        teacher = ByteLM("addr_msl64")
        t_bpb = train_stream(teacher, tr, va, fs=fs, channel="direct",
                             steps=TEACHER_STEPS, device=device, seed=seed)
        gt = gauges(teacher, train_recs, test_recs)
        rec = {"exp": "19b", "channel": "teacher", "n": n_train, "seed": seed,
               "steps": TEACHER_STEPS, "gauges": gt, "bpb": round(t_bpb, 4)}
        ledger.write(json.dumps(rec) + "\n"); ledger.flush()
        print(f"[19b teacher s{seed}] {gt} bpb={t_bpb:.4f}", flush=True)
        torch.save({"seed": seed, "state_dict": {k: v.cpu() for k, v in
                                                 teacher.state_dict().items()}},
                   os.path.join(EXP19_DIR, f"rule_teacher_s{seed}.pt"))
        for ch in ("direct", "kd_facts", "kd_general"):
            torch.manual_seed(1000 + seed)
            m = ByteLM("addr_msl64")
            bpb = train_stream(
                m, tr, va, fs=fs if ch in ("direct", "kd_facts") else None,
                teacher=teacher if ch.startswith("kd") else None,
                channel=ch, steps=STUDENT_STEPS, device=device, seed=1000 + seed)
            g0 = gauges(m, train_recs, test_recs)
            bpb2 = train_stream(m, tr, va, channel="none",
                                steps=INTERFERE_STEPS, device=device,
                                seed=2000 + seed)
            g1 = gauges(m, train_recs, test_recs)
            rec = {"exp": "19b", "channel": ch, "n": n_train, "seed": seed,
                   "steps": STUDENT_STEPS, "gauges": g0, "gauges_interf": g1,
                   "bpb": round(bpb, 4), "bpb_after_interf": round(bpb2, 4)}
            ledger.write(json.dumps(rec) + "\n"); ledger.flush()
            print(f"[19b {ch} s{seed}] {g0} interf_heldout="
                  f"{g1['heldout']} bpb={bpb:.4f}", flush=True)
            torch.save({"channel": ch, "seed": seed,
                        "state_dict": {k: v.cpu() for k, v in
                                       m.state_dict().items()}},
                       os.path.join(EXP19_DIR, f"rule_{ch}_s{seed}.pt"))
            del m
            torch.cuda.empty_cache()
        del teacher
        torch.cuda.empty_cache()
    ledger.close()


def smoke():
    recs = make_facts(8, seed=0)
    assert len(recs) == 8 and all(len(k) == KEY_LEN and len(v) == VAL_LEN
                                  for k, v in recs)
    fs = fact_stream(recs, copies=3, seed=0)
    assert fs.dtype == torch.uint8 and fs.numel() == 3 * 8 * (KEY_LEN + VAL_LEN + 4)
    m = ByteLM("addr_msl64", d=96, layers=2, block=64)
    r = recall(m, recs, device="cpu", max_eval=8, batch=4)
    assert 0.0 <= r["exact"] <= 1.0 and 0.0 <= r["byte_acc"] <= 1.0
    g = torch.Generator().manual_seed(0)
    x, y, mask = _mix_batch(torch.randint(0, 256, (50000,),
                                          dtype=torch.uint8, generator=g),
                            fs, 8, 64, "cpu", g)
    assert x.shape == (8, 64) and mask.shape == (8,)
    assert (x[:, 1:] == y[:, :-1]).all()          # stream alignment
    # 19b: rule facts are rule-consistent + disjoint; variant recall runs
    tr8, te4 = make_rule_facts(8, 4, seed=0)
    assert len(tr8) == 8 and len(te4) == 4
    assert not set(k for k, _ in tr8) & set(k for k, _ in te4)
    k0, v0 = tr8[0]
    assert len(v0) == VAL_LEN and v0[:KEY_LEN] == v0[KEY_LEN:2 * KEY_LEN]
    rv = recall_fmt(m, tr8, FMT_VARIANT, device="cpu", max_eval=8, batch=4)
    assert 0.0 <= rv["exact"] <= 1.0
    print(f"exp019 smoke passed (untrained recall exact={r['exact']} "
          f"byte={r['byte_acc']} ~ chance; 19b rule+variant OK)")


def _in_notebook():
    try:
        get_ipython()  # type: ignore[name-defined]  # noqa: F821
        return True
    except NameError:
        return False


if __name__ == "__main__":
    smoke() if not _in_notebook() else (smoke(),
        print("Notebook: run_retention() on GPU."))