#!/usr/bin/env python
"""
KazBERT benchmark — Part 2: full Kazakh UNT (ЕНТ) exam, zero-shot, no training.

Fully AI-generated (Claude, Hermes ML research loop). Inference-only on a Kaggle T4.
No tokens/credentials: all models + datasets are PUBLIC; results pushed from local.

14,850 real Unified National Testing questions across 7 subjects. Each encoder answers
zero-shot by (a) pseudo-log-likelihood scoring of each option under its MLM head, and
(b) cosine of mean-pooled embeddings. Options vary from 4 to 8 per question.
"""
import os, json, time, gc, warnings
import numpy as np
import torch, torch.nn.functional as F
warnings.filterwarnings("ignore")
import matplotlib; matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = "/kaggle/working"; os.makedirs(OUT, exist_ok=True)
plt.rcParams.update({"figure.dpi": 130, "font.size": 11, "axes.grid": True, "grid.alpha": 0.25,
                     "axes.spines.top": False, "axes.spines.right": False})
dev = "cuda" if torch.cuda.is_available() else "cpu"
print("device:", dev, torch.cuda.get_device_name(0) if dev == "cuda" else "")
import subprocess, sys
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "-U", "transformers", "datasets"], check=True)
from transformers import AutoTokenizer, AutoModelForMaskedLM, AutoModel
from datasets import load_dataset, get_dataset_config_names, get_dataset_split_names

MODELS = {
    "KazBERT":         "Eraly-ml/KazBERT",
    "mBERT":           "google-bert/bert-base-multilingual-cased",
    "XLM-R base":      "FacebookAI/xlm-roberta-base",
    "kaz-roberta":     "kz-transformers/kaz-roberta-conversational",
    "KazakhBERTmulti": "amandyk/KazakhBERTmulti",
}
LET = {L: i for i, L in enumerate("ABCDEFGH")}
LET.update({"А":0,"В":1,"С":2,"Д":3,"Е":4,"Н":7})  # Cyrillic homoglyphs
OPTCOLS = list("ABCDEFGH")

# ---------- load full UNT ----------
DS = "kz-transformers/kazakh-unified-national-testing-mc"
subjects = get_dataset_split_names(DS)
print("subjects:", subjects)
items = []   # (subject, question, [options], gold_idx)
for sub in subjects:
    d = load_dataset(DS, split=sub)
    for r in d:
        opts = [str(r[c]).strip() for c in OPTCOLS if r.get(c) not in (None, "None", "")]
        g = LET.get(str(r["correct_answer"]).strip()[:1].upper())
        if g is None or g >= len(opts) or len(opts) < 2 or not r.get("question"): continue
        items.append((sub, str(r["question"]).strip(), opts, g))
print(f"total questions: {len(items)}")
rand_baseline = float(np.mean([1/len(o) for _,_,o,_ in items]))
print(f"random baseline (avg 1/#opts): {rand_baseline:.3f}")

# ---------- scoring ----------
@torch.no_grad()
def pll_all(mlm, tok, q, opts, max_len=160):
    """One batched forward for the whole question: mask each option token, return per-option normalised PLL."""
    cls, sep, msk, pad = tok.cls_token_id, tok.sep_token_id, tok.mask_token_id, (tok.pad_token_id or 0)
    q_ids = tok.encode(q, add_special_tokens=False)[:100]
    rows, meta = [], []   # meta: (opt_idx, target_token)
    for oi, o in enumerate(opts):
        o_ids = tok.encode(o, add_special_tokens=False)[:40]
        if not o_ids: meta.append(None); continue
        base = ([cls] if cls is not None else []) + q_ids + ([sep] if sep is not None else []) + o_ids + ([sep] if sep is not None else [])
        base = base[:max_len]
        start = (1 if cls is not None else 0) + len(q_ids) + (1 if sep is not None else 0)
        pos = [p for p in range(start, start+len(o_ids)) if p < len(base)]
        for p in pos:
            r = base.copy(); r[p] = msk
            rows.append((r, p)); meta.append((oi, base[p]))
    if not rows: return [-1e9]*len(opts)
    ml = max(len(r) for r,_ in rows)
    ids = torch.tensor([r+[pad]*(ml-len(r)) for r,_ in rows], device=dev)
    att = torch.tensor([[1]*len(r)+[0]*(ml-len(r)) for r,_ in rows], device=dev)
    lp = F.log_softmax(mlm(input_ids=ids, attention_mask=att).logits, dim=-1)
    scores = [[] for _ in opts]
    for k, ((r, p), m) in enumerate(zip(rows, [x for x in meta if x is not None])):
        oi, tgt = m
        scores[oi].append(lp[k, p, tgt].item())
    return [np.mean(s) if s else -1e9 for s in scores]

@torch.no_grad()
def emb(model, tok, texts, bs=128):
    vs = []
    for i in range(0, len(texts), bs):
        b = tok(texts[i:i+bs], return_tensors="pt", padding=True, truncation=True, max_length=64).to(dev)
        o = model(**b).last_hidden_state
        m = b["attention_mask"].unsqueeze(-1).float()
        vs.append(F.normalize((o*m).sum(1)/m.sum(1).clamp(min=1), dim=-1).cpu())
    return torch.cat(vs)

