"""Frame prototype v0.1 — training runner (phases 1, 2, 2b + twin).

Implements the v0.1 spec (plan: 2026-08-18_lexicon_translation_
learning_DRAFT.md). Byte-state arm: E_C init from C's own byte-
structural features; teacher aligns TO the codebook (Route B);
logit-adjusted cosine InfoNCE (prior in frozen bias, never in
geometry); student phase against frozen E_C with unit-sphere MSE.
Index-space twin: same losses/budget, codebook keyed by student
anchor token id, state posterior via P(state|key). Gauges: per-class
identification (prior-free + prior-added), masked-reading delta,
teacher ceiling, oracle-codebook control, parity baseline, drift,
RSA frame attribution, gallery decay, counterfactual entry.
"""
import json
import sys
import zlib

sys.path.insert(0, r"E:\mirel\geolip-bytelex")
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
import numpy as np
import torch
import torch.nn.functional as F
import transformers
from transformers import AutoTokenizer

from geolip.bytelex.frame import (apply_whitening, fit_whitening,
                                  gallery_decay, split_sites,
                                  state_byte_features, procrustes,
                                  top_k_accuracy)

D = r"E:\mirel\data\bytelex\proto_frame"
WORDS = r"E:\mirel\data\bytelex\words_of_C.json"
DEV = "cuda"
SEED = zlib.crc32(b"frame-proto-v01") & 0xFFFFFFFF
EPOCHS_T, EPOCHS_S, BS = 6, 6, 512
torch.manual_seed(SEED)

states = json.load(open(WORDS, encoding="utf-8"))
t5_walk = json.load(open(rf"{D}\t5_walk_of_C.json", encoding="utf-8"))
dump = np.load(rf"{D}\frame_dump_v2.npz")
ctx = np.load(rf"{D}\ctx_profiles.npz")
sid = dump["sid"].astype(np.int64)
KK = dump["KK"]
ok = KK[:, 0] > 0
NS = 999

# walk-based class axis (frame-purity fix 7a)
def cls_of(s):
    t = s["text"]
    if t.isdigit():
        return "digit"
    if t[0].isupper():
        return "Name"
    return "word"
CLS = np.array([cls_of(s) for s in states])
site_cls = CLS[sid]

sp = split_sites(sid[ok], seed=SEED)
gix = {k: np.flatnonzero(ok)[v] for k, v in sp.items()}
print(f"[F] sites fit/train/eval = "
      f"{len(gix['fit'])}/{len(gix['train'])}/{len(gix['eval'])}",
      flush=True)

# ---------- Phase 0 stats (fp64, fit/train only)
H_B = dump["H_B"][:, 0].astype(np.float64)      # byte-anchored first
H_T = dump["H_T"][:, 0].astype(np.float64)
muB, wB = fit_whitening(H_B[gix["fit"]])
muT, wT = fit_whitening(H_T[gix["fit"]])
ZB = apply_whitening(H_B, muB, wB)
ZT = apply_whitening(H_T, muT, wT)
evB = np.linalg.eigvalsh(np.cov(H_B[gix["fit"]] - muB, rowvar=False))
prB = float((evB.sum() ** 2) / (evB ** 2).sum() / len(evB))

# byte-structural codebook init (PCA-whiten fp64 -> 256, unit rows)
feats = state_byte_features([s["text"].encode() for s in states],
                            ctx["prev_ctx"], ctx["next_ctx"])
fc = feats - feats.mean(0)
u, sv, vt = np.linalg.svd(fc.astype(np.float64), full_matrices=False)
E0 = u[:, :256] * 1.0                            # decorrelated rows
E0 = E0 / np.linalg.norm(E0, axis=1, keepdims=True)

# teacher whitened state-means on TRAIN
def state_means(z, ix):
    m = np.zeros((NS, z.shape[1]))
    for s in range(NS):
        r = ix[sid[ix] == s]
        if len(r):
            m[s] = z[r].mean(0)
    return m
mT_train = state_means(ZB, gix["train"])
R_T = procrustes(mT_train, E0)                   # 768 -> 256, fp64

