#!/usr/bin/env python
"""
KazBERT benchmark — ModernBERT-style comparison of Kazakh encoders.

Fully AI-generated (Claude, via the Hermes ML research loop). Runs inference-only
on a single Kaggle T4. NO tokens or credentials are used or embedded: every model
and dataset below is PUBLIC, and results are pushed to the Hub from the local
machine afterwards (not from this script).

Metrics (all fair across tokenizers, with caveats documented in the card):
  1. Tokenizer fertility  — subword tokens per Kazakh word (lower = more efficient)
  2. MC accuracy (PLL)    — zero-shot multiple-choice via pseudo-log-likelihood
  3. MC accuracy (embed)  — zero-shot MC via cosine of mean-pooled embeddings
  4. Embedding separation — cos(question, correct) − mean cos(question, distractor)
Plus qualitative fill-mask examples (not scored — tokenizer-dependent).
"""
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

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",
}
N_MC   = 1000   # multiple-choice items
N_FERT = 3000   # sentences for fertility
FILLMASK_SENTS = [
    "Астана — Қазақстанның [MASK] қаласы.",
    "Мен қазақ [MASK] сөйлеймін.",
    "Абай Құнанбаев — ұлы қазақ [MASK].",
]

# ---------- data ----------
print("loading dastur-mc ...")
mc = load_dataset("kz-transformers/kazakh-dastur-mc", split="test")
mc = mc.select(range(min(N_MC, len(mc))))
# "Correct Answer" mixes Latin and Cyrillic homoglyph letters (A/А, B/В, C/С, D/Д).
LET = {"A": 0, "А": 0, "B": 1, "В": 1, "C": 2, "С": 2, "D": 3, "Д": 3}
def mc_item(r):
    opts = [str(r[f"Option {L}"]).strip() for L in ["A", "B", "C", "D"]]
    gold = LET.get(str(r["Correct Answer"]).strip()[:1].upper())
    if gold is None or any(not o for o in opts):
        return None
    return r["Question"].strip(), opts, gold
MC = [x for x in (mc_item(r) for r in mc) if x is not None]
print(f"  {len(MC)} MC items")

print("loading Kazakh sentences for fertility ...")
wiki = load_dataset("amandyk/kazakh_wiki_articles", split="train", streaming=True)
sents = []
for r in wiki:
    for line in str(r.get("text", "")).split("."):
        line = line.strip()
        if 6 <= len(line.split()) <= 40:
            sents.append(line)
        if len(sents) >= N_FERT: break
    if len(sents) >= N_FERT: break
print(f"  {len(sents)} sentences")

# ---------- per-model eval ----------
@torch.no_grad()
def embed(model, tok, texts, bs=64):
    vecs = []
    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)
        out = model(**b).last_hidden_state
        m = b["attention_mask"].unsqueeze(-1).float()
        v = (out * m).sum(1) / m.sum(1).clamp(min=1)
        vecs.append(F.normalize(v, dim=-1).cpu())
    return torch.cat(vecs)

@torch.no_grad()
def pll_score(mlm, tok, q, option, max_len=128):
    """Pseudo-log-likelihood of `option` conditioned on `q` (mask each option token)."""
    q_ids = tok.encode(q, add_special_tokens=False)
    o_ids = tok.encode(option, add_special_tokens=False)[:40]
    if not o_ids: return -1e9
    cls, sep, msk = tok.cls_token_id, tok.sep_token_id, tok.mask_token_id
    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)]
    if not pos: return -1e9
    rows, tgt = [], []
    for p in pos:
        r = base.copy(); r[p] = msk; rows.append(r); tgt.append(base[p])
    ml = max(len(r) for r in rows)
    pad = tok.pad_token_id or 0
    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)
    logits = mlm(input_ids=ids, attention_mask=att).logits
    lp = F.log_softmax(logits, dim=-1)
    s = sum(lp[i, pos[i], tgt[i]].item() for i in range(len(pos)))
    return s / len(pos)   # length-normalised

