# CUERDAS-015: transferencia (valores/consultas nuevas), vocabulario aleatorizado, K=1 vs K=4. CPU.
import json,random,time,re,sys,statistics as st,itertools
import torch,torch.nn as nn,torch.nn.functional as F
from transformers import AutoModelForCausalLM,AutoTokenizer,AutoConfig
torch.set_num_threads(8)
import os
SRC=os.environ.get('SRC','/workspace/cuerdas014/x/cuerdas014'); MP=os.environ.get('MP','/workspace/cuerdas014/x/models/smollm2-135m-instruct'); OD=os.environ.get('OD','/workspace/cuerdas015')
DEV='cuda' if torch.cuda.is_available() else 'cpu'
tok=AutoTokenizer.from_pretrained(MP); lm=AutoModelForCausalLM.from_pretrained(MP,torch_dtype=torch.float32).eval().requires_grad_(False).to(DEV)
if DEV=='cuda': torch.set_default_device('cuda')
Emb=lm.get_input_embeddings().weight; D=Emb.shape[1]
meta=json.load(open(f'{SRC}/data/expanded.json')); V=meta['vocabulary']; FIELDS=['action','target','status','window']
rows=meta['rows']; SPL=meta['splits']
TPL="The action is {action}. The target is {target}. The status is {status}. The window is {window}."
SUF_MAIN="\nAccording to the record, the {field} is"
SUF_TRAIN=[SUF_MAIN,"\nThe {field} is","\nIn this record, the {field} is","\nQuestion: what is the {field}? Answer:"]
SUF_HELD={"recorded":"\nThe recorded {field} is","value":"\nThe {field} value is","past":"\nThe {field} was"}
def t1(w): return tok.encode(' '+w,add_special_tokens=False)
# ---- vocabularios nuevos (nunca en entrenamiento) ----
NEW={'action':['inspect','repair','cancel','approve','archive','scan','print','send'],
     'target':['sensor','battery','printer','server','truck','door','screen','cable'],
     'status':['broken','stable','late','empty','full','warm','missing','urgent'],
     'window':['night','evening','weekend','noon','spring','summer','winter','midnight']}
PHRASE={'action':['shut down','back up','sign off','log out'],'target':['air filter','fuel tank','water heater','hard drive'],
        'status':['out of stock','on hold','under review','half full'],'window':['next week','last night','this month','later today']}
# limpiar: primer token único dentro de (vocab original ∪ nuevos ∪ frases) por campo
for j,f in enumerate(FIELDS):
    used={t1(w)[0] for w in V[j]}; keep=[];keepp=[]
    for w in NEW[f]:
        if len(t1(w))==1 and t1(w)[0] not in used and w not in V[j]: keep.append(w); used.add(t1(w)[0])
    for w in PHRASE[f]:
        if len(t1(w))>1 and t1(w)[0] not in used: keepp.append(w); used.add(t1(w)[0])
    NEW[f]=keep; PHRASE[f]=keepp
BANNED={t1(w)[0] for f in FIELDS for w in NEW[f]+PHRASE[f]}|{t1(w)[0] for vs in V for w in vs}
# pool de palabras aleatorias para entrenamiento (meta-aprendizaje de binding)
POOL=[]
for t,i in tok.get_vocab().items():
    if t.startswith('Ġ') and re.fullmatch(r'[a-z]{3,9}',t[1:]) and i not in BANNED: POOL.append(t[1:])
random.Random(0).shuffle(POOL); POOL_VAL=POOL[3000:3200]; POOL=POOL[:3000]
# ---- splits: los de 014 + test FRESCO nuevo del universo no usado ----
used_vals={tuple(r['values']) for r in rows}
fit=[tuple(rows[i]['values']) for i in SPL['fit']]; seen_at={v[:2] for v in fit}; seen_sw={v[2:] for v in fit}
rng=random.Random(15015); uni=[v for v in itertools.product(*[range(len(x)) for x in V]) if v not in used_vals]; rng.shuffle(uni)
fresh_id=[v for v in uni if v[:2] in seen_at and v[2:] in seen_sw][:200]; fresh_ood=[v for v in uni if not(v[:2] in seen_at and v[2:] in seen_sw)][:400]
def W(v): return [V[j][v[j]] for j in range(4)]
FIT=[W(v) for v in fit]; VAL=[W(tuple(rows[i]['values'])) for i in SPL['validation']]
TEST014={'id':[W(tuple(rows[i]['values'])) for i in SPL['id_test']],'ood':[W(tuple(rows[i]['values'])) for i in SPL['ood_test']]}
FRESH={'id':[W(v) for v in fresh_id],'ood':[W(v) for v in fresh_ood]}
r2=random.Random(77)
def novel_set(nnew,n=96,phr=False):  # nnew campos con valor nuevo (los demás del vocab), tomados de FRESH
    out=[]
    for k in range(n):
        base=list(FRESH['ood'][k%len(FRESH['ood'])]); js=r2.sample(range(4),nnew)
        for j in js: base[j]=r2.choice(PHRASE[FIELDS[j]] if phr else NEW[FIELDS[j]])
        out.append(base)
    return out
