# 013h+: maestro fuerte (>=97% en calibración) + per-field diagnóstico + escalado de pasos + timing del ganador.
import sys; exec(open('iterate.py').read().split("CFG=")[0].replace('FMT="natural_labeled"','FMT="natural_labeled"').split("_,TL=eval_text")[0])
FORMATS["rival_like"]=lambda m:f"The action is {A[m[0]]}. The target is {T[m[1]]}. The state is {S[m[2]]}. The window is {W[m[3]]}."
FORMATS["natural_labeled"]=FORMATS["natural_labeled"]
QS={"qa":[ids(f"\nQ: what is the {f}?\nA:") for f in FN],"cloze":[ids(f"\nThe {f} is") for f in FN],"cloze_colon":[ids(f"\n{f.capitalize()}:") for f in FN]}
random.seed(2); calS=random.sample(cov_cal,100)+random.sample(unc_cal,100)
sweep=[]
@torch.no_grad()
def perfield(fmt,ms):
    pf=[0]*4; ok=0
    for i in range(0,len(ms),50):
        b=ms[i:i+50]; L={}
        for m in b: L.setdefault(len(ids(" "+FORMATS[fmt](m))),[]).append(m)
        for sub in L.values():
            r=readout(E(torch.stack([ids(" "+FORMATS[fmt](m)) for m in sub]))); t=torch.tensor(sub); o=torch.ones(len(sub),dtype=torch.bool)
            for f in range(4): c=r[f].argmax(1)==t[:,f]; pf[f]+=int(c.sum()); o&=c
            ok+=int(o.sum())
    return ok/len(ms),[x/len(ms) for x in pf]
DEMO=(5,6,7,4)  # ejemplo fijo; se excluye de todas las evaluaciones? es cubierto/no: se reporta
PRES={"none":ids("Message:")}
for fmt in ["key_value","natural_labeled"]:
    d=FORMATS[fmt](DEMO); PRES["1shot_"+fmt]=ids("Message: "+d+"".join(f"\nQ: what is the {f}?\nA: {FL[i][DEMO[i]]}" for i,f in enumerate(FN))+"\n\nMessage:")
Q[:]=QS["qa"]
for pn,pv in PRES.items():
    PRE=pv; qn="qa"
    for fmt in ["key_value","natural_labeled","json"]:
        if pn.startswith("1shot") and not pn.endswith(fmt) : continue
        a,pf=perfield(fmt,calS); r=dict(pre=pn,query=qn,fmt=fmt,cal_acc=a,per_field=dict(zip(FN,pf)),avg_tokens=sum(len(ids(" "+FORMATS[fmt](m))) for m in calS)/len(calS))
        sweep.append(r); print(json.dumps(r),flush=True)
json.dump(sweep,open("sweep_013h.json","w"),indent=1)
good=[r for r in sweep if r["cal_acc"]>=0.97]
best=min(good,key=lambda r:r["avg_tokens"]) if good else max(sweep,key=lambda r:r["cal_acc"])
print("BEST",best,flush=True); PRE=PRES[best["pre"]]; Q[:]=QS[best["query"]]; FMT=best["fmt"]; pos=lambda m:len(ids(" "+FORMATS[FMT](m)))
# errores por valor en calibración completa
okT_cal,_=eval_text(FMT,cal_all); okT_dev,_=eval_text(FMT,dev)
errs={}
for m,o in okT_cal.items():
    if not o:
        for f in range(4): errs[f"{FN[f]}={FL[f][m[f]]}"]=errs.get(f"{FN[f]}={FL[f][m[f]]}",0)+1
json.dump(dict(best=best,cal_full_acc=sum(okT_cal.values())/len(cal_all),dev_acc=sum(okT_dev.values())/len(dev),values_in_failed_msgs=errs),open("baseline_013h.json","w"),indent=1)
_,TL=eval_text(FMT,cov_train,keep_logits=True)
exec(open("defs_iter.py").read())
RUNS=[("013h","fact4",200),("013i","add1",200),("013j","mlp4",160),("013k","mlp4",640),("013l","mlp1",640)]
ADS={}
for name,mode,steps in RUNS:
    torch.manual_seed(0); random.seed(0)
    ad=Fact(mode) if mode in("fact4","add1") else MLP(int(mode[-1]))
    t0=time.time(); ad=train(ad,steps,True); ADS[name]=ad
    r=evaluate(name,ad,cal_all,dev,okT_cal,okT_dev); r.update(mode=mode,steps=steps,train_s=round(time.time()-t0),params=sum(p.numel() for p in ad.parameters()),teacher=f'{FMT}/{best["pre"]}')
    open("iteraciones.jsonl","a").write(json.dumps(r)+"\n"); torch.save(ad.state_dict(),f"adapter_{name}.pt"); print(name,json.dumps(r),flush=True)
