# CUERDAS-021 A: K=2 retrained with varied pool; sealed eval once; probes; controls.
from c021 import *
from transformers import AutoModelForCausalLM, AutoConfig
log = opener('resultsA.jsonl'); PRED = open(f'{OD}/predictionsA.jsonl', 'w'); SEEDS = [17] if SMOKE else [17, 29, 43]
log({'sealed_sha256': H, 'torch': torch.__version__, 'pool_sizes': {'words': len(WORDS), 'names': len(NAMES), 'subs': len(SUBS)}})
def ev(sysname, rowfn):
    res = {}
    for s in SETS:
        r = per_ex(rowfn, S[s], CAND); res[s] = r
        for i, c in enumerate(r): PRED.write(json.dumps({'sys': sysname, 'set': s, 'i': i, 'correct': bool(c)}) + '\n')
    PRED.flush(); log({'sys': sysname, **{s: sum(v) / len(v) for s, v in res.items()}}); return res
ev('text_template', lambda mb: text_rows(mb, R.SUF_MAIN, R.TPL))
for sd in SEEDS:
    c = R.PackK(2, sd); c.load_state_dict(torch.load(f'{INP}/k2long_s{sd}.pt', map_location=R.DEV)); c.eval()
    ev(f'k2_018_s{sd}', lambda mb, c=c: k2_rows(c, mb, R.SUF_MAIN))
cods = {}
for sd in SEEDS:
    teach = load_teacher(sd); k2 = R.PackK(2, sd); k2.load_state_dict(torch.load(f'{INP}/k2_s{sd}.pt', map_location=R.DEV))
    step, v = train(list(k2.parameters()), lambda mb, tpl: k2_rows(k2, mb, tpl), teach, sd, f'k2_021_s{sd}', log, set_mode=lambda t: k2.train(t))
    k2.eval().requires_grad_(False); torch.save(k2.state_dict(), f'{OD}/k2v_s{sd}.pt'); cods[sd] = k2
    log({'selected': f'k2_021_s{sd}', 'step': step, 'val': v})
    ev(f'k2_021_s{sd}', lambda mb, k2=k2: k2_rows(k2, mb, R.SUF_MAIN))
# ---- probes on retrained slots ----
VALS = [V[j] + NEW[j] for j in range(4)]; CEV = [F.normalize(torch.stack([R.lex(w) for w in VALS[j]]), dim=-1) for j in range(4)]
g = random.Random(2103); AUG = sample_msgs(g, 300 if SMOKE else 20000)
with torch.no_grad(): YA = torch.stack([torch.stack([R.lex(w) for w in m]) for m in AUG])
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)])
def dec(P, ms):
    ok = torch.ones(len(ms), dtype=torch.bool, device='cpu')
    for j in range(4): ok &= (F.normalize(P[:, j], dim=-1) @ CEV[j].T).argmax(1).cpu() == torch.tensor([VALS[j].index(m[j]) for m in ms], device='cpu')
    return ok.float().tolist()
def train_mlp(X, Y, shuffle=False, steps=None, seed=0):
    steps = steps or (5 if SMOKE else 4000); 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 * 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, 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, D).detach()
for sd in SEEDS:
    cod = cods[sd]; XA = slots(cod, AUG)
    A = torch.linalg.solve(XA.T @ XA + 10 * torch.eye(XA.shape[1]), XA.T @ YA.flatten(1))
    fns = {'ridge_aug': lambda Z: (Z @ A).view(-1, 4, D), 'mlp_aug': train_mlp(XA, YA, seed=sd), 'mlp_ctrl_shuffled': train_mlp(XA, YA, shuffle=True, seed=sd)}
    for name, fn in fns.items():
        for s in ('sealed_ood', 'sealed_1_new', 'sealed_2_new', 'sealed_4_new'):
            r = dec(fn(slots(cod, S[s])), S[s])
            for i, c in enumerate(r): PRED.write(json.dumps({'sys': f'{name}_s{sd}', 'set': s, 'i': i, 'correct': bool(c)}) + '\n')
            log({'probe': name, 'seed': sd, 'set': s, 'acc': sum(r) / len(r)})
# ---- 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(lambda mb: text_rows(mb, R.SUF_MAIN, R.TPL), S[s], CAND)) / len(S[s])}
    for sd in SEEDS: e[f'k2_021_s{sd}'] = sum(per_ex(lambda mb, c=cods[sd]: k2_rows(c, mb, R.SUF_MAIN), S[s], CAND)) / len(S[s])
    log(e)
PRED.close()