TRANSFER={'1_new':novel_set(1),'2_new':novel_set(2),'4_new':novel_set(4),'1_phrase':novel_set(1,phr=True)}
CAND=[sorted({t1(w)[0] for w in V[j]+NEW[f]+PHRASE[f]}) for j,f in enumerate(FIELDS)]  # opciones válidas: abiertas (vocab+nuevos)
CAND=[torch.tensor(c) for c in CAND]
def lex(w): return Emb[t1(w)].mean(0)
def gold_idx(j,w): return int((CAND[j]==t1(w)[0]).nonzero()[0])
# ---- codecs ----
class Binder(nn.Module):  # K=4: ranura por campo, binding de cualquier palabra a un rol
    def __init__(s,seed):
        super().__init__(); torch.manual_seed(seed); s.K=4
        anc=torch.stack([Emb[t1(f)[0]] for f in FIELDS]); s.roles=nn.Parameter(anc.clone())
        s.net=nn.Sequential(nn.LayerNorm(D),nn.Linear(D,128),nn.GELU(),nn.Linear(128,D)); nn.init.normal_(s.net[-1].weight,std=.002); nn.init.zeros_(s.net[-1].bias)
        s.register_buffer('sc',Emb.std().clone())
    def forward(s,L): c=L+s.roles[None]; return c+s.sc*s.net(c)        # L: [B,4,D]
class Pack1(nn.Module):  # K=1: empaqueta los 4 campos en UNA posición
    def __init__(s,seed):
        super().__init__(); torch.manual_seed(seed); s.K=1
        s.roles=nn.Parameter(torch.randn(4,D)*.02)
        s.net=nn.Sequential(nn.LayerNorm(4*D),nn.Linear(4*D,512),nn.GELU(),nn.Linear(512,D)); s.register_buffer('sc',Emb.std().clone())
        s.mix=nn.Parameter(torch.full((4,),.25))
    def forward(s,L):
        c=L+s.roles[None]; base=(s.mix[None,:,None]*c).sum(1)
        return (base+s.sc*s.net(c.flatten(1)))[:,None]
class Their014(nn.Module):  # su RoleAdapter, con léxico extensible (mismo cálculo que su forward)
    def __init__(s,seed):
        super().__init__(); ck=torch.load(f'{SRC}/checkpoints/expanded_role_ce_k4_seed{seed}/best.pt',map_location=DEV,weights_only=False)['adapter_state']
        s.K=4; s.roles=nn.Parameter(ck['roles']); s.sc=ck['scale']
        s.net=nn.Sequential(nn.LayerNorm(D),nn.Linear(D,128),nn.GELU(),nn.Linear(128,D)); s.net.load_state_dict({k[4:]:v for k,v in ck.items() if k.startswith('net.')})
    def forward(s,L): c=L+s.roles[None]; return c+s.sc*s.net(c)
# ---- lectura por el LLM congelado ----
def sufids(tpl,j): return tok.encode(tpl.format(field=FIELDS[j]),add_special_tokens=False)
def run_rows(embs,tails):  # listas de tensores [L,D] -> logits última posición (padding a la derecha)
    L=max(e.shape[0] for e in embs); B=len(embs)
    x=torch.zeros(B,L,D); m=torch.zeros(B,L,dtype=torch.long)
    for i,e in enumerate(embs): x[i,:e.shape[0]]=e; m[i,:e.shape[0]]=1
    h=lm.model(inputs_embeds=x,attention_mask=m).last_hidden_state[torch.arange(B),m.sum(1)-1]
    return lm.lm_head(h)