# ---------- per-model ----------
results = {}
subj_list = sorted(set(s for s,_,_,_ in items))
for name, mid in MODELS.items():
    print(f"\n=== {name} ({mid}) ===", flush=True)
    t0 = time.time(); r = {"model_id": mid}
    try:
        tok = AutoTokenizer.from_pretrained(mid)
        # ---- PLL ----
        try:
            mlm = AutoModelForMaskedLM.from_pretrained(mid).to(dev).eval()
            hit = {s: [0,0] for s in subj_list}; tot_c = 0
            for j,(sub,q,opts,g) in enumerate(items):
                sc = pll_all(mlm, tok, q, opts)
                ok = int(np.argmax(sc)) == g
                hit[sub][0] += ok; hit[sub][1] += 1; tot_c += ok
                if (j+1) % 2000 == 0: print(f"    PLL {j+1}/{len(items)}", flush=True)
            r["pll_overall"] = round(tot_c/len(items), 4)
            r["pll_by_subject"] = {s: round(hit[s][0]/hit[s][1], 4) for s in subj_list if hit[s][1]}
            print(f"  PLL overall={r['pll_overall']}")
            del mlm; gc.collect(); torch.cuda.empty_cache()
        except Exception as e:
            print("  MLM head unavailable:", str(e)[:120]); r["pll_overall"]=None; r["pll_by_subject"]={}
        # ---- embeddings ----
        enc = AutoModel.from_pretrained(mid).to(dev).eval()
        qs = [q for _,q,_,_ in items]; q_emb = emb(enc, tok, qs)
        flat, spans = [], []
        for _,_,opts,_ in items:
            spans.append((len(flat), len(flat)+len(opts))); flat.extend(opts)
        o_emb = emb(enc, tok, flat)
        hit = {s:[0,0] for s in subj_list}; tot_c=0
        for i,(sub,q,opts,g) in enumerate(items):
            a,b = spans[i]; cos = (q_emb[i].unsqueeze(0)*o_emb[a:b]).sum(-1)
            ok = int(cos.argmax())==g; hit[sub][0]+=ok; hit[sub][1]+=1; tot_c+=ok
        r["emb_overall"] = round(tot_c/len(items), 4)
        r["emb_by_subject"] = {s: round(hit[s][0]/hit[s][1],4) for s in subj_list if hit[s][1]}
        r["vocab_size"] = tok.vocab_size
        print(f"  EMB overall={r['emb_overall']}")
        del enc; gc.collect(); torch.cuda.empty_cache()
    except Exception as e:
        print("  FAILED:", str(e)[:200]); r["error"]=str(e)[:200]
    r["seconds"] = round(time.time()-t0,1); results[name]=r
    json.dump({"random_baseline":round(rand_baseline,4),"n_questions":len(items),"subjects":subj_list,"models":results},
              open(f"{OUT}/unt_results.json","w"), ensure_ascii=False, indent=2)

# ---------- plots ----------
names = [n for n in MODELS if results.get(n,{}).get("pll_overall") is not None]
# overall bar
fig, ax = plt.subplots(figsize=(9,4.4))
x = np.arange(len(names)); w=0.38
pll = [results[n]["pll_overall"]*100 for n in names]
emb_ = [results[n]["emb_overall"]*100 for n in names]
b1=ax.bar(x-w/2, pll, w, label="PLL", color=["#2a9d3f" if n=="KazBERT" else "#4C72B0" for n in names])
b2=ax.bar(x+w/2, emb_, w, label="embeddings", color=["#7fce8f" if n=="KazBERT" else "#a7c0e0" for n in names])
ax.axhline(rand_baseline*100, ls="--", c="grey", label=f"random {rand_baseline*100:.0f}%")
for bars in (b1,b2):
    for rr in bars: ax.text(rr.get_x()+rr.get_width()/2, rr.get_height(), f"{rr.get_height():.1f}", ha="center", va="bottom", fontsize=8)
ax.set_xticks(x); ax.set_xticklabels(names, rotation=15); ax.set_ylabel("accuracy %")
ax.set_title(f"Full ЕНТ zero-shot accuracy — {len(items):,} questions", fontweight="bold"); ax.legend()
fig.tight_layout(); fig.savefig(f"{OUT}/unt_overall.png", bbox_inches="tight"); plt.close(fig)

# per-subject heatmap (PLL)
subs = subj_list
M = np.array([[results[n]["pll_by_subject"].get(s, np.nan) for s in subs] for n in names], float)
fig, ax = plt.subplots(figsize=(1.1*len(subs)+3, 0.7*len(names)+2))
im = ax.imshow(M, cmap="RdYlGn", vmin=rand_baseline, vmax=max(0.6, np.nanmax(M)), aspect="auto")
ax.set_xticks(range(len(subs))); ax.set_xticklabels([s.replace("_"," ") for s in subs], rotation=35, ha="right")
ax.set_yticks(range(len(names))); ax.set_yticklabels(names)
for i in range(len(names)):
    for j in range(len(subs)):
        v=M[i,j]; ax.text(j,i,"—" if np.isnan(v) else f"{v*100:.0f}", ha="center", va="center", fontsize=8)
ax.set_title("Per-subject accuracy (PLL, %) — green > random", fontweight="bold"); ax.grid(False)
fig.colorbar(im, fraction=0.03, pad=0.02); fig.tight_layout()
fig.savefig(f"{OUT}/unt_by_subject.png", bbox_inches="tight"); plt.close(fig)

print("\nDONE. artifacts:", sorted(os.listdir(OUT)))
print(json.dumps({n:{"pll":results[n].get("pll_overall"),"emb":results[n].get("emb_overall")} for n in results}, indent=2))