# fresca: sólo ganador por dev (h si prometedor)
cands=[json.loads(l) for l in open("iteraciones.jsonl") if l.startswith('{"iter": "013') and json.loads(l).get("teacher")]
win=max(cands,key=lambda r:(r["promising_native"],r["native_acc_sysholdout"],-r["K"]))["iter"]
okT_f,_=eval_text(FMT,fresh); rf=evaluate(win+"_FRESH",ADS[win],cal_all,fresh,okT_cal,okT_f)
open("iteraciones.jsonl","a").write(json.dumps(rf)+"\n"); print("FRESH",json.dumps(rf),flush=True)
# ABLACIONES del ganador en fresh: (a) LLM aleatorio misma arquitectura, (b) decodificación directa por vecino más cercano sin LLM
ad=ADS[win]; AB={}
with torch.no_grad():
    t=torch.tensor(fresh); v=ad(t)
    if v.shape[1]==4:
        okn=torch.ones(len(fresh),dtype=torch.bool)
        for f in range(4): okn&=(v[:,f]@E.weight[cand[f]].T).argmax(1)==t[:,f]
        AB["nn_decode_no_llm_acc"]=float(okn.float().mean())
    AB["holdout_pairs_in_train"]=sum(1 for m in cov_train if (m[0],m[1]) in set(H_AT) or (m[2],m[3]) in set(H_SW))
    AB["demo_msg_in_test"]=DEMO in set(fresh+dev)
from transformers import AutoConfig
real=lm; lm=AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(M)).eval(); E_real=E
okr,_=eval_ad(ad,fresh); AB["random_llm_native_acc"]=sum(okr.values())/len(fresh)
okrt,_=eval_text(FMT,fresh); AB["random_llm_text_acc"]=sum(okrt.values())/len(fresh)
lm=real; json.dump(AB,open("ablaciones_013.json","w"),indent=1); print("ABL",AB,flush=True)
# timing: texto vs nativo ganador (sin gate: todo nativo) vs híbrido gate (cubierto->nativo), bs=1, 3 reps, mediana/p90; incluye gate+adaptador
ad=ADS[win]; random.seed(3); ts_=random.sample([m for m in fresh if covered(m)],57)+random.sample([m for m in fresh if not covered(m)],43)
@torch.no_grad()
def one(m,route):
    t0=time.perf_counter(); g=0.0
    if route=="text" or (route=="hyb" and not covered(m)):
        if route=="hyb": g=time.perf_counter()-t0
        e=E(ids(" "+FORMATS[FMT](m)))[None]
    else:
        e=ad(torch.tensor([m])); g=time.perf_counter()-t0
    out=lm(inputs_embeds=torch.cat([E(PRE)[None],e],1),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*1e3,(time.perf_counter()-t0)*1e3,g*1e3
for m in ts_[:10]: one(m,"text"); one(m,"native")
TIM={}; q=lambda v:sorted(v)[int(0.9*len(v))]
data={k:[] for k in ("text","native","hyb")}
for rep in range(3):
    for m in ts_:
        for k in random.sample(list(data),3): data[k].append(one(m,k))  # orden intercalado aleatorio
for k,v in data.items():
    tt=[x[0] for x in v]; tl=[x[1] for x in v]; tg=[x[2] for x in v]
    TIM[k]=dict(ttft_med_ms=st.median(tt),ttft_p90_ms=q(tt),total_med_ms=st.median(tl),total_p90_ms=q(tl),gate_adapter_mean_ms=st.mean(tg),n=len(v))
import resource
json.dump(dict(winner=win,timing_cpu=TIM,positions=dict(text=sum(map(pos,ts_))/len(ts_),native=ad.K,hyb=sum((ad.K if covered(m) else pos(m)) for m in ts_)/len(ts_)),peak_rss_MB=resource.getrusage(resource.RUSAGE_SELF).ru_maxrss/1024),open("timing_013.json","w"),indent=1)
print("TIM",json.dumps(TIM),"DONE",flush=True)