def codec_logits(cod,msgs,tpl,js=range(4)):
    S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs]))
    embs=[];meta_=[]
    for i,m in enumerate(msgs):
        for j in js: embs.append(torch.cat([S[i],Emb[sufids(tpl,j)]])); meta_.append((i,j))
    return run_rows(embs,None),meta_
def text_logits(msgs,tpl):
    embs=[];meta_=[]
    for i,m in enumerate(msgs):
        for j in range(4): embs.append(Emb[tok.encode(TPL.format(**dict(zip(FIELDS,m)))+tpl.format(field=FIELDS[j]),add_special_tokens=False)]); meta_.append((i,j))
    return run_rows(embs,None),meta_
@torch.no_grad()
def acc(route,msgs,tpl=SUF_MAIN,cod=None,bs=24):
    ok=[]
    for b in range(0,len(msgs),bs):
        mb=msgs[b:b+bs]; lg,mt=(text_logits(mb,tpl) if route=='text' else codec_logits(cod,mb,tpl))
        good=torch.ones(len(mb),dtype=torch.bool)
        for r,(i,j) in enumerate(mt): good[i]&=int(lg[r][CAND[j]].argmax())==gold_idx(j,mb[i][j])
        ok+=good.tolist()
    return sum(ok)/len(ok)
@torch.no_grad()
def nolm(cod,msgs):  # anti-lookup: decodificar valor desde las ranuras por coseno, sin LLM
    S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs])); ok=0
    for i,m in enumerate(msgs):
        g=True
        for j in range(4):
            cands=F.normalize(Emb[CAND[j]],dim=-1); s=S[i,j] if S.shape[1]==4 else S[i].mean(0)
            g&=int((F.normalize(s,dim=-1)@cands.T).argmax())==gold_idx(j,m[j])
        ok+=g
    return ok/len(msgs)
def train(cod,steps,seed,p_rand=.5,p_phr=.15,bs=12):
    g=random.Random(seed); opt=torch.optim.AdamW(cod.parameters(),lr=2e-3,weight_decay=1e-3)
    for s in range(steps):
        msgs=[]
        for _ in range(bs):
            m=list(g.choice(FIT))
            for j in range(4):
                u=g.random()
                if u<p_phr: m[j]=g.choice(POOL)+' '+g.choice(POOL)
                elif u<p_phr+p_rand: m[j]=g.choice(POOL)
            msgs.append(m)
        tpl=g.choice(SUF_TRAIN); lg,mt=codec_logits(cod,msgs,tpl)
        gold=torch.tensor([t1(msgs[i][j])[0] for i,j in mt]); loss=F.cross_entropy(lg,gold)  # CE sobre TODO el vocabulario: copiar la palabra ligada al rol
        opt.zero_grad(); loss.backward(); opt.step()
    return cod.eval()
OUT=OD+'/results018.jsonl'
def log(d): open(OUT,'a').write(json.dumps(d)+"\n"); print(json.dumps(d),flush=True)
def evaluate_all(name,cod=None,quick=False):
    route='text' if cod is None else 'codec'; r={'name':name}
    r['val']=acc(route,VAL[:96],cod=cod)
    if not quick:
        for k,v in FRESH.items(): r['fresh_'+k]=acc(route,v,cod=cod)
        for k,v in TRANSFER.items(): r['transfer_'+k]=acc(route,v,cod=cod)
        sub=FRESH['id'][:48]+FRESH['ood'][:48]
        for k,t in SUF_HELD.items(): r['query_'+k]=acc(route,sub,t,cod=cod)
        if cod is not None: r['noLLM_fresh_ood']=nolm(cod,FRESH['ood'][:200]); r['noLLM_transfer_4new']=nolm(cod,TRANSFER['4_new'])
    log(r); return r

def _branch(cache):
    if hasattr(cache,'to_legacy_cache'): cache=cache.to_legacy_cache()
    return tuple((k.repeat_interleave(4,0),v.repeat_interleave(4,0)) for k,v in cache)
