#!/usr/bin/env python
"""
Did the Sinhala CPT work?  Compare the BASE model with the CPT model on identical test material.

Three steps (each model is scored in its own process so GPU memory is clean):
  python eval_cpt_success.py score   --model unsloth/gemma-4-E4B        --tag base
  python eval_cpt_success.py score   --model YOURNAME/gemma4-e4b-sinhala-cpt --tag cpt
  python eval_cpt_success.py compare --base base_scores.json --cpt cpt_scores.json

Tests (all automatic; none needs me to have written Sinhala answer keys):
  1. Held-out bits-per-character (BPC) on Sinhala sentences the model NEVER saw in training
     (the rows not in the 1M training sample), overall and per domain, with bootstrap CIs.
  2. Grammar / word-order probes: real sentence vs. a corrupted copy. Does the model give the real one
     the higher likelihood?  (adjacent-word swap, full word shuffle, one wrong consonant, word replaced
     by a random corpus word).  Chance = 50%.
  3. English forgetting check: BPC on a few English sentences (CPT should not get much worse).
  4. Qualitative: greedy continuations of Sinhala prompts, side by side (for a native speaker to read).
Use the SAME --max_sentences / --seed / --val_sentences as the training run so "unseen" is really unseen.
"""
import argparse
import json
import math
import os
import random
import sys

import numpy as np

ENGLISH = [
    "The committee met on Tuesday to discuss the budget for the coming year.",
    "Water boils at one hundred degrees Celsius at sea level.",
    "She opened the window and listened to the rain falling on the roof.",
    "Researchers have developed a new method for detecting early signs of disease.",
    "The train from the capital arrives at the station every morning at seven.",
    "He has been studying computer science at the university for three years.",
    "Many farmers in the region depend on the monsoon for their harvest.",
    "The museum is open to visitors from Monday to Saturday.",
    "A balanced diet includes vegetables, fruit, protein and whole grains.",
    "The company announced that it would open two new offices next spring.",
    "Please send me the report as soon as you have finished reviewing it.",
    "The river flooded after several days of heavy rain in the hills.",
    "Learning a second language can improve memory and problem solving skills.",
    "The doctor advised him to rest and drink plenty of fluids.",
    "Students must submit their assignments before the end of the month.",
    "The old bridge was repaired last year and now carries both cars and pedestrians.",
    "Prices of essential goods have risen steadily over the past six months.",
    "The film tells the story of a young girl who moves to a small coastal town.",
    "Electric cars are becoming more common as charging stations spread across the country.",
    "After the meeting, everyone agreed that the plan needed more time.",
]

PROMPTS = [
    "ශ්‍රී ලංකාව", "ඊයේ පාසලේදී", "ජනාධිපතිවරයා පැවසුවේ", "ගොවීන්ට පොහොර",
    "බුදුන් වහන්සේ", "කොළඹ නගරයේ", "දරුවන්ගේ අධ්‍යාපනය", "වෛද්‍යවරුන්ට අනුව",
]

KINDS = ["swap", "shuffle", "charsub", "wordsub"]
CONSONANTS = [chr(c) for c in range(0x0D9A, 0x0DC7)]


# ------------------------------------------------------------------ corruptions
def c_swap(text, rng, vocab=None):
    ws = text.split(" ")
    if len(ws) < 5:
        return None
    for _ in range(10):
        i = rng.randrange(len(ws) - 1)
        if ws[i] != ws[i + 1]:
            w = ws[:]
            w[i], w[i + 1] = w[i + 1], w[i]
            return " ".join(w)
    return None


def c_shuffle(text, rng, vocab=None):
    ws = text.split(" ")
    if len(ws) < 5:
        return None
    for _ in range(10):
        w = ws[:]
        rng.shuffle(w)
        if w != ws:
            return " ".join(w)
    return None


def c_charsub(text, rng, vocab=None):
    ws = text.split(" ")
    cand = [i for i, w in enumerate(ws) if len(w) >= 4 and any(ch in CONSONANTS for ch in w)]
    if not cand:
        return None
    i = rng.choice(cand)
    w = ws[i]
    pos = [k for k, ch in enumerate(w) if ch in CONSONANTS]
    k = rng.choice(pos)
    new = rng.choice([c for c in CONSONANTS if c != w[k]])
    ws[i] = w[:k] + new + w[k + 1:]
    return " ".join(ws)


