"""Erk-32B ara degerlendirme — HELD-OUT 750 (probe'un 150 sorusu HARIC).
Mevcut LoRA checkpoint'i taban modele takip gercek sayiyi olcer."""
import os, json, math, torch, torch.nn as nn
os.environ["HF_DATASETS_OFFLINE"]="1"
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

MID="/ari/users/ytanis/qwen3-32b"
# Olculecek checkpoint dizini disaridan verilebilir (CPT / SFT ayni betikle
# ayni protokolde olculsun diye).
OUT=os.environ.get("CK_DIR", "/ari/users/ytanis/erk32b-cpt")
R=256; ALPHA=512
tok=AutoTokenizer.from_pretrained(MID)
model=AutoModelForCausalLM.from_pretrained(MID, torch_dtype=torch.bfloat16, device_map="auto").eval()

class LoRA(nn.Module):
    def __init__(self, base, r=R, alpha=ALPHA):
        super().__init__(); self.base=base
        d=base.weight.device
        self.A=nn.Parameter(torch.zeros(r, base.in_features, device=d, dtype=torch.bfloat16))
        self.B=nn.Parameter(torch.zeros(base.out_features, r, device=d, dtype=torch.bfloat16))
        self.s=alpha/r
    def forward(self,x):
        return self.base(x)+((x@self.A.t())@self.B.t())*self.s
for layer in model.model.layers:
    for mod,names in ((layer.self_attn,["q_proj","k_proj","v_proj","o_proj"]),
                      (layer.mlp,["gate_proj","up_proj","down_proj"])):
        for nm in names: setattr(mod,nm,LoRA(getattr(mod,nm)))

st=json.load(open(f"{OUT}/state.json"))
lsd=torch.load(f"{OUT}/lora_ck.pt", map_location="cpu", weights_only=True)
named=dict(model.named_parameters()); n=0
for k,v in lsd.items():
    if k in named:
        with torch.no_grad(): named[k].copy_(v.to(named[k].device))
        n+=1
print(f"LoRA yuklendi: {n} tensor | adim {st['step']} ({st['tokens']/1e9:.2f}B token)", flush=True)

ds=load_dataset("AYueksel/TurkishMMLU","All",split="test")
letters="ABCDE"; lid=[tok(" "+c, add_special_tokens=False).input_ids[-1] for c in letters]
probe=set(list(range(1,len(ds),6))[:150])
held=[i for i in range(min(900,len(ds))) if i not in probe]
# --- Olcum protokolu (adim 23500'de sabitlendi, olculdu):
#   1) probe'un 150 sorusu haric  -> secim sizintisi yok
#   2) kirli sorular haric        -> kulliyatta 13-kelime birebir gecenler
#      (held-out 750'nin 98'i; bkz. kirlilik_CPT.json)
#   3) sik sirasi dondurulur      -> harf/konum onyargisi elenir
# Bu ucu birden uygulanmadan olculen fark ~2 puan sisiyor (6.27 -> 4.14).
KIR=set()
try:
    KIR=set(json.load(open("/ari/users/ytanis/kirlilik_CPT.json"))["kirli_13"])
except Exception as e:
    print("UYARI: kirlilik listesi okunamadi (%s) — temiz alt kume raporlanmayacak" % e, flush=True)
temiz=[i for i in held if i not in KIR]

def run(kaydir):
    v=[]
    with torch.no_grad():
        for i in held:
            ex=ds[i]
            try:
                a=ex["answer"]; ans=int(a) if isinstance(a,int) else letters.index(str(a).strip())
            except: v.append(None); continue
            ch=list(ex["choices"])[:5]; k=len(ch)
            if kaydir:
                ch=ch[-kaydir:]+ch[:-kaydir]; ans=(ans+kaydir)%k
            s="Asagidaki coktan secmeli soruyu cevapla.\n\nSoru: "+str(ex["question"])+"\n"
            for j,c in enumerate(ch): s+=f"{letters[j]}) {c}\n"
            ids=tok(s+"Cevap:", return_tensors="pt", truncation=True, max_length=1024).input_ids.to(model.device)
            lg=model(ids).logits[0,-1]
            v.append(int(int(torch.stack([lg[t] for t in lid[:k]]).argmax())==ans))
    return v

def oran(v, idx):
    z=[v[held.index(i)] for i in idx if v[held.index(i)] is not None]
    return 100*sum(z)/max(len(z),1), len(z)

sonuc={"step":st["step"],"tokens":st["tokens"]}
# Asama etiketi checkpoint dizininden turetilir; sabit "CPT" yazmak SFT
# olcumlerinde yaniltici log uretiyordu.
ASAMA = "SFT" if OUT.rstrip("/").endswith("sft") else "CPT"
print(f"\n=== ERK-32B ({ASAMA}, adim {st['step']}, {st['tokens']/1e9:.2f}B token) ===", flush=True)
for ad,kay in (("orijinal",0),("dondurulmus",1)):
    v=run(kay); sonuc["vec_"+ad]=v
    a750,n750=oran(v,held); atem,ntem=oran(v,temiz)
    sonuc[ad]={"held750":a750,"temiz":atem,"n_temiz":ntem}
    print(f"  {ad:12s} held-out {n750}: %{a750:.2f} | TEMIZ {ntem}: %{atem:.2f}", flush=True)
print(f"\nRAPORLANACAK SAYI (temiz + dondurulmus): %{sonuc['dondurulmus']['temiz']:.2f}", flush=True)
print(f"  stok Qwen3-32B ayni kosulda: %67.64 (adim 23500 olcumunden)", flush=True)
json.dump(sonuc, open(f"{OUT}/heldout_step{st['step']}.json","w"))
print("BITTI", flush=True)