@torch.no_grad()
def one_event(route,m,cod,dev):
    tails=[sufids(SUF_MAIN,j) for j in range(4)]; sn=0
    while all(t[sn]==tails[0][sn] for t in tails): sn+=1
    if route=='text':
        pre=Emb[tok.encode(TPL.format(**dict(zip(FIELDS,m))),add_special_tokens=False)+tails[0][:sn]][None]
    else:
        pre=torch.cat([cod(torch.stack([lex(w) for w in m])[None].to(dev))[0],Emb[tails[0][:sn]]])[None]
    o=lm.model(inputs_embeds=pre,use_cache=True)
    h=lm.model(input_ids=torch.tensor([t[sn:] for t in tails],device=dev),past_key_values=_branch(o.past_key_values),use_cache=True).last_hidden_state[:,-1]
    out=[int(lm.lm_head(h)[j][CAND[j].to(dev)].argmax()) for j in range(4)]
    return pre.shape[1]+4*len(tails[0][sn:])
def latency(cod,dev,n=40,reps=3):
    sync=(lambda: torch.cuda.synchronize()) if dev!='cpu' else (lambda:None)
    ev=FRESH['ood'][:n]; T={'text':[],'codec':[]}; P={}
    for m in ev[:5]: one_event('text',m,cod,dev); one_event('codec',m,cod,dev)
    rr=random.Random(3)
    for _ in range(reps):
        for m in ev:
            for r in rr.sample(['text','codec'],2):
                sync(); t0=time.perf_counter(); P[r]=one_event(r,m,cod,dev); sync(); T[r].append(1e3*(time.perf_counter()-t0))
    q=lambda x:sorted(x)[int(.9*len(x))]
    return {k:dict(median_ms=st.median(v),p90_ms=q(v),positions=P[k]) for k,v in T.items()}

class PackK(nn.Module):  # superposición aprendida: ranura_k = sum_j a_kj (e_j + U_j V_j^T e_j + rol_j) + MLP(concat)
    def __init__(s,K,seed,r=64):
        super().__init__(); torch.manual_seed(seed); s.K=K
        s.roles=nn.Parameter(torch.randn(4,D)*.02); s.U=nn.Parameter(torch.randn(4,D,r)*.01); s.V=nn.Parameter(torch.randn(4,D,r)*.01)
        s.a=nn.Parameter(torch.full((K,4),1.0/4*K if K<4 else 1.0))
        s.net=nn.Sequential(nn.LayerNorm(4*D),nn.Linear(4*D,1024),nn.GELU(),nn.Linear(1024,K*D)); nn.init.normal_(s.net[-1].weight,std=.002); nn.init.zeros_(s.net[-1].bias)
        s.register_buffer('sc',Emb.std().clone())
    def forward(s,L):
        c=L+torch.einsum('bjd,jdr,jer->bje',L,s.V,s.U)+s.roles[None]
        base=torch.einsum('kj,bjd->bkd',s.a,c)
        return base+s.sc*s.net(c.flatten(1)).view(-1,s.K,D)
def train_kd(cod,steps,seed,p_rand=.3,bs=12,lr=1e-3):
    g=random.Random(seed); opt=torch.optim.AdamW(cod.parameters(),lr=lr,weight_decay=1e-3)
    sch=torch.optim.lr_scheduler.OneCycleLR(opt,max_lr=lr,total_steps=steps,pct_start=.05)
    for s_ in range(steps):
        msgs=[]
        for _ in range(bs):
            m=list(g.choice(FIT))
            for j in range(4):
                if g.random()<p_rand: m[j]=g.choice(POOL)
            msgs.append(m)
        tpl=g.choice(SUF_TRAIN)
        S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs]))
        embs=[];hs=[];tg=[]
        for i,m in enumerate(msgs):
            for j in range(4): embs.append(torch.cat([S[i],Emb[sufids(tpl,j)]])); tg.append(t1(m[j])[0])
        with torch.no_grad():
            te=[Emb[tok.encode(TPL.format(**dict(zip(FIELDS,m)))+tpl.format(field=FIELDS[j]),add_special_tokens=False)] for m in msgs for j in range(4)]
            Lt=max(e.shape[0] for e in te); x=torch.zeros(len(te),Lt,D); mk=torch.zeros(len(te),Lt,dtype=torch.long)
            for i,e in enumerate(te): x[i,:e.shape[0]]=e; mk[i,:e.shape[0]]=1
            ht=lm.model(inputs_embeds=x,attention_mask=mk).last_hidden_state[torch.arange(len(te)),mk.sum(1)-1]; lt=lm.lm_head(ht)
        L=max(e.shape[0] for e in embs); x=torch.zeros(len(embs),L,D); mk=torch.zeros(len(embs),L,dtype=torch.long)
        for i,e in enumerate(embs): x[i,:e.shape[0]]=e; mk[i,:e.shape[0]]=1
        hsd=lm.model(inputs_embeds=x,attention_mask=mk).last_hidden_state[torch.arange(len(embs)),mk.sum(1)-1]; ls=lm.lm_head(hsd)
        loss=F.cross_entropy(ls,torch.tensor(tg))+F.kl_div(F.log_softmax(ls,-1),F.log_softmax(lt,-1),log_target=True,reduction='batchmean')+.5*(1-F.cosine_similarity(hsd,ht,-1).mean())
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(cod.parameters(),1.0); opt.step(); sch.step()
        if s_%500==0: print(json.dumps({'step':s_,'loss':float(loss)}),flush=True)
    return cod.eval()