def c_wordsub(text, rng, vocab):
    ws = text.split(" ")
    if len(ws) < 5 or not vocab:
        return None
    i = rng.randrange(len(ws))
    same_len = [v for v in vocab if abs(len(v) - len(ws[i])) <= 1 and v != ws[i]]
    if not same_len:
        return None
    ws[i] = rng.choice(same_len)
    return " ".join(ws)


CORRUPT = {"swap": c_swap, "shuffle": c_shuffle, "charsub": c_charsub, "wordsub": c_wordsub}


# ------------------------------------------------------------------ data
def build_data(args):
    import pandas as pd
    from huggingface_hub import hf_hub_download

    path = hf_hub_download(args.repo_id, args.data_file, repo_type="dataset")
    print(f"[data] reading {path}")
    df = pd.read_csv(path, usecols=["sentence", "predicted_domain"])
    # identical cleaning to cpt_gemma4_sinhala.py
    df = df.dropna(subset=["sentence"])
    df["sentence"] = df["sentence"].astype(str).str.strip()
    df = df[df["sentence"].str.len() >= 15]
    df = df.drop_duplicates(subset="sentence")
    total = min(args.max_sentences + args.val_sentences, len(df))
    sampled = df.sample(n=total, random_state=args.seed)       # exactly the training sample
    unseen = df.drop(sampled.index)
    unseen = unseen[unseen["sentence"].str.len() <= 300]
    unseen = unseen.sample(n=min(args.n, len(unseen)), random_state=args.seed + 1)
    print(f"[data] training sample {len(sampled):,}; unseen pool {len(df) - len(sampled):,}; using {len(unseen):,}")

    heldout = [{"text": s, "domain": d} for s, d in zip(unseen["sentence"], unseen["predicted_domain"])]

    rng = random.Random(args.seed)
    vocab = sorted({w for s in unseen["sentence"].head(3000) for w in s.split(" ") if 2 <= len(w) <= 14})
    pairs = {k: [] for k in KINDS}
    order = list(range(len(heldout)))
    for k in KINDS:
        rng.shuffle(order)
        for i in order:
            if len(pairs[k]) >= args.n_pairs:
                break
            nw = len(heldout[i]["text"].split(" "))
            if not (5 <= nw <= 40):
                continue
            bad = CORRUPT[k](heldout[i]["text"], rng, vocab)
            if bad and bad != heldout[i]["text"]:
                pairs[k].append({"idx": i, "bad": bad})
        print(f"[data] pairs[{k}] = {len(pairs[k])}")
    data = {"heldout": heldout, "pairs": pairs, "english": ENGLISH, "prompts": PROMPTS,
            "meta": {"max_sentences": args.max_sentences, "seed": args.seed}}
    json.dump(data, open(args.data, "w"), ensure_ascii=False)
    print(f"[data] saved {args.data}")
    return data


