# 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"]
best=dict(pre="1shot_natural_labeled",query="qa",fmt="natural_labeled",cal_acc=0.98)
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

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
0 and 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)

exec(open("defs_iter.py").read())
ADS={}; ad=Fact("fact4"); ad.load_state_dict(torch.load("adapter_013h.pt")); ADS["013h"]=ad.eval(); win="013h"
# 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)).float().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),flush=True)
print("DONE",flush=True)