# priors (frozen, TRAIN)
cnt = np.bincount(sid[gix["train"]], minlength=NS).astype(np.float64)
b_prior = np.log(np.maximum(cnt, 0.5) / cnt.sum())

# ---------- torch setup
tZB = torch.tensor(ZB, dtype=torch.float32, device=DEV)
tZT = torch.tensor(ZT, dtype=torch.float32, device=DEV)
tsid = torch.tensor(sid, device=DEV)
E_C = torch.nn.Parameter(torch.tensor(E0, dtype=torch.float32,
                                      device=DEV))
W_Tm = torch.nn.Parameter(torch.tensor(R_T, dtype=torch.float32,
                                       device=DEV))
s_T = torch.nn.Parameter(torch.tensor(10.0, device=DEV))
bp = torch.tensor(b_prior, dtype=torch.float32, device=DEV)


def id_logits(z, W, E, s, prior):
    zf = F.normalize(z @ W, dim=-1)
    lg = s.clamp(1, 100) * (zf @ F.normalize(E, dim=-1).T)
    return lg + bp if prior else lg


def train_phase(params, loss_fn, ixs, epochs, tag):
    opt = torch.optim.Adam(params, lr=1e-3, weight_decay=0.0)
    rng = np.random.default_rng(SEED + 7)
    for ep in range(epochs):
        order = rng.permutation(ixs)
        tot = nb = 0
        for i in range(0, len(order), BS):
            b = torch.tensor(order[i:i + BS], device=DEV)
            opt.zero_grad(set_to_none=True)
            l = loss_fn(b)
            l.backward()
            opt.step()
            with torch.no_grad():
                E_C.data = F.normalize(E_C.data, dim=-1)
            tot += float(l)
            nb += 1
        print(f"[F {tag}] ep{ep} loss={tot / nb:.4f}", flush=True)


# ---------- Phase 1: teacher anchors (L_id^T)
def loss_T(b):
    lg = id_logits(tZB[b], W_Tm, E_C, s_T, True)
    return F.cross_entropy(lg, tsid[b])

E_before = E_C.detach().cpu().numpy().copy()
train_phase([W_Tm, E_C, s_T], loss_T, gix["train"], EPOCHS_T, "P1")
E_C.requires_grad_(False)
s_T.requires_grad_(False)

# drift by count decile + dead rows
drift = np.linalg.norm(E_C.detach().cpu().numpy() - E_before, axis=1)
dec = np.digitize(cnt, np.quantile(cnt, np.linspace(0, 1, 11)[1:-1]))
drift_dec = [round(float(drift[dec == d].mean()), 4) for d in range(10)]

# ---------- Phase 2: student aligns to frozen center
mS_train = state_means(ZT, gix["train"])
R_S0 = procrustes(mS_train, E_C.detach().cpu().numpy().astype(
    np.float64))
W_Sm = torch.nn.Parameter(torch.tensor(R_S0, dtype=torch.float32,
                                       device=DEV))
with torch.no_grad():
    zT_frame = F.normalize(tZB @ W_Tm, dim=-1)   # frozen teacher

def eval_id(z, W, ix, prior, k=1):
    with torch.no_grad():
        lg = id_logits(z[torch.tensor(ix, device=DEV)], W, E_C, s_T,
                       prior).cpu().numpy()
    return top_k_accuracy(lg, sid[ix], k), lg

parity1, _ = eval_id(tZT, W_Sm, gix["eval"], False)

def loss_S(b):
    lg = id_logits(tZT[b], W_Sm, E_C, s_T, True)
    zs = F.normalize(tZT[b] @ W_Sm, dim=-1)
    return (F.cross_entropy(lg, tsid[b])
            + ((zs - zT_frame[b]) ** 2).sum(-1).mean())

train_phase([W_Sm], loss_S, gix["train"], EPOCHS_S, "P2")

# ---------- gauges
led = {"_env": {"transformers": transformers.__version__,
                "torch": torch.__version__, "seed": int(SEED)},
       "splits": {k: int(len(v)) for k, v in gix.items()},
       "whitening": {"participation_ratio_B": round(prB, 4)},
       "drift_by_decile": drift_dec,
       "dead_rows": int((drift < 1e-4).sum()),
       "parity_baseline_eval_top1": round(float(parity1), 4)}