@torch.no_grad()
def probe(cod):  # anti-lookup fuerte: sonda lineal (ridge, sin LLM) entrenada en slots de FIT, evaluada en FRESH ood y 4_new
    def slots(ms): return cod(torch.stack([torch.stack([lex(w) for w in m]) for m in ms])).flatten(1)
    Xtr=torch.cat([slots(FIT[i:i+100]) for i in range(0,len(FIT),100)]); res={}
    for name,ms in (('fresh_ood',FRESH['ood']),('transfer_4_new',TRANSFER['4_new'])):
        Xte=slots(ms); ok=torch.ones(len(ms),dtype=torch.bool,device='cpu')
        for j in range(4):  # regresión al embedding del valor, decodificación por coseno contra opciones (abierta a palabras nuevas)
            Y=torch.stack([lex(m[j]) for m in FIT]); A=torch.linalg.solve(Xtr.T@Xtr+10*torch.eye(Xtr.shape[1]),Xtr.T@Y)
            P=F.normalize(Xte@A,dim=-1)@F.normalize(Emb[CAND[j]],dim=-1).T
            ok&=P.argmax(1).cpu()==torch.tensor([gold_idx(j,m[j]) for m in ms],device='cpu')
        res['probe_'+name]=float(ok.float().mean())
    return res

rv=random.Random(99)
VAL_T=[[rv.choice(POOL_VAL) for j in range(4)] for _ in range(64)]   # validación de transferencia: palabras nunca usadas en entrenamiento ni en test
CAND_FULL=None
@torch.no_grad()
def acc_open(cod,msgs,bs=16):  # exactitud sobre TODO el vocabulario (palabras de validación no están en CAND)
    ok=[]
    for b in range(0,len(msgs),bs):
        mb=msgs[b:b+bs]; lg,mt=codec_logits(cod,mb,SUF_MAIN); good=torch.ones(len(mb),dtype=torch.bool)
        for r,(i,j) in enumerate(mt): good[i]&=int(lg[r].argmax())==t1(mb[i][j])[0]
        ok+=good.tolist()
    return sum(ok)/len(ok)
def readout_h(cod,msgs,tpl):
    S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in msgs])); embs=[]
    for i,m in enumerate(msgs):
        for j in range(4): embs.append(torch.cat([S[i],Emb[sufids(tpl,j)]]))
    L=max(e.shape[0] for e in embs); x=torch.zeros(len(embs),L,D); mk=torch.zeros(len(embs),L,dtype=torch.long)
    for i,e in enumerate(embs): x[i,:e.shape[0]]=e; mk[i,:e.shape[0]]=1
    h=lm.model(inputs_embeds=x,attention_mask=mk).last_hidden_state[torch.arange(len(embs)),mk.sum(1)-1]; return h,lm.lm_head(h)
