# CUERDAS-020: sealed test (evaluated once), compact text baselines, stronger no-LLM decoders.
import os, sys, json, hashlib, random, math, platform
import torch, torch.nn as nn, torch.nn.functional as F, transformers
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import run018 as R
from transformers import AutoModelForCausalLM, AutoConfig
INP = os.environ['INP']; OD = os.environ['OD']; NBOOT = 10000; SEEDS = [17, 29, 43]
raw = open(f'{INP}/sealed020.json').read()
H = hashlib.sha256(raw.encode()).hexdigest(); assert H == '4c7737f66364a44786b4d9acb4686a6f31829d78cb5914f19eca1df74fe73393', H
S = json.loads(raw); NEW = S['new_words']
SUMM = open(f'{OD}/results020.jsonl', 'w'); PRED = open(f'{OD}/predictions020.jsonl', 'w')
def log(d): SUMM.write(json.dumps(d) + '\n'); SUMM.flush(); print(json.dumps(d), flush=True)
log({'env': {'python': platform.python_version(), 'torch': torch.__version__, 'transformers': transformers.__version__, 'device': torch.cuda.get_device_name() if R.DEV == 'cuda' else 'cpu'}, 'sealed_sha256': H})
for j in range(4):
    for w in NEW[j]: assert len(R.t1(w)) == 1, w
CAND = [torch.tensor(sorted({R.t1(w)[0] for w in R.V[j] + NEW[j]})) for j in range(4)]
assert all(len(CAND[j]) == len(R.V[j]) + 8 for j in range(4))
def gidx(j, w): return int((CAND[j] == R.t1(w)[0]).nonzero()[0])
SETS = ['sealed_id', 'sealed_ood', 'sealed_1_new', 'sealed_2_new', 'sealed_4_new']
FMT = {'template': R.TPL, 'kv': "action: {action}; target: {target}; status: {status}; window: {window}.",
       'abbrev': "act={action} tgt={target} st={status} win={window}", 'csv': "action,target,status,window\n{action},{target},{status},{window}",
       'values_csv': "{action}, {target}, {status}, {window}", 'values': "{action} {target} {status} {window}"}
def text_logits(msgs, fmt):
    embs, mt = [], []
    for i, m in enumerate(msgs):
        for j in range(4):
            embs.append(R.Emb[R.tok.encode(fmt.format(**dict(zip(R.FIELDS, m))) + R.SUF_MAIN.format(field=R.FIELDS[j]), add_special_tokens=False)]); mt.append((i, j))
    return R.run_rows(embs, None), mt
@torch.no_grad()
def per_ex(msgs, fmt=None, cod=None, bs=24):
    out = []
    for b in range(0, len(msgs), bs):
        mb = msgs[b:b+bs]; lg, mt = text_logits(mb, fmt) if cod is None else R.codec_logits(cod, mb, R.SUF_MAIN)
        good = [True] * len(mb)
        for r, (i, j) in enumerate(mt): good[i] &= int(lg[r][CAND[j]].argmax()) == gidx(j, mb[i][j])
        out += good
    return [float(g) for g in out]
def wilson(k, n, z=1.96):
    p = k / n; d = 1 + z*z/n; c = (p + z*z/(2*n)) / d; h = z * math.sqrt(p*(1-p)/n + z*z/(4*n*n)) / d
    return [round(100*(c-h), 1), round(100*(c+h), 1)]
def boot(a, b, seed=0):
    g = torch.Generator().manual_seed(seed); a = torch.tensor(a, dtype=torch.float64, device='cpu'); b = torch.tensor(b, dtype=torch.float64, device='cpu')
    idx = torch.randint(0, len(a), (NBOOT, len(a)), generator=g, device='cpu'); d = (a[idx].mean(1) - b[idx].mean(1)) * 100
    return round(float((a.mean()-b.mean())*100), 2), [round(float(d.quantile(.025)), 2), round(float(d.quantile(.975)), 2)]