ev = gix["eval"]
accs = {}
for tag, z, W in (("teacher", tZB, W_Tm), ("student", tZT, W_Sm)):
    for prior in (False, True):
        a1, lg = eval_id(z, W, ev, prior)
        a5 = top_k_accuracy(lg, sid[ev], 5)
        key = f"{tag}_{'prior' if prior else 'balanced'}"
        accs[key] = {"top1": round(a1, 4), "top5": round(a5, 4)}
        per = {}
        for c in ("word", "Name", "digit"):
            m = site_cls[ev] == c
            if m.any():
                per[c] = {"n": int(m.sum()),
                          "top1": round(top_k_accuracy(
                              lg[m], sid[ev][m]), 4)}
        accs[key]["by_class"] = per
led["identification"] = accs

# masked-reading delta (control): masked readouts through same maps
MB = apply_whitening(dump["M_B"].astype(np.float64), muB, wB)
MT = apply_whitening(dump["M_T"].astype(np.float64), muT, wT)
tMB = torch.tensor(MB, dtype=torch.float32, device=DEV)
tMT = torch.tensor(MT, dtype=torch.float32, device=DEV)
delta = {}
for tag, zu, zm, W in (("teacher", tZB, tMB, W_Tm),
                       ("student", tZT, tMT, W_Sm)):
    au, lgu = eval_id(zu, W, ev, False)
    am, lgm = eval_id(zm, W, ev, False)
    per = {}
    for c in ("word", "Name", "digit"):
        m = site_cls[ev] == c
        if m.any():
            per[c] = round(top_k_accuracy(lgu[m], sid[ev][m])
                           - top_k_accuracy(lgm[m], sid[ev][m]), 4)
    delta[tag] = {"overall": round(au - am, 4), "by_class": per}
led["reading_delta"] = delta

# oracle-codebook control (student vs never-trained student means)
mS_unit = mS_train / np.maximum(
    np.linalg.norm(mS_train, axis=1, keepdims=True), 1e-9)
scr = ZT[ev] @ mS_unit.T
led["oracle_codebook_student_top1"] = round(
    top_k_accuracy(scr, sid[ev]), 4)

# RSA frame attribution: whose geometry is final E_C?
Ef = F.normalize(E_C, dim=-1).cpu().numpy().astype(np.float64)
def rsa(a, b, rng):
    iu = np.triu_indices(NS, 1)
    pick = rng.choice(len(iu[0]), size=100_000, replace=False)
    va = (a @ a.T)[iu][pick]
    vb = (b @ b.T)[iu][pick]
    ra, rb = np.argsort(np.argsort(va)), np.argsort(np.argsort(vb))
    return float(np.corrcoef(ra, rb)[0, 1])
rr = np.random.default_rng(SEED + 13)
mB_unit = mT_train / np.maximum(
    np.linalg.norm(mT_train, axis=1, keepdims=True), 1e-9)
led["rsa"] = {"E_C_vs_byte_features": round(rsa(Ef, E0, rr), 4),
              "E_C_vs_teacher_means": round(rsa(Ef, mB_unit, rr), 4)}

# gallery decay (prior-free, student)
_, lgS = eval_id(tZT, W_Sm, ev, False)
sub = rr.choice(len(ev), size=min(2000, len(ev)), replace=False)
led["gallery_decay_student"] = gallery_decay(lgS[sub], sid[ev][sub],
                                             seed=SEED + 17)

# ---------- index-space twin (same losses/budget, key = t5 anchor id)
tkA = AutoTokenizer.from_pretrained("google/flan-t5-small")
lines = open(r"E:\mirel\data\bytelex\codex_v1.txt",
             "rb").read().decode("ascii").split("\n")
anchor_id = np.full(len(sid), -1, dtype=np.int64)
by_line = {}
for k in range(len(sid)):
    by_line.setdefault(int(dump["line"][k]), []).append(k)