results = {}
for name, mid in MODELS.items():
    print(f"\n=== {name} ({mid}) ===")
    t0 = time.time()
    r = {"model_id": mid}
    try:
        tok = AutoTokenizer.from_pretrained(mid)
        # fertility
        tot_t = sum(len(tok.tokenize(s)) for s in sents)
        tot_w = sum(len(s.split()) for s in sents)
        r["fertility"] = round(tot_t / tot_w, 4)
        r["vocab_size"] = tok.vocab_size
        print(f"  fertility={r['fertility']}  vocab={tok.vocab_size}")

        # MLM head for PLL + fill-mask
        try:
            mlm = AutoModelForMaskedLM.from_pretrained(mid).to(dev).eval()
            correct = 0
            for q, opts, gold in MC:
                sc = [pll_score(mlm, tok, q, o) for o in opts]
                if int(np.argmax(sc)) == gold: correct += 1
            r["mc_acc_pll"] = round(correct / len(MC), 4)
            print(f"  MC(PLL)  acc={r['mc_acc_pll']}")
            # qualitative fill-mask
            fm = {}
            for s in FILLMASK_SENTS:
                try:
                    ss = s.replace("[MASK]", tok.mask_token)
                    b = tok(ss, return_tensors="pt").to(dev)
                    mi = (b["input_ids"][0] == tok.mask_token_id).nonzero()[0].item()
                    top = mlm(**b).logits[0, mi].topk(3).indices.tolist()
                    fm[s] = [tok.decode([t]).strip() for t in top]
                except Exception as e:
                    fm[s] = [f"<err>"]
            r["fill_mask"] = fm
            del mlm; gc.collect(); torch.cuda.empty_cache()
        except Exception as e:
            print("  MLM head unavailable:", str(e)[:100]); r["mc_acc_pll"] = None; r["fill_mask"] = {}

        # embeddings (base encoder)
        enc = AutoModel.from_pretrained(mid).to(dev).eval()
        correct = 0; seps = []
        qs = [m[0] for m in MC]
        q_emb = embed(enc, tok, qs)
        all_opts = [o for m in MC for o in m[1]]
        o_emb = embed(enc, tok, all_opts).view(len(MC), 4, -1)
        for i, (q, opts, gold) in enumerate(MC):
            cos = (q_emb[i].unsqueeze(0) * o_emb[i]).sum(-1)
            if int(cos.argmax()) == gold: correct += 1
            seps.append((cos[gold] - (cos.sum()-cos[gold])/3).item())
        r["mc_acc_embed"] = round(correct / len(MC), 4)
        r["embed_separation"] = round(float(np.mean(seps)), 4)
        print(f"  MC(embed) acc={r['mc_acc_embed']}  sep={r['embed_separation']}")
        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(results, open(f"{OUT}/results.json", "w"), ensure_ascii=False, indent=2)

# ---------- plots ----------
names = [n for n in MODELS if "fertility" in results.get(n, {})]
def bar(vals, title, ylabel, fname, better_low=False, hline=None, fmt="{:.3f}"):
    xs = [n for n in names if results[n].get(vals) is not None]
    ys = [results[n][vals] for n in xs]
    if not xs: return
    order = np.argsort(ys); order = order if better_low else order[::-1]
    xs = [xs[i] for i in order]; ys = [ys[i] for i in order]
    cols = ["#2a9d3f" if x == "KazBERT" else "#4C72B0" for x in xs]
    fig, ax = plt.subplots(figsize=(8, 4.2))
    b = ax.bar(xs, ys, color=cols)
    if hline is not None: ax.axhline(hline, ls="--", c="grey", alpha=.7, label=f"random {hline}")
    for rect, v in zip(b, ys): ax.text(rect.get_x()+rect.get_width()/2, v, fmt.format(v), ha="center", va="bottom", fontsize=9)
    ax.set_title(title, fontweight="bold"); ax.set_ylabel(ylabel)
    if hline is not None: ax.legend()
    plt.xticks(rotation=15); fig.tight_layout(); fig.savefig(f"{OUT}/{fname}", bbox_inches="tight"); plt.close(fig)

bar("fertility", "Tokenizer fertility on Kazakh (lower = better)", "tokens / word", "fertility.png", better_low=True)
bar("mc_acc_pll", "Zero-shot MC accuracy — PLL (dastur-mc)", "accuracy", "mc_pll.png", hline=0.25)
bar("mc_acc_embed", "Zero-shot MC accuracy — embeddings", "accuracy", "mc_embed.png", hline=0.25)
bar("embed_separation", "Embedding separation (correct − distractor cosine)", "Δ cosine", "embed_sep.png", fmt="{:.3f}")

# summary heatmap
metcols = [("fertility", True), ("mc_acc_pll", False), ("mc_acc_embed", False), ("embed_separation", False)]
M = []
for n in names:
    row = []
    for k, low in metcols:
        v = results[n].get(k)
        row.append(np.nan if v is None else v)
    M.append(row)
M = np.array(M, float)
norm = np.zeros_like(M)
for j, (k, low) in enumerate(metcols):
    col = M[:, j]; lo, hi = np.nanmin(col), np.nanmax(col)
    z = (col - lo) / (hi - lo + 1e-9)
    norm[:, j] = 1 - z if low else z
fig, ax = plt.subplots(figsize=(7.5, 0.7*len(names)+2))
im = ax.imshow(norm, cmap="RdYlGn", vmin=0, vmax=1, aspect="auto")
ax.set_xticks(range(len(metcols))); ax.set_xticklabels(["fertility↓", "MC-PLL↑", "MC-embed↑", "embed-sep↑"])
ax.set_yticks(range(len(names))); ax.set_yticklabels(names)
for i in range(len(names)):
    for j in range(len(metcols)):
        v = M[i, j]; ax.text(j, i, "—" if np.isnan(v) else f"{v:.3f}", ha="center", va="center", fontsize=9)
ax.set_title("KazBERT benchmark — summary (green = better)", fontweight="bold"); ax.grid(False)
fig.tight_layout(); fig.savefig(f"{OUT}/summary_heatmap.png", bbox_inches="tight"); plt.close(fig)

print("\nDONE. artifacts:", sorted(os.listdir(OUT)))
print(json.dumps(results, ensure_ascii=False, indent=2)[:1500])