def positions(msgs, fmt): return sum(len(R.tok.encode(fmt.format(**dict(zip(R.FIELDS, m))), add_special_tokens=False)) for m in msgs) / len(msgs) + 15
# ---- text formats ----
TXT = {}
for fn, fmt in FMT.items():
    for s in SETS:
        r = per_ex(S[s], fmt); TXT[(fn, s)] = r; k = int(sum(r))
        for i, c in enumerate(r): PRED.write(json.dumps({'route': 'text_' + fn, 'set': s, 'i': i, 'correct': bool(c)}) + '\n')
        e = {'route': 'text_' + fn, 'set': s, 'n': len(r), 'acc': k/len(r), 'wilson95': wilson(k, len(r)), 'positions_per_event': round(positions(S[s], fmt), 2)}
        if fn != 'template': e['diff_vs_template'], e['boot95'] = boot(r, TXT[('template', s)], seed=7)
        log(e)
# ---- K=2, sealed, evaluated once ----
cods = {}; K2 = {}
for sd in SEEDS:
    cod = R.PackK(2, sd); cod.load_state_dict(torch.load(f'{INP}/k2long_s{sd}.pt', map_location=R.DEV)); cod.eval().requires_grad_(False); cods[sd] = cod
    for s in SETS:
        r = per_ex(S[s], cod=cod); K2[(sd, s)] = r; k = int(sum(r))
        for i, c in enumerate(r): PRED.write(json.dumps({'route': f'k2_s{sd}', 'set': s, 'i': i, 'correct': bool(c)}) + '\n')
        d, ci = boot(r, TXT[('template', s)], seed=sd)
        log({'route': f'k2_s{sd}', 'set': s, 'n': len(r), 'acc': k/len(r), 'wilson95': wilson(k, len(r)), 'positions_per_event': 17, 'diff_vs_template': d, 'boot95': ci})
for s in SETS:
    a = [sum(K2[(sd, s)][i] for sd in SEEDS)/3 for i in range(len(S[s]))]
    for fn in FMT:
        d, ci = boot(a, TXT[(fn, s)], seed=1); log({'route': 'k2_mean3', 'set': s, 'acc': sum(a)/len(a), 'vs': 'text_' + fn, 'diff': d, 'boot95': ci})
# ---- no-LLM decoders on K=2 slots ----
def slots(cod, ms):
    with torch.no_grad(): return torch.cat([cod(torch.stack([torch.stack([R.lex(w) for w in m]) for m in ms[i:i+512]])).flatten(1) for i in range(0, len(ms), 512)])
g = random.Random(2020); AUG = []
for _ in range(20000):
    m = list(g.choice(R.FIT))
    for j in range(4):
        if g.random() < .4: m[j] = g.choice(R.POOL)
    AUG.append(m)
YA = torch.stack([torch.stack([R.lex(w) for w in m]) for m in AUG])  # [N,4,D]
CE = [F.normalize(R.Emb[CAND[j]], dim=-1) for j in range(4)]
def dec_emb(P, ms):  # P [n,4,D] predicted value embeddings -> exact-match
    ok = torch.ones(len(ms), dtype=torch.bool, device='cpu')
    for j in range(4): ok &= ((F.normalize(P[:, j], dim=-1) @ CE[j].T).argmax(1).cpu() == torch.tensor([gidx(j, m[j]) for m in ms], device='cpu'))
    return ok.float().tolist()
def train_mlp(X, Y, shuffle=False, steps=4000, seed=0):
    torch.manual_seed(seed); Y = Y[torch.randperm(len(Y))] if shuffle else Y
    mu, sdv = X.mean(0), X.std(0) + 1e-6
    net = nn.Sequential(nn.Linear(X.shape[1], 1024), nn.GELU(), nn.Linear(1024, 1024), nn.GELU(), nn.Linear(1024, 4*R.D)); opt = torch.optim.Adam(net.parameters(), 1e-3)
    for _ in range(steps):
        b = torch.randint(0, len(X), (256,)); p = net((X[b]-mu)/sdv).view(-1, 4, R.D)
        loss = (1 - F.cosine_similarity(p, Y[b], -1)).mean(); opt.zero_grad(); loss.backward(); opt.step()
    net.eval(); return lambda Z: net((Z-mu)/sdv).view(-1, 4, R.D).detach()
