"""Frame prototype Phase 2b — the entry channel (before model entry).

g: frame -> T5 input-embedding space, restricted to t5-whole-walk
states (length-preserving substitution). Init = closed-form lstsq
(embedding-MSE); train = behavior preservation THROUGH the frozen
encoder. Gauges: substitution fidelity, COUNTERFACTUAL entry (inject
state u' at a u site; does the frozen frame readout at that position
flip to u'?), and the random-codebook control (same fit against
random rows — must fail the counterfactual or the probe is vacuous).
"""
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, T5EncoderModel
import transformers.utils.logging as hlog

hlog.set_verbosity_error()
D = r"E:\mirel\data\bytelex\proto_frame"
DEV = "cuda"
SEED = zlib.crc32(b"frame-proto-v01") & 0xFFFFFFFF
L_T5 = 6
torch.manual_seed(SEED)
from geolip.bytelex.frame import (apply_whitening, fit_whitening,
                                  split_sites)

dump = np.load(rf"{D}\frame_dump_v2.npz")
anch = np.load(rf"{D}\frame_anchors_v01.npz")
walk = json.load(open(rf"{D}\t5_walk_of_C.json", encoding="utf-8"))
sid = dump["sid"].astype(np.int64)
KK = dump["KK"]
ok = KK[:, 0] > 0
sp = split_sites(sid[ok], seed=SEED)
gix = {k: np.flatnonzero(ok)[v] for k, v in sp.items()}

whole = np.array([w["whole"] for w in walk])
tok1 = np.array([w["ids"][0] if w["whole"] else -1 for w in walk])
print(f"[2b] t5-whole states: {int(whole.sum())}/999", flush=True)

tkA = AutoTokenizer.from_pretrained("google/flan-t5-small")
mT = T5EncoderModel.from_pretrained("google/flan-t5-small").to(DEV)
mT.eval()
for p in mT.parameters():
    p.requires_grad_(False)
emb = mT.get_input_embeddings().weight.detach()

E_C = torch.tensor(anch["E_C"], dtype=torch.float32, device=DEV)
E_Cn = F.normalize(E_C, dim=-1)
W_S = torch.tensor(anch["W_S"], dtype=torch.float32, device=DEV)
sT = float(anch["s"])

# whitening for the frame readout of encoder states (recompute, fit)
H_T = dump["H_T"][:, 0].astype(np.float64)
muT, wTw = fit_whitening(H_T[gix["fit"]])
tmu = torch.tensor(muT, dtype=torch.float32, device=DEV)
tw = torch.tensor(wTw, dtype=torch.float32, device=DEV)


def frame_readout(h):
    return F.normalize(((h - tmu) @ tw) @ W_S, dim=-1)


# per-state embedding targets for whole states — GENERALIZATION
# split: g is fit on gfit states only; counterfactual probes inject
# ONLY gprobe states the fit never saw. An interpolating fit (any
# full-rank codebook, n_rows <= 256) passes the naive probe — the
# held-out split is what lets the random control FAIL (Gate 4).
ws = np.flatnonzero(whole)
rngw = np.random.default_rng(SEED + 41)
perm = rngw.permutation(len(ws))
gfit = ws[perm[:int(0.6 * len(ws))]]
gprobe = ws[perm[int(0.6 * len(ws)):]]
print(f"[2b] whole states {len(ws)}/999 -> gfit {len(gfit)} / "
      f"gprobe {len(gprobe)}", flush=True)


def lstsq_fit(E_use):
    a = E_use[torch.tensor(gfit, device=DEV)].double().cpu().numpy()
    b = emb[torch.tensor(tok1[gfit], device=DEV)
            ].double().cpu().numpy()
    g, *_ = np.linalg.lstsq(a, b, rcond=None)
    rel = float(np.linalg.norm(a @ g - b) / np.linalg.norm(b))
    ah = E_use[torch.tensor(gprobe, device=DEV)
               ].double().cpu().numpy()
    bh = emb[torch.tensor(tok1[gprobe], device=DEV)
             ].double().cpu().numpy()
    relh = float(np.linalg.norm(ah @ g - bh) / np.linalg.norm(bh))
    return torch.tensor(g, dtype=torch.float32, device=DEV), rel, relh

G0, rel0, relh0 = lstsq_fit(E_C)
E_rand = F.normalize(torch.randn_like(E_C), dim=-1)
G_r, rel_r, relh_r = lstsq_fit(E_rand)
print(f"[2b] lstsq rel-residual fit/HELD-OUT byte={rel0:.4f}/"
      f"{relh0:.4f} random={rel_r:.4f}/{relh_r:.4f}", flush=True)

# ---- behavior-preservation training of g on whole-walk train sites
lines = open(r"E:\mirel\data\bytelex\codex_v1.txt",
             "rb").read().decode("ascii").split("\n")
gfit_set = set(int(x) for x in gfit)
cand = [k for k in gix["train"] if int(sid[k]) in gfit_set and KK[k, 1] == 1]
ev_cand = [k for k in gix["eval"] if whole[sid[k]] and KK[k, 1] == 1]
rng = np.random.default_rng(SEED + 23)
cand = rng.permutation(cand)[:4000]
ev_cand = rng.permutation(ev_cand)[:800]
G = torch.nn.Parameter(G0.clone())
opt = torch.optim.Adam([G], lr=3e-4, weight_decay=0.0)