# ------------------------------------------------------------------ scoring
def score_texts(model, tok, texts, bs=16, max_len=256):
    """Sum of negative log-likelihood (nats) per text, first real token conditioned on BOS."""
    import torch
    import torch.nn.functional as F
    tok.padding_side = "right"
    out = [0.0] * len(texts)
    order = sorted(range(len(texts)), key=lambda i: len(texts[i]))
    for s in range(0, len(order), bs):
        idx = order[s:s + bs]
        enc = tok([texts[i] for i in idx], return_tensors="pt", padding=True,
                  truncation=True, max_length=max_len).to("cuda")
        with torch.no_grad():
            logits = model(input_ids=enc["input_ids"], attention_mask=enc["attention_mask"]).logits
        lg = logits[:, :-1].float()
        tg = enc["input_ids"][:, 1:]
        m = enc["attention_mask"][:, 1:].bool()
        nll = F.cross_entropy(lg.reshape(-1, lg.size(-1)), tg.reshape(-1), reduction="none").view(tg.shape)
        nll = (nll * m).sum(1).tolist()
        for j, i in enumerate(idx):
            out[i] = nll[j]
        if (s // bs) % 25 == 0:
            print(f"   scored {min(s + bs, len(order))}/{len(order)}", flush=True)
    return out


def run_score(args):
    data = json.load(open(args.data)) if os.path.exists(args.data) else build_data(args)
    from unsloth import FastModel  # must be imported before transformers
    print(f"[model] loading {args.model}")
    model, tokenizer = FastModel.from_pretrained(model_name=args.model, max_seq_length=512, dtype=None,
                                                 load_in_4bit=False, full_finetuning=False)
    tok = getattr(tokenizer, "tokenizer", tokenizer)
    model.eval()

    res = {"tag": args.tag, "model": args.model}
    print("[score] held-out Sinhala")
    res["heldout_nll"] = score_texts(model, tok, [x["text"] for x in data["heldout"]])
    print("[score] English")
    res["english_nll"] = score_texts(model, tok, data["english"])
    res["pairs"] = {}
    for k in KINDS:
        print(f"[score] pairs: {k}")
        res["pairs"][k] = score_texts(model, tok, [p["bad"] for p in data["pairs"][k]])

    print("[generate] prompts")
    res["gens"] = {}
    for p in data["prompts"]:
        enc = tok(p, return_tensors="pt").to("cuda")
        try:
            out = model.generate(**enc, max_new_tokens=60, do_sample=False, repetition_penalty=1.15)
            res["gens"][p] = tok.decode(out[0], skip_special_tokens=True)
        except Exception as e:  # noqa: BLE001
            res["gens"][p] = f"<generation failed: {type(e).__name__}: {e}>"
    json.dump(res, open(args.out or f"{args.tag}_scores.json", "w"), ensure_ascii=False)
    print(f"[done] wrote {args.out or args.tag + '_scores.json'}")


# ------------------------------------------------------------------ compare
def boot_ci(fn, n, B=2000, seed=0):
    rng = np.random.default_rng(seed)
    vals = [fn(rng.integers(0, n, n)) for _ in range(B)]
    return np.percentile(vals, [2.5, 97.5])


def run_compare(args):
    data = json.load(open(args.data))
    A, C = json.load(open(args.base)), json.load(open(args.cpt))
    lines = []

    def out(s=""):
        print(s)
        lines.append(s)

    chars = np.array([len(x["text"]) for x in data["heldout"]], float)
    a, c = np.array(A["heldout_nll"]), np.array(C["heldout_nll"])
    bpc = lambda nll, ch: nll.sum() / math.log(2) / ch.sum()          # noqa: E731
    rel = lambda idx: (a[idx].sum() - c[idx].sum()) / a[idx].sum()    # noqa: E731
    lo, hi = boot_ci(rel, len(a))
    out(f"# CPT success report: {C['model']}  vs  {A['model']}\n")
    out("## 1. Held-out Sinhala bits per character (lower is better)")
    out(f"- sentences: {len(a)} (never seen in training)")
    out(f"- base: {bpc(a, chars):.4f}   cpt: {bpc(c, chars):.4f}")
    out(f"- relative improvement: {rel(np.arange(len(a))) * 100:.1f}%  (95% CI {lo * 100:.1f}% to {hi * 100:.1f}%)\n")
    out("| domain | n | base BPC | cpt BPC | improvement |")
    out("|---|---|---|---|---|")
    doms = np.array([x["domain"] for x in data["heldout"]])
    for d in sorted(set(doms)):
        m = doms == d
        out(f"| {d} | {m.sum()} | {bpc(a[m], chars[m]):.4f} | {bpc(c[m], chars[m]):.4f} | "
            f"{(a[m].sum() - c[m].sum()) / a[m].sum() * 100:.1f}% |")

    ech = np.array([len(t) for t in data["english"]], float)
    ea, ec = np.array(A["english_nll"]), np.array(C["english_nll"])
    erel = (ec.sum() - ea.sum()) / ea.sum()
    elo, ehi = boot_ci(lambda i: (ec[i].sum() - ea[i].sum()) / ea[i].sum(), len(ea))
    out("\n## 2. English forgetting (BPC; lower is better)")
    out(f"- base: {bpc(ea, ech):.4f}   cpt: {bpc(ec, ech):.4f}   change: {erel * 100:+.1f}% "
        f"(95% CI {elo * 100:+.1f}% to {ehi * 100:+.1f}%)  [only {len(ea)} sentences: noisy]")

    out("\n## 3. Grammar / word-choice probes: % where the REAL sentence beats the corrupted one (chance 50%)")
    out("| probe | n | base | cpt | change (pp) | 95% CI |")
    out("|---|---|---|---|---|---|")
    diffs = []
    descr = {"swap": "two adjacent words swapped", "shuffle": "all words shuffled",
             "charsub": "one consonant replaced", "wordsub": "one word replaced by a random word"}
    for k in KINDS:
        items = data["pairs"][k]
        if not items:
            continue
        idx = np.array([p["idx"] for p in items])
        ba, bc = np.array(A["pairs"][k]), np.array(C["pairs"][k])
        wa, wc = (a[idx] < ba).astype(float), (c[idx] < bc).astype(float)
        d = lambda i: (wc[i].mean() - wa[i].mean()) * 100   # noqa: E731
        plo, phi = boot_ci(d, len(idx))
        diffs.append(d(np.arange(len(idx))))
        out(f"| {descr[k]} | {len(idx)} | {wa.mean() * 100:.1f}% | {wc.mean() * 100:.1f}% | "
            f"{d(np.arange(len(idx))):+.1f} | {plo:+.1f} to {phi:+.1f} |")

    out("\n## 4. Verdict (my rules of thumb, not a standard)")
    ok1 = lo > 0 and rel(np.arange(len(a))) >= 0.03
    ok2 = erel <= 0.10
    ok3 = bool(diffs) and float(np.mean(diffs)) >= 0
    out(f"- [{'PASS' if ok1 else 'FAIL'}] Sinhala BPC improved by >= 3% with CI above 0")
    out(f"- [{'PASS' if ok2 else 'FAIL'}] English BPC got worse by <= 10%")
    out(f"- [{'PASS' if ok3 else 'FAIL'}] grammar probes did not get worse on average")
    out("- Perplexity gains show the model fits Sinhala text better. They do not by themselves prove better "
        "exam-style reasoning: confirm with the shared-task dev set accuracy.")

    out("\n## 5. Continuations (greedy, 60 tokens). Have a native speaker judge fluency and sense.")
    for p in data["prompts"]:
        out(f"\n**Prompt:** {p}\n\n- base: {A['gens'].get(p, '').replace(chr(10), ' ')}\n- cpt:  {C['gens'].get(p, '').replace(chr(10), ' ')}")
    open(args.report, "w").write("\n".join(lines))
    print(f"\n[done] report written to {args.report}")


# ------------------------------------------------------------------ cli
def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    for name in ("prepare", "score"):
        p = sub.add_parser(name)
        p.add_argument("--repo_id", default="Minuri/diverse_sinhala_dataset")
        p.add_argument("--data_file", default="high_confidence_7_7M.csv")
        p.add_argument("--max_sentences", type=int, default=1_000_000, help="same as training")
        p.add_argument("--val_sentences", type=int, default=10_000, help="same as training")
        p.add_argument("--seed", type=int, default=3407, help="same as training")
        p.add_argument("--n", type=int, default=4000, help="unseen Sinhala sentences to score")
        p.add_argument("--n_pairs", type=int, default=1000, help="pairs per probe type")
        p.add_argument("--data", default="eval_data.json")
        if name == "score":
            p.add_argument("--model", required=True, help="HF repo or local folder (base, adapter or merged)")
            p.add_argument("--tag", required=True, help="base or cpt")
            p.add_argument("--out", default="")
    p = sub.add_parser("compare")
    p.add_argument("--base", default="base_scores.json")
    p.add_argument("--cpt", default="cpt_scores.json")
    p.add_argument("--data", default="eval_data.json")
    p.add_argument("--report", default="cpt_report.md")
    args = ap.parse_args()
    {"prepare": build_data, "score": run_score, "compare": run_compare}[args.cmd](args)


if __name__ == "__main__":
    main()