PROBE = {}
for sd in SEEDS:
    cod = cods[sd]; XA = slots(cod, AUG); XF = slots(cod, R.FIT); res = {'route': f'probes_k2_s{sd}'}
    A = torch.linalg.solve(XA.T@XA + 10*torch.eye(XA.shape[1]), XA.T@YA.flatten(1)); ridge = lambda Z: (Z@A).view(-1, 4, R.D)
    mlp = train_mlp(XA, YA, seed=sd); ctrl = train_mlp(XA, YA, shuffle=True, seed=sd)
    YF = torch.tensor([[R.V[j].index(m[j]) for j in range(4)] for m in R.FIT])
    torch.manual_seed(sd); cls = nn.Sequential(nn.Linear(XF.shape[1], 1024), nn.GELU(), nn.Linear(1024, sum(len(v) for v in R.V))); opt = torch.optim.Adam(cls.parameters(), 1e-3)
    off = [0, 12, 24, 40, 50]; muF, sdF = XF.mean(0), XF.std(0)+1e-6
    for _ in range(2000):
        b = torch.randint(0, len(XF), (128,)); o = cls((XF[b]-muF)/sdF); loss = sum(F.cross_entropy(o[:, off[j]:off[j+1]], YF[b, j]) for j in range(4)); opt.zero_grad(); loss.backward(); opt.step()
    XFn = F.normalize(XF, dim=-1)
    for s in ('sealed_ood', 'sealed_4_new', 'sealed_1_new'):
        ms = S[s]; Z = slots(cod, ms)
        for name, fn in (('ridge_aug', ridge), ('mlp_aug', mlp), ('mlp_ctrl_shuffled', ctrl)):
            r = dec_emb(fn(Z), ms); PROBE[(sd, name, s)] = r; res[f'{name}:{s}'] = sum(r)/len(r)
        with torch.no_grad(): o = cls((Z-muF)/sdF)
        okc = []; okk1 = []; okk5 = []; sim = F.normalize(Z, dim=-1) @ XFn.T; top = sim.topk(5, dim=1).indices.cpu()
        for i, m in enumerate(ms):
            okc.append(float(all(m[j] in R.V[j] and int(o[i, off[j]:off[j+1]].argmax()) == R.V[j].index(m[j]) for j in range(4))))
            n1 = R.FIT[int(top[i, 0])]; okk1.append(float(all(n1[j] == m[j] for j in range(4))))
            vote = [max(set(R.FIT[int(t)][j] for t in top[i]), key=lambda v: sum(R.FIT[int(t)][j] == v for t in top[i])) for j in range(4)]
            okk5.append(float(all(vote[j] == m[j] for j in range(4))))
        for name, r in (('mlp_cls_fit', okc), ('knn1', okk1), ('knn5', okk5)): PROBE[(sd, name, s)] = r; res[f'{name}:{s}'] = sum(r)/len(r)
    log(res)
for s in ('sealed_ood', 'sealed_4_new', 'sealed_1_new'):
    for name in ('ridge_aug', 'mlp_aug', 'mlp_ctrl_shuffled', 'mlp_cls_fit', 'knn1', 'knn5'):
        a = [sum(PROBE[(sd, name, s)][i] for sd in SEEDS)/3 for i in range(len(S[s]))]; l = [sum(K2[(sd, s)][i] for sd in SEEDS)/3 for i in range(len(S[s]))]
        d, ci = boot(a, l, seed=3); log({'probe_mean3': name, 'set': s, 'acc': sum(a)/len(a), 'llm_k2_acc': sum(l)/len(l), 'diff_vs_llm': d, 'boot95': ci})
# ---- random-weight receiver ----
torch.manual_seed(0)
rnd = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(R.MP)).float().eval().requires_grad_(False).to(R.DEV)
rnd.model.embed_tokens.weight.data.copy_(R.Emb.data); R.lm = rnd
for s in ('sealed_ood', 'sealed_4_new'):
    e = {'control': 'random_receiver', 'set': s, 'text_template': sum(per_ex(S[s], R.TPL))/len(S[s])}
    for sd in SEEDS: e[f'k2_s{sd}'] = sum(per_ex(S[s], cod=cods[sd]))/len(S[s])
    log(e)
SUMM.close(); PRED.close()