line_tok = {}
def enc_line(li):
    if li not in line_tok:
        e = tkA(lines[li], return_offsets_mapping=True)
        line_tok[li] = (e["input_ids"], e["offset_mapping"])
    return line_tok[li]

def site_pos(k):
    li = int(dump["line"][k])
    ids, off = enc_line(li)
    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]
    return li, ids, ix[0] if ix else None

BS = 48
for ep in range(2):
    tot = nb = 0
    for i in range(0, len(cand), BS):
        ks, seqs, poss, sids_ = [], [], [], []
        for k in cand[i:i + BS]:
            li, ids, p = site_pos(int(k))
            if p is None:
                continue
            ks.append(int(k)); seqs.append(ids); poss.append(p)
            sids_.append(sid[int(k)])
        if not ks:
            continue
        mx = max(len(s) for s in seqs)
        ids = torch.full((len(ks), mx), tkA.pad_token_id,
                         dtype=torch.long)
        att = torch.zeros((len(ks), mx), dtype=torch.long)
        for j, s in enumerate(seqs):
            ids[j, :len(s)] = torch.tensor(s)
            att[j, :len(s)] = 1
        ids, att = ids.to(DEV), att.to(DEV)
        with torch.no_grad():
            clean = mT(input_ids=ids, attention_mask=att,
                       output_hidden_states=True).hidden_states[L_T5]
        base = emb[ids]
        inj = E_C[torch.tensor(sids_, device=DEV)] @ G
        x = base.clone()
        x[torch.arange(len(ks)), torch.tensor(poss, device=DEV)] = inj
        subh = mT(inputs_embeds=x, attention_mask=att,
                  output_hidden_states=True).hidden_states[L_T5]
        loss = ((subh - clean) ** 2).mean() / (clean ** 2).mean()
        opt.zero_grad(set_to_none=True)
        loss.backward()
        opt.step()
        tot += float(loss.detach()); nb += 1
    print(f"[2b] ep{ep} preserve-loss={tot / max(nb,1):.4f}", flush=True)

# ---- gauges on eval sites: fidelity + counterfactual + control
def probe(G_use, E_use, tag):
    fid = flips = flips_to_true = n = 0
    rngp = np.random.default_rng(SEED + 31)
    for i in range(0, len(ev_cand), BS):
        ks, seqs, poss, sids_, alts = [], [], [], [], []
        for k in ev_cand[i:i + BS]:
            li, ids, p = site_pos(int(k))
            if p is None:
                continue
            u = sid[int(k)]
            up = int(gprobe[rngp.integers(0, len(gprobe))])
            while up == u:
                up = int(gprobe[rngp.integers(0, len(gprobe))])
            ks.append(int(k)); seqs.append(ids); poss.append(p)
            sids_.append(u); alts.append(up)
        if not ks:
            continue
        mx = max(len(s) for s in seqs)
        ids = torch.full((len(ks), mx), tkA.pad_token_id,
                         dtype=torch.long)
        att = torch.zeros((len(ks), mx), dtype=torch.long)
        for j, s in enumerate(seqs):
            ids[j, :len(s)] = torch.tensor(s)
            att[j, :len(s)] = 1
        ids, att = ids.to(DEV), att.to(DEV)
        with torch.no_grad():
            clean = mT(input_ids=ids, attention_mask=att,
                       output_hidden_states=True).hidden_states[L_T5]
            base = emb[ids]
            ar = torch.arange(len(ks))
            pp = torch.tensor(poss, device=DEV)
            # counterfactual: inject u'
            x = base.clone()
            x[ar, pp] = E_use[torch.tensor(alts, device=DEV)] @ G_use
            subh = mT(inputs_embeds=x, attention_mask=att,
                      output_hidden_states=True).hidden_states[L_T5]
            other = (((subh - clean) ** 2).sum(-1) * att
                     ).sum() / att.sum()
            fid += float(other / (clean ** 2).sum(-1).mean())
            zr = frame_readout(subh[ar, pp])
            pred = (zr @ E_Cn.T).argmax(-1).cpu().numpy()
            flips += int((pred == np.array(alts)).sum())
            flips_to_true += int((pred == np.array(sids_)).sum())
            n += len(ks)
    return {"n": n, "flip_to_injected": round(flips / max(n, 1), 4),
            "stuck_on_true": round(flips_to_true / max(n, 1), 4),
            "rel_divergence": round(fid / max(1, (len(ev_cand)//BS+1)), 4)}

res = {"byte_codebook": probe(G, E_C, "byte"),
       "random_codebook_control": probe(G_r, E_rand, "rand"),
       "lstsq_rel_residual": {"byte_fit": round(rel0, 4),
                              "byte_heldout": round(relh0, 4),
                              "random_fit": round(rel_r, 4),
                              "random_heldout": round(relh_r, 4)},
       "g_split": {"gfit": int(len(gfit)), "gprobe": int(len(gprobe))},
       "_env": {"transformers": transformers.__version__,
                "torch": torch.__version__}}
with open(rf"{D}\entry_ledger_2b.json", "w", encoding="utf-8") as f:
    json.dump(res, f, indent=1)
np.savez(rf"{D}\entry_g_2b.npz", G=G.detach().cpu().numpy(),
         G_rand=G_r.cpu().numpy())
print(json.dumps(res, indent=1), flush=True)
print("[2b] ENTRY COMPLETE", flush=True)