# 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
R={"model":M,"device":"CPU del box (8 hilos, compartido, sin GPU)","splits":SPL}
# 1) sweep de formato SOLO en calibración (muestra de cov_cal + unc_cal)
random.seed(2); calS=random.sample(cov_cal,100)+random.sample(unc_cal,min(100,len(unc_cal)))
sweep={}
for fmt in FORMATS:
    t0=time.time(); ok,_=eval_text(fmt,calS)
    sweep[fmt]=dict(cal_acc=sum(ok.values())/len(ok),avg_tokens=sum(len(ids(" "+FORMATS[fmt](m))) for m in calS)/len(calS),secs=round(time.time()-t0,1))
    print(fmt,sweep[fmt],flush=True)
R["format_sweep_cal"]=sweep
best=max(sweep,key=lambda f:(round(sweep[f]["cal_acc"],2),-sweep[f]["avg_tokens"]))  # exactitud primero, luego tokens
R["teacher_format"]=best; FMT=best; print("teacher",best,flush=True)
pos=lambda m:len(ids(" "+FORMATS[FMT](m)))
# 2) logits del maestro en train (para destilar)
_,TL=eval_text(FMT,cov_train,keep_logits=True)
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
def train(K,epochs=int(os.environ.get("EPOCHS",8)),bs=32,Tk=2.0):
    ad=Adapter(K); opt=torch.optim.AdamW(ad.parameters(),3e-3); data=list(cov_train)
    for ep in range(epochs):
        random.shuffle(data); tot=0
        for i in range(0,len(data),bs):
            b=data[i:i+bs]; mb=torch.tensor(b); r=readout(ad(mb)); loss=0
            for f in range(4):
                tl=torch.stack([TL[m][f] for m in b])
                loss=loss+F.cross_entropy(r[f],mb[:,f])+Tk*Tk*F.kl_div(F.log_softmax(r[f]/Tk,1),F.softmax(tl/Tk,1),reduction="batchmean")
            opt.zero_grad(); loss.backward(); opt.step(); tot+=loss.item()
        print(f"K={K} ep{ep} loss {tot:.2f}",flush=True)
    return ad.eval()
@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
cal_all=cov_cal+unc_cal; test_all=cov_test+unc_test
tok_cal,_=eval_text(FMT,cal_all); tok_test,_=eval_text(FMT,test_all)
R["baseline_test"]=dict(acc=sum(tok_test.values())/len(test_all),pos_avg=sum(map(pos,test_all))/len(test_all),acc_cov=sum(tok_test[m] for m in cov_test)/len(cov_test),acc_unc=sum(tok_test[m] for m in unc_test)/len(unc_test))
print(R["baseline_test"],flush=True)
ADS={}
for K in (1,2):
    ts=time.time(); ad=train(K); ADS[K]=ad
    okc,cc=eval_ad(ad,cal_all); okt,ct=eval_ad(ad,test_all)
    # gate: cobertura (manifiesto) Y confianza>=thr; thr elegido en calibración: el mínimo con degradación <=1pp vs texto
    def hyb(okA,conf,okT,ms,thr):
        nat=[m for m in ms if covered(m) and conf[m]>=thr]; s=set(nat)
        acc=sum((okA[m] if m in s else okT[m]) for m in ms)/len(ms)
        p=sum((K if m in s else pos(m)) for m in ms)/len(ms); return acc,p,len(nat)/len(ms)
    base_cal=sum(tok_cal.values())/len(cal_all); thr_ok=None
    for thr in [i/100 for i in range(0,100,2)]:
        a,_,_=hyb(okc,cc,tok_cal,cal_all,thr)
        if base_cal-a<=0.01: thr_ok=thr; break
    if thr_ok is None: thr_ok=1.01
    a,p,fr=hyb(okt,ct,tok_test,test_all,thr_ok)
    # sólo-cobertura (sin confianza) para referencia
    a0,p0,fr0=hyb(okt,ct,tok_test,test_all,0.0)
    R[f"K{K}"]=dict(params=sum(x.numel() for x in ad.parameters()),train_s=round(time.time()-ts),
        adapter_acc_cov_test=sum(okt[m] for m in cov_test)/len(cov_test),adapter_acc_unc_test=sum(okt[m] for m in unc_test)/len(unc_test),
        gate_threshold_from_cal=thr_ok,hybrid_test=dict(acc=a,pos_avg=p,native_frac=fr),coverage_only_test=dict(acc=a0,pos_avg=p0,native_frac=fr0))
    R[f"K{K}"]["_thr"]=thr_ok; print(R[f"K{K}"],flush=True)
    json.dump(R,open("resultados-013b.json","w"),indent=1)
# 3) tiempos CPU: por mensaje, bs=1, incluye tokenizar/gate/adaptador; 3 repeticiones; mediana y p90
random.seed(3); ts_=random.sample(cov_test,60)+random.sample(unc_test,40)  # aprox. 57/43
@torch.no_grad()
def one(m,route,K):
    t0=time.perf_counter(); g=0.0
    if route=="hyb":
        tg=time.perf_counter(); ad=ADS[K]; nat=False
        if covered(m):
            v=ad(torch.tensor([m])); r=readout(v)  # confianza requiere la lectura; se reutiliza si pasa
            c=min(float(F.softmax(x,1).max()) for x in r); nat=c>=R[f"K{K}"]["_thr"]
        g=time.perf_counter()-tg
        if nat: return time.perf_counter()-t0,g,g
    tf=time.perf_counter(); e=E(ids(" "+FORMATS[FMT](m)))[None]
    x=torch.cat([E(PRE)[None],e],1); out=lm(inputs_embeds=x,use_cache=True)
    ttft=None
    for f in range(4):
        c=copy.deepcopy(out.past_key_values); lm(inputs_embeds=E(Q[f])[None],past_key_values=c)
        if ttft is None: ttft=time.perf_counter()-t0
    return ttft,time.perf_counter()-t0,g
for m in ts_[:10]: one(m,"text",1); one(m,"hyb",1)
TIM={}
for name,route,K in (("texto",'text',1),("hibrido_K1",'hyb',1),("hibrido_K2",'hyb',2)):
    tt=[];tl=[];tg=[];c0=os.times()
    for rep in range(3):
        for m in ts_: a,b,g=one(m,route,K); tt.append(a*1e3);tl.append(b*1e3);tg.append(g*1e3)
    c1=os.times(); q=lambda v:sorted(v)[int(0.9*len(v))]
    TIM[name]=dict(ttft_med_ms=st.median(tt),ttft_p90_ms=q(tt),total_med_ms=st.median(tl),total_p90_ms=q(tl),total_mean_ms=st.mean(tl),gate_adapter_mean_ms=st.mean(tg),cpu_s_per_msg=((c1.user-c0.user)+(c1.system-c0.system))/len(tl))
    print(name,TIM[name],flush=True)
R["timing_cpu"]=TIM
import resource; R["peak_rss_MB"]=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024
json.dump(R,open("resultados-013b.json","w"),indent=1); print("DONE",flush=True)