# Cuerdas 013b: LLM real congelado, splits disjuntos, sweep de formato, destilación, gate calibrado. CPU.
import torch, time, random, json, itertools, os, copy, statistics as st
import torch.nn as nn, torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM
torch.manual_seed(0); random.seed(0); torch.set_num_threads(8)
M=os.environ.get("MODEL","HuggingFaceTB/SmolLM2-135M")
tok=AutoTokenizer.from_pretrained(M); lm=AutoModelForCausalLM.from_pretrained(M,dtype=torch.float32).eval()
for p in lm.parameters(): p.requires_grad=False
E=lm.get_input_embeddings(); D=E.weight.shape[1]
A=["open","close","move","copy","delete","send","check"]
T=["file","door","box","email","report","car","light"]
S=["new","old","red","blue","hot","cold","big","small"]
W=["today","monday","friday","night","morning"]
FN=["action","target","state","window"]; FL=[A,T,S,W]
msgs=list(itertools.product(range(7),range(7),range(8),range(5)))
at=list(itertools.product(range(7),range(7))); sw=list(itertools.product(range(8),range(5)))
random.shuffle(at); random.shuffle(sw)
H_AT=at[:9]; H_SW=sw[:12]                      # holdout sistemático: 40x28 = 1120 cubiertas
CAL_AT=set(H_AT[:3]); TST_AT=set(H_AT[3:]); CAL_SW=set(H_SW[:4]); TST_SW=set(H_SW[4:])
def covered(m): return (m[0],m[1]) not in set(H_AT) and (m[2],m[3]) not in set(H_SW)
cov=[m for m in msgs if covered(m)]; random.shuffle(cov)
cov_train=cov[:672]; cov_cal=cov[672:896]; cov_test=cov[896:]
def ugroup(m):
    tags={("c" if (m[0],m[1]) in CAL_AT else "t") for _ in [0] if (m[0],m[1]) in set(H_AT)}
    if (m[2],m[3]) in set(H_SW): tags.add("c" if (m[2],m[3]) in CAL_SW else "t")
    return tags
unc=[m for m in msgs if not covered(m)]
unc_cal=[m for m in unc if ugroup(m)=={"c"}]; unc_test=[m for m in unc if ugroup(m)=={"t"}]
SPL=dict(cov_train=len(cov_train),cov_cal=len(cov_cal),cov_test=len(cov_test),unc_cal=len(unc_cal),unc_test=len(unc_test),unc_discarded_mixed=len(unc)-len(unc_cal)-len(unc_test))
print(SPL,flush=True)
FORMATS={
 "semicolon": lambda m:f"{A[m[0]]};{T[m[1]]};{S[m[2]]};{W[m[3]]}",
 "bare_words": lambda m:f"{A[m[0]]} {T[m[1]]} {S[m[2]]} {W[m[3]]}",
 "key_value": lambda m:f"action={A[m[0]]} target={T[m[1]]} state={S[m[2]]} window={W[m[3]]}",
 "json": lambda m:json.dumps(dict(action=A[m[0]],target=T[m[1]],state=S[m[2]],window=W[m[3]])),
 "natural": lambda m:f"{A[m[0]]} the {S[m[2]]} {T[m[1]]} {W[m[3]]}.",
 "natural_labeled": lambda m:f"Action: {A[m[0]]}. Target: {T[m[1]]}. State: {S[m[2]]}. Window: {W[m[3]]}.",
}
def ids(s): return tok(s,add_special_tokens=False,return_tensors="pt").input_ids[0]
cand=[torch.tensor([ids(" "+w)[0].item() for w in L]) for L in FL]
for c in cand: assert len(set(c.tolist()))==len(c)
Q=[ids(f"\nQ: what is the {f}?\nA:") for f in FN]; PRE=ids("Message:")
QE=None
def readout(msg_emb):  # prefill compartido (KV cache) + 4 preguntas -> lista de logits [B,ncand]
    B=msg_emb.shape[0]
    if torch.is_grad_enabled():  # entrenamiento: sin cache (deepcopy no admite grafo); misma función
        return [lm(inputs_embeds=torch.cat([E(PRE).expand(B,-1,-1),msg_emb,E(Q[f]).expand(B,-1,-1)],1)).logits[:,-1][:,cand[f]] for f in range(4)]; x=torch.cat([E(PRE).expand(B,-1,-1),msg_emb],1)
    x=torch.cat([E(PRE).expand(B,-1,-1),msg_emb],1)
    out=lm(inputs_embeds=x,use_cache=True); res=[]
    for f in range(4):
        c=copy.deepcopy(out.past_key_values)
        res.append(lm(inputs_embeds=E(Q[f]).expand(B,-1,-1),past_key_values=c).logits[:,-1][:,cand[f]])
    return res
@torch.no_grad()
def eval_text(fmt,ms,keep_logits=False):
    from collections import defaultdict
    g=defaultdict(list); ok={}; lg={}
    for m in ms: g[len(ids(" "+FORMATS[fmt](m)))].append(m)
    for L,sub in g.items():
        for i in range(0,len(sub),64):
            b=sub[i:i+64]; r=readout(E(torch.stack([ids(" "+FORMATS[fmt](m)) for m in b]))); t=torch.tensor(b)
            o=torch.ones(len(b),dtype=torch.bool)
            for f in range(4): o&=r[f].argmax(1)==t[:,f]
            for j,m in enumerate(b):
                ok[m]=bool(o[j]); 
                if keep_logits: lg[m]=[r[f][j] for f in range(4)]
    return ok,lg
class Adapter(nn.Module):
    def __init__(s,K):
        super().__init__(); s.K=K
        s.emb=nn.ModuleList([nn.Embedding(len(L),64) for L in FL])
        s.mlp=nn.Sequential(nn.Linear(256,512),nn.GELU(),nn.Linear(512,K*D))
        s.scale=nn.Parameter(torch.tensor(float(E.weight.norm(dim=1).mean())))
    def forward(s,mb):
        h=torch.cat([s.emb[i](mb[:,i]) for i in range(4)],-1)
        return F.normalize(s.mlp(h).view(-1,s.K,D),dim=-1)*s.scale
@torch.no_grad()
def eval_ad(ad,ms):
    ok={}; conf={}
    for i in range(0,len(ms),64):
        b=ms[i:i+64]; t=torch.tensor(b); r=readout(ad(t))
        o=torch.ones(len(b),dtype=torch.bool); c=torch.ones(len(b))
        for f in range(4): o&=r[f].argmax(1)==t[:,f]; c=torch.minimum(c,F.softmax(r[f],1).max(1).values)
        for j,m in enumerate(b): ok[m]=bool(o[j]); conf[m]=float(c[j])
    return ok,conf