def distill(stu,teach,steps,seed,tagname,bs=12,lr=1e-3,every=250):
    g=random.Random(seed); opt=torch.optim.AdamW(stu.parameters(),lr=lr,weight_decay=1e-3)
    sch=torch.optim.lr_scheduler.OneCycleLR(opt,max_lr=lr,total_steps=steps,pct_start=.05); best=(-1,None,0)
    for s_ in range(1,steps+1):
        msgs=[]
        for _ in range(bs):
            m=list(g.choice(FIT))
            for j in range(4):
                u=g.random()
                if u<.1: m[j]=g.choice(POOL)+' '+g.choice(POOL)
                elif u<.5: m[j]=g.choice(POOL)
            msgs.append(m)
        tpl=g.choice(SUF_TRAIN)
        with torch.no_grad(): ht,lt=readout_h(teach,msgs,tpl)
        hs,ls=readout_h(stu,msgs,tpl); tg=torch.tensor([t1(m[j])[0] for m in msgs for j in range(4)])
        loss=F.cross_entropy(ls,tg)+F.kl_div(F.log_softmax(ls,-1),F.log_softmax(lt,-1),log_target=True,reduction='batchmean')+.5*(1-F.cosine_similarity(hs,ht,-1).mean())
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(stu.parameters(),1.0); opt.step(); sch.step()
        if s_%every==0:
            stu.eval(); v1=acc('codec',VAL[:96],cod=stu); v2=acc_open(stu,VAL_T); stu.train()
            sc=v1+v2; log({'run':tagname,'step':s_,'val_contract':v1,'val_transfer_open':v2})
            if sc>best[0]: best=(sc,{k:v.detach().clone() for k,v in stu.state_dict().items()},s_)
    stu.load_state_dict(best[1]); stu.eval(); return stu,best[2]
def grow_down(src,K):  # inicializar K pequeño desde K mayor (curriculum)
    dst=PackK(K,0); sd=src.state_dict(); dd=dst.state_dict()
    for k in ('roles','U','V','net.0.weight','net.0.bias','net.1.weight','net.1.bias','sc'): dd[k]=sd[k].clone()
    Ks=src.K; dd['a']=sd['a'].view(K,Ks//K,4).sum(1) if Ks%K==0 else sd['a'].mean(0,keepdim=True).expand(K,4).clone()
    W=sd['net.3.weight'].view(Ks,D,-1); b=sd['net.3.bias'].view(Ks,D)
    dd['net.3.weight']=W.view(K,Ks//K,D,-1).sum(1).reshape(K*D,-1); dd['net.3.bias']=b.view(K,Ks//K,D).sum(1).reshape(-1)
    dst.load_state_dict(dd); return dst
@torch.no_grad()
def batched_latency(cod,bsz,reps=5):
    sync=(lambda: torch.cuda.synchronize()) if DEV=='cuda' else (lambda:None)
    ms=FRESH['ood'][:bsz]; res={}
    for route in ('text','codec'):
        for _ in range(2): (text_logits(ms,SUF_MAIN) if route=='text' else codec_logits(cod,ms,SUF_MAIN))
        T=[]
        for _ in range(reps):
            sync(); t0=time.perf_counter(); (text_logits(ms,SUF_MAIN) if route=='text' else codec_logits(cod,ms,SUF_MAIN)); sync(); T.append(1e3*(time.perf_counter()-t0))
        res[route]=dict(median_ms=st.median(T),max_ms=max(T))
    return res

# validación de transferencia RESTRINGIDA: opciones = vocab del campo + 40 palabras reservadas por campo (disjuntas de train y test)
CANDV=[];VPOOL=[]
for j in range(4):
    ws=POOL_VAL[j*50:(j+1)*50]; ids=sorted({t1(w)[0] for w in V[j]}|{t1(w)[0] for w in ws}); CANDV.append(torch.tensor(ids)); VPOOL.append(ws)
rv2=random.Random(123); VAL_TR=[[rv2.choice(VPOOL[j]) for j in range(4)] for _ in range(96)]
@torch.no_grad()
def acc_valT(cod,bs=24):
    ok=[]
    for b in range(0,len(VAL_TR),bs):
        mb=VAL_TR[b:b+bs]; lg,mt=codec_logits(cod,mb,SUF_MAIN); good=torch.ones(len(mb),dtype=torch.bool)
        for r,(i,j) in enumerate(mt): good[i]&=int(CANDV[j][lg[r][CANDV[j]].argmax()])==t1(mb[i][j])[0]
        ok+=good.tolist()
    return sum(ok)/len(ok)
def distill2(stu,teach,steps,seed,tagname,bs=12,lr=5e-4,every=250):
    g=random.Random(seed+1000); opt=torch.optim.AdamW(stu.parameters(),lr=lr,weight_decay=1e-3)
    sch=torch.optim.lr_scheduler.OneCycleLR(opt,max_lr=lr,total_steps=steps,pct_start=.05)
    stu.eval(); v1=acc('codec',VAL[:96],cod=stu); v2=acc_valT(stu); stu.train()
    best=(v1+v2,{k:v.detach().clone() for k,v in stu.state_dict().items()},0); log({'run':tagname,'step':0,'val_contract':v1,'val_transfer_restricted':v2})
    for s_ in range(1,steps+1):
        msgs=[]
        for _ in range(bs):
            m=list(g.choice(FIT))
            for j in range(4):
                u=g.random()
                if u<.1: m[j]=g.choice(POOL)+' '+g.choice(POOL)
                elif u<.4: m[j]=g.choice(POOL)
            msgs.append(m)
        tpl=g.choice(SUF_TRAIN)
        with torch.no_grad(): ht,lt=readout_h(teach,msgs,tpl)
        hs,ls=readout_h(stu,msgs,tpl); tg=torch.tensor([t1(m[j])[0] for m in msgs for j in range(4)])
        loss=F.cross_entropy(ls,tg)+F.kl_div(F.log_softmax(ls,-1),F.log_softmax(lt,-1),log_target=True,reduction='batchmean')+.5*(1-F.cosine_similarity(hs,ht,-1).mean())
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(stu.parameters(),1.0); opt.step(); sch.step()
        if s_%every==0:
            stu.eval(); v1=acc('codec',VAL[:96],cod=stu); v2=acc_valT(stu); stu.train()
            log({'run':tagname,'step':s_,'val_contract':v1,'val_transfer_restricted':v2})
            if v1+v2>best[0]: best=(v1+v2,{k:v.detach().clone() for k,v in stu.state_dict().items()},s_)
    stu.load_state_dict(best[1]); stu.eval(); return stu,best[2]
@torch.no_grad()
def batched_shared(cod,bsz,reps=5):  # prefijo KV compartido en batch, ambas rutas
    sync=(lambda: torch.cuda.synchronize()) if DEV=='cuda' else (lambda:None)
    ms=FRESH['ood'][:bsz]; tails=[sufids(SUF_MAIN,j) for j in range(4)]; sn=0
    while all(t[sn]==tails[0][sn] for t in tails): sn+=1
    TT=torch.tensor([t[sn:] for t in tails]).repeat(bsz,1)
    def run(route):
        if route=='text':
            ids=[tok.encode(TPL.format(**dict(zip(FIELDS,m))),add_special_tokens=False)+tails[0][:sn] for m in ms]
            assert len(set(map(len,ids)))==1; pre=Emb[torch.tensor(ids)]
        else:
            S=cod(torch.stack([torch.stack([lex(w) for w in m]) for m in ms])); pre=torch.cat([S,Emb[tails[0][:sn]][None].expand(bsz,-1,-1)],1)
        o=lm.model(inputs_embeds=pre,use_cache=True)
        h=lm.model(input_ids=TT,past_key_values=_branch(o.past_key_values),use_cache=True).last_hidden_state[:,-1]
        lg=lm.lm_head(h); return pre.shape[1]
    res={}
    for route in ('text','codec'):
        for _ in range(2): run(route)
        T=[]
        for _ in range(reps): sync(); t0=time.perf_counter(); p=run(route); sync(); T.append(1e3*(time.perf_counter()-t0))
        res[route]=dict(median_ms=st.median(T),max_ms=max(T),prefix_positions=p)
    return res
if __name__=='__main__':
    INP=os.environ.get('INP',''); seeds=[int(x) for x in os.environ.get('SEEDS','17,29,43').split(',')]
    for sd in seeds:
        t0=time.time(); teach=Binder(sd); teach.load_state_dict(torch.load(f'{INP}/binder4_1000_s{sd}.pt',map_location=DEV)); teach.eval().requires_grad_(False)
        k2=PackK(2,sd); k2.load_state_dict(torch.load(f'{INP}/k2_s{sd}.pt',map_location=DEV))
        log({'name':f'k2_017_s{sd}','val_transfer_restricted':acc_valT(k2.eval())})
        k2,b=distill2(k2,teach,int(os.environ.get('S2',3000)),sd,f'k2long_s{sd}'); torch.save(k2.state_dict(),f'{OD}/k2long_s{sd}.pt')
        r=evaluate_all(f'k2long_s{sd}',k2); log({'name':r['name'],'best_step':b,'train_s':round(time.time()-t0),**probe(k2)})
        if sd==seeds[0]:
            for bz in (1,16,64): log({'latency_shared':'k2long','batch':bz,'gpu':torch.cuda.get_device_name() if DEV=='cuda' else 'cpu',**batched_shared(k2,bz)})