for li, ks in by_line.items():
    e = tkA(lines[li], add_special_tokens=False,
            return_offsets_mapping=True)
    off = e["offset_mapping"]
    for k in ks:
        lo, hi = int(dump["lo"][k]), int(dump["hi"][k])
        ix = [i for i, (s, t) in enumerate(off)
              if t > s and s < hi and t > lo]
        if ix:
            anchor_id[k] = e["input_ids"][ix[0]]
keys, key_inv = np.unique(anchor_id[anchor_id >= 0],
                          return_inverse=False), None
key_of = {int(a): i for i, a in enumerate(keys)}
ksite = np.array([key_of.get(int(a), -1) for a in anchor_id])
NK = len(keys)
# P(state|key) from TRAIN
post = np.zeros((NK, NS))
for k in gix["train"]:
    if ksite[k] >= 0:
        post[ksite[k], sid[k]] += 1
post = post / np.maximum(post.sum(1, keepdims=True), 1)
mK_train = np.zeros((NK, 768))
for kk in range(NK):
    r = gix["train"][ksite[gix["train"]] == kk]
    if len(r):
        mK_train[kk] = ZB[r].mean(0)
uk, sk, vk = np.linalg.svd(
    (mK_train - mK_train.mean(0)).astype(np.float64),
    full_matrices=False)
EK0 = uk[:, :256]
EK0 = EK0 / np.maximum(np.linalg.norm(EK0, axis=1, keepdims=True),
                       1e-9)
E_K = torch.nn.Parameter(torch.tensor(EK0, dtype=torch.float32,
                                      device=DEV))
R_K = procrustes(mK_train, EK0)                  # same-quality init
W_Ki = torch.nn.Parameter(torch.tensor(R_K, dtype=torch.float32,
                                       device=DEV))
s_K = torch.nn.Parameter(torch.tensor(10.0, device=DEV))
tks = torch.tensor(ksite, device=DEV)
cntK = np.bincount(ksite[gix["train"]][ksite[gix["train"]] >= 0],
                   minlength=NK).astype(np.float64)
bK = torch.tensor(np.log(np.maximum(cntK, .5) / cntK.sum()),
                  dtype=torch.float32, device=DEV)
trK = gix["train"][ksite[gix["train"]] >= 0]

def loss_K(b):
    zf = F.normalize(tZB[b] @ W_Ki, dim=-1)
    lg = s_K.clamp(1, 100) * (zf @ F.normalize(E_K, dim=-1).T) + bK
    return F.cross_entropy(lg, tks[b])

train_phase([W_Ki, E_K, s_K], loss_K, trK, EPOCHS_T, "TWIN")
evK = ev[ksite[ev] >= 0]
with torch.no_grad():
    zf = F.normalize(tZB[torch.tensor(evK, device=DEV)] @ W_Ki, -1)
    lgK = (s_K.clamp(1, 100) * zf @ F.normalize(E_K, -1).T).cpu().numpy()
state_scores = np.exp(lgK - lgK.max(1, keepdims=True)) @ post
twin = {"overall": round(top_k_accuracy(state_scores, sid[evK]), 4),
        "n_keys": int(NK)}
for c in ("word", "Name", "digit"):
    m = site_cls[evK] == c
    if m.any():
        twin[c] = {"n": int(m.sum()),
                   "top1": round(top_k_accuracy(state_scores[m],
                                                sid[evK][m]), 4)}
led["index_twin_state_id"] = twin

with open(rf"{D}\frame_ledger_v01.json", "w", encoding="utf-8") as f:
    json.dump(led, f, indent=1)
np.savez(rf"{D}\frame_anchors_v01.npz",
         E_C=E_C.detach().cpu().numpy(),
         W_T=W_Tm.detach().cpu().numpy(),
         W_S=W_Sm.detach().cpu().numpy(),
         s=float(s_T), b_prior=b_prior, E0=E0)
print(json.dumps(led["identification"], indent=1)[:1500], flush=True)
print("[F] TWIN:", json.dumps(twin), flush=True)
print("[F] RSA:", json.dumps(led["rsa"]),
      "| reading_delta:", json.dumps(led["reading_delta"]), flush=True)
print("[F] PHASES 1-2 + TWIN COMPLETE", flush=True)