#!/usr/bin/env python
"""
Toolkit for building a Sinhala multiple-choice fine-tuning dataset (SinhalaMMLU-style).

Record schema (one JSON object per line, UTF-8). Mirrors the test file; `answer` is 1-based:
  {"id": "ol2019_sci_p1_q07", "subject": "science", "category": "stem",
   "question": "...", "choices": ["...", "...", "...", "..."], "answer": 2,
   "metadata": {"source": "<paper id or URL>", "exam_type": "OL", "year": 2019, "grade": 11,
                "difficulty": "medium", "extraction": "text_pdf|ocr|synthetic|translated"}}

Commands (run in this order):
  python mcq_dataset.py example  example.jsonl            # writes 2 dummy records to see the schema
  python mcq_dataset.py validate raw.jsonl                # report problems, changes nothing
  python mcq_dataset.py clean    raw.jsonl clean.jsonl --rejects rejects.jsonl
  python mcq_dataset.py dedupe   clean.jsonl dedup.jsonl --blocklist public_test.json   # also decontaminates
  python mcq_dataset.py split    dedup.jsonl --train train.jsonl --val val.jsonl --val_frac 0.1
  python mcq_dataset.py stats    train.jsonl
  python mcq_dataset.py build-sft train.jsonl sft_train.jsonl --shuffle_copies 1 --strip_zwj_prob 0.5
  python mcq_dataset.py build-sft val.jsonl   sft_val.jsonl   --shuffle_copies 0 --strip_zwj_prob 0 --templates si_instruction

The --blocklist file is ONLY used to find and drop overlapping questions. It is never copied into the output.
"""
import argparse
import collections
import hashlib
import json
import random
import re
import sys
import unicodedata

ZWJ = "‍"

# options that refer to other options: never reorder these
POSITIONAL = [r"ඉහත", r"සියල්ල", r"කිසිවක්", r"[1-5]\s*(සහ|හා|and|&)\s*[1-5]", r"\ball of the above\b", r"\bnone of\b"]
FIGURE = [r"රූපය", r"රූප සටහන", r"රූපයේ", r"රූපවල", r"ප්රස්තාර", r"සිතියම", r"\bfigure\b", r"\bdiagram\b", r"\bgraph\b"]

SI_INSTR = "පහත බහුවරණ ප්‍රශ්නයට නිවැරදි පිළිතුරේ අංකය ලියන්න."
EN_INSTR = "Answer the following multiple-choice question. Reply with only the number of the correct option."


# ------------------------------------------------------------------ io / helpers
def load(path):
    with open(path, encoding="utf-8") as f:
        txt = f.read().strip()
    if not txt:
        return []
    if txt[0] == "[":
        return json.loads(txt)
    return [json.loads(l) for l in txt.splitlines() if l.strip()]


def dump(path, rows):
    with open(path, "w", encoding="utf-8") as f:
        for r in rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def getpath(d, path, default=None):
    for k in path.split("."):
        if not isinstance(d, dict) or k not in d:
            return default
        d = d[k]
    return d


def strip_zwj(s):
    return s.replace(ZWJ, "")


def norm_key(s):
    """Comparison key: letters/marks/digits only, no ZWJ, no spaces/punctuation, lower-case."""
    s = unicodedata.normalize("NFC", s).lower()
    return "".join(ch for ch in s if unicodedata.category(ch)[0] in "LMN")


def clean_text(s):
    s = unicodedata.normalize("NFC", str(s))
    s = s.replace("\\r\\n", " ").replace("\\n", " ").replace("\\r", " ").replace("\\t", " ")  # literal escapes
    # U+200C (ZWNJ) is a legacy-font-converter artifact after word-final virama; real test text has none
    s = s.replace("​", "").replace("‌", "").replace("", "").replace("�", "")
    s = "".join(ch if (ch in "\n\t" or unicodedata.category(ch) != "Cc") else " " for ch in s)
    return re.sub(r"\s+", " ", s).strip()


def strip_numbering(choices):
    """Remove '(1) ', '2. ', '3) ' prefixes, but only if ALL choices carry the matching sequential prefix."""
    pats = [re.compile(rf"^\s*[\(\[]?{i + 1}[\)\]\.]\s*") for i in range(len(choices))]
    if all(p.match(c) for p, c in zip(pats, choices)):
        return [p.sub("", c, count=1) for p, c in zip(pats, choices)]
    return choices


def garbage_ratio(s):
    if not s:
        return 0.0
    bad = sum(1 for ch in s if "" <= ch <= "" or ch == "�" or "À" <= ch <= "ÿ")
    return bad / len(s)


def sinhala_ratio(s):
    letters = [ch for ch in s if unicodedata.category(ch).startswith("L")]
    if not letters:
        return 1.0
    return sum(1 for ch in letters if "඀" <= ch <= "෿") / len(letters)


def any_match(patterns, text):
    t = strip_zwj(text)
    return any(re.search(p, t, flags=re.I) for p in patterns)


def problems_of(r):
    """Hard errors (record unusable) and soft warnings."""
    errs, warns = [], []
    q, ch, a = r.get("question"), r.get("choices"), r.get("answer")
    if not isinstance(q, str) or len(q.strip()) < 5:
        errs.append("question missing/too short")
    if not isinstance(ch, list) or not (2 <= len(ch) <= 6) or not all(isinstance(c, str) and c.strip() for c in ch):
        errs.append("choices must be 2-6 non-empty strings")
    elif len({norm_key(c) for c in ch}) < len(ch):
        errs.append("duplicate choices")
    if not isinstance(a, int) or isinstance(a, bool):
        errs.append("answer must be an integer (1-based)")
    elif isinstance(ch, list) and not (1 <= a <= len(ch)):
        errs.append("answer out of range (must be 1-based)")
    if not r.get("subject"):
        warns.append("no subject")
    if not getpath(r, "metadata.source"):
        warns.append("no metadata.source (needed for splitting and for removing overlaps later)")
    blob = (q or "") + " " + " ".join(ch if isinstance(ch, list) else [])
    if "\\n" in blob:
        warns.append("literal backslash-n in text")
    if re.search(r"<[a-zA-Z/][^>]*>", blob):
        warns.append("HTML tags in text")
    if garbage_ratio(blob) > 0.05:
        errs.append("garbled text (legacy font / OCR garbage)")
    if sinhala_ratio(blob) < 0.25:
        warns.append("little Sinhala text (English medium or garbled?)")
    if isinstance(q, str) and any_match(FIGURE, q):
        warns.append("refers to a figure/diagram")
    return errs, warns


# ------------------------------------------------------------------ commands
def cmd_example(a):
    rows = [
        {"id": "demo_1", "subject": "science", "category": "stem", "question": "උදාහරණ ප්‍රශ්නය එක කුමක්ද?",
         "choices": ["පිළිතුර එක", "පිළිතුර දෙක", "පිළිතුර තුන", "පිළිතුර හතර"], "answer": 2,
         "metadata": {"source": "demo_paper_A", "exam_type": "OL", "year": 2019, "grade": 11,
                      "difficulty": "medium", "extraction": "text_pdf"}},
        {"id": "demo_2", "subject": "science", "category": "stem", "question": "උදාහරණ ප්‍රශ්නය දෙක කුමක්ද?",
         "choices": ["එක", "දෙක", "තුන", "ඉහත සියල්ලම"], "answer": 4,
         "metadata": {"source": "demo_paper_B", "exam_type": "Scholarship", "year": 2018, "grade": 5,
                      "difficulty": "easy", "extraction": "text_pdf"}},
    ]
    dump(a.out, rows)
    print(f"wrote {a.out}")


JUNK_CHARS = set("-—–~|।*!'\"“”‘’.•_:;,")
LATIN_OK = {"UN", "UNO", "UNESCO", "NATO", "SAARC", "ASEAN", "GDP", "IT", "ICT", "TV", "DNA", "HIV", "AIDS", "WHO",
            "CD", "DVD", "BC", "AD", "pH", "CO2", "H2O", "SMS", "PDF", "USA", "UK", "EU", "AI"}


LATIN_OK |= {"kg", "cm", "mm", "km", "ml", "ms", "cm3", "m2", "m3", "mg", "Hz", "kHz", "MHz", "GHz", "KB", "MB", "GB", "TB",
             "bps", "Kbps", "Mbps", "Ctrl", "Alt", "Shift", "Enter", "Esc", "Tab", "Del", "Delete", "Home", "End", "Page",
             "key", "Key", "ASCII", "USB", "RAM", "ROM", "CPU", "ALU", "BIOS", "OS", "LAN", "WAN", "IP", "MSB", "LSB", "BCD",
             "www", "http", "https", "HTML", "URL", "Data", "Duplex", "control", "AND", "OR", "NOT", "XOR", "NAND", "NOR",
             "II", "III", "IV", "VI", "VII", "VIII", "IX", "Excel", "Word", "Windows", "Access", "PowerPoint", "Internet",
             "Email", "Wi", "Fi", "pH", "DC", "AC", "LED", "LCD", "PIN", "ATM", "VAT", "PAYE", "CEO", "NGO", "SMS", "GPS"}


def chem_like(t):
    """Chemical formula such as HCl, NaOH, CaCO, CO, Mg: element symbols with <=2-digit counts, >=2 distinct chars
    (so 'OO' and 'S888' are NOT accepted)."""
    return bool(re.fullmatch(r"(?:[A-Z][a-z]?\d{0,2}){1,6}", t)) and len(set(t)) > 1


LATIN_OK |= {"of", "to", "in", "and", "or", "the", "gov", "lk", "com", "moe", "edu", "org", "net", "www", "mol", "aq", "ppm", "kgm",
             "Web", "Hub", "Mem", "Sri", "TT", "tt", "CC", "Cc", "Tt", "ii", "iii", "iv", "vi", "vii", "viii", "ix"}


def looks_like_ocr_garbage(t):
    """Real English words / acronyms / formulas are fine inside Sinhala exam text (ICT, science, URLs...).
    Garbage looks different: tokens containing digits ('S888'), very short stray tokens ('wo', 'Qed'),
    or letter soup with no vowel / a long consonant run."""
    if t in LATIN_OK or chem_like(t):
        return False
    if re.search(r"\d", t):
        return True
    if len(t) <= 3:
        return True
    if not re.search(r"[AEIOUaeiouYy]", t):
        return True
    if re.search(r"[^AEIOUaeiouYy]{5,}", t):
        return True
    return False


def latin_garbage_tokens(s):
    """Latin tokens (2+ chars) that look like OCR garbage (e.g. 'S888', 'Qed', 'OO', 'wo')."""
    return [t for t in re.findall(r"[A-Za-z][A-Za-z0-9]*", s) if len(t) >= 2 and looks_like_ocr_garbage(t)]


def trim_tail(s):
    """Remove OCR debris from the END of a question/option: lone punctuation, '|', '।', short quote-led tokens,
    and Latin garbage tokens such as 'Qed,'."""
    toks = s.split(" ")
    while toks:
        t = toks[-1]
        junk_only = t == "" or all(ch in JUNK_CHARS for ch in t)
        quote_led = t[:1] in "“”‘’'\"|।*~" and len(t) <= 3
        latin_junk = bool(latin_garbage_tokens(t)) and len(re.sub(r"[A-Za-z0-9]", "", t)) <= 1
        if junk_only or quote_led or latin_junk:
            toks.pop()
        else:
            break
    return " ".join(toks)


SHARED_STEM = [r"ප්‍රශ්නවලට", r"ප්රශ්නවලට", r"ඇසුරෙන්", r"ඉහත\s*(ප්‍රකාශ|ප්රකාශ|වගුව|සටහන|කොටස|ඡේදය)", r"පහත\s*(ඡේදය|කොටස|ප්‍රකාශ|ප්රකාශ)"]
CONF = {"low": 0, "medium": 1, "high": 2}


def cmd_convert(a):
    """Convert the PDF-extraction format (answer_1..answer_5 + correct_answer letter) to the MCQ schema."""
    rows, out, rej = load(a.inp), [], collections.Counter()
    seen = collections.Counter()
    latin_tok, latin_rows = collections.Counter(), []
    for r in rows:
        conf = str(r.get("extraction_confidence") or "unknown").lower()
        seen[("confidence", conf)] += 1
        seen[("needs_review", bool(r.get("needs_review")))] += 1
        qt = r.get("question_type")
        if qt not in (None, "", "mcq"):
            rej[f"question_type={qt}"] += 1
            continue
        if CONF.get(conf, 0) < CONF[a.min_confidence]:
            rej["below min_confidence"] += 1
            continue
        if a.skip_needs_review and r.get("needs_review"):
            rej["needs_review"] += 1
            continue
        # the training label must exist: rows with a null/empty correct_answer are unusable
        if not str(r.get("correct_answer") or "").strip():
            rej["no correct_answer (null)"] += 1
            continue
        q = clean_text(r.get("original_question") or "")
        # OCR junk after the final question mark, e.g.  "...ද? ී i “"
        if "?" in q:
            head, tail = q.rsplit("?", 1)
            if len(tail.strip()) <= 8:
                q = head + "?"
        q = trim_tail(q)
        raw_choices, spill = [], False
        filled = [i for i in range(1, 6) if r.get(f"answer_{i}") is not None and str(r.get(f"answer_{i}")).strip()]
        last_i = filled[-1] if filled else 0
        for i in range(1, 6):
            v = r.get(f"answer_{i}")
            if v is None or not str(v).strip():
                raw_choices.append(None)
                continue
            lines = [ln.strip() for ln in str(v).strip().split("\n") if ln.strip()]
            keep = [lines[0]]
            for ln in lines[1:]:
                if i < last_i:
                    # another option follows, so nothing can have leaked in: a newline is a wrapped line
                    keep.append(ln)
                else:
                    # last option: continuation lines are short; next-question text is long / a question / numbered
                    looks_like_leak = (len(ln) > 80 or "?" in ln or re.match(r"^\(?\d+[\.\)]", ln)
                                       or any_match(SHARED_STEM, ln))
                    if looks_like_leak:
                        spill = True
                        break
                    keep.append(ln)
            first = clean_text(" ".join(keep))
            first = trim_tail(first).rstrip(" ,;:")
            raw_choices.append(first or None)
        # options must be contiguous from 1
        n = 0
        while n < 5 and raw_choices[n]:
            n += 1
        if any(raw_choices[n:]):
            rej["non-contiguous options"] += 1
            continue
        choices = raw_choices[:n]
        if n < 3:
            rej["fewer than 3 options"] += 1
            continue
        ca = str(r.get("correct_answer") or "").strip().upper()
        if ca in "ABCDE" and ca:
            ans = "ABCDE".index(ca) + 1
        elif ca.isdigit():
            ans = int(ca)
        else:
            rej["missing/invalid correct_answer"] += 1
            continue
        if not (1 <= ans <= n):
            rej["correct_answer outside the options"] += 1
            continue
        if len(q) < 8:
            rej["question too short"] += 1
            continue
        if any_match(SHARED_STEM, q) and not a.keep_shared_stem:
            rej["depends on a shared stem/passage"] += 1
            continue
        if max(len(c) for c in choices) > 250:
            rej["option too long (extraction spill)"] += 1
            continue
        if not a.keep_latin:
            allow = {t for t in (a.latin_allow or [])}
            bad = [t for t in latin_garbage_tokens(q + " " + " ".join(choices)) if t not in allow]
            if bad:
                rej["OCR garbage (Latin letters/digits inside Sinhala text)"] += 1
                latin_tok.update(bad)
                latin_rows.append({"id": r.get("id"), "subject": r.get("subject"), "tokens": bad,
                                   "question": q, "choices": choices})
                continue
        out.append({
            "id": r.get("id"), "subject": r.get("subject"), "question": q, "choices": choices, "answer": ans,
            "metadata": {"source": r.get("source_pdf"), "page": r.get("source_page"), "year": r.get("year"),
                         "exam_type": r.get("level"), "medium": r.get("medium"),
                         "extraction_confidence": conf, "needs_review": bool(r.get("needs_review")),
                         "option_spill_fixed": spill}})
    dump(a.out, out)
    print(f"read {len(rows)} -> kept {len(out)} ({len(out) / max(1, len(rows)):.0%}) -> {a.out}")
    for (k, v), c in sorted(seen.items(), key=lambda kv: str(kv[0])):
        print(f"  input {k}={v}: {c}")
    for k, v in rej.most_common():
        print(f"  dropped: {k}: {v}")
    if latin_rows:
        print("  most common Latin tokens in the rows dropped as 'OCR garbage' (legit terms like CPU/Excel/NaCl "
              "should be allowed with --latin_allow):")
        print("   ", ", ".join(f"{t}:{c}" for t, c in latin_tok.most_common(40)))
        if a.rejects:
            dump(a.rejects, latin_rows)
            print(f"  dropped Latin rows saved to {a.rejects}")
    if out and all(x["metadata"]["needs_review"] for x in out):
        print("WARNING: every kept record is flagged needs_review. Label noise from extraction is the main risk: "
              "hand-check ~100 random records (python mcq_dataset.py sample ...) before training.")


def cmd_sample(a):
    rows = load(a.inp)
    random.Random(a.seed).shuffle(rows)
    lines = []
    for r in rows[:a.n]:
        md = r.get("metadata") or {}
        lines.append("-" * 60 + f" {r.get('id')}  [{md.get('extraction_confidence', '?')}]")
        lines.append("Q: " + r["question"])
        for i, c in enumerate(r["choices"], 1):
            lines.append(f"  {'*' if i == r['answer'] else ' '} {i}. {c}")
    text = "\n".join(lines) + "\n"
    if a.out:
        # write the file directly as UTF-8 (a PowerShell '>' redirect re-encodes and garbles Sinhala)
        with open(a.out, "w", encoding="utf-8-sig") as f:
            f.write(text)
        print(f"wrote {min(a.n, len(rows))} records -> {a.out}")
    else:
        print(text, end="")


def cmd_validate(a):
    rows = load(a.inp)
    ec, wc = collections.Counter(), collections.Counter()
    bad = 0
    for r in rows:
        errs, warns = problems_of(r)
        bad += bool(errs)
        ec.update(errs)
        wc.update(warns)
    print(f"{len(rows)} records, {bad} with hard errors ({bad / max(1, len(rows)):.1%})")
    for k, v in ec.most_common():
        print(f"  ERROR  {v:6d}  {k}")
    for k, v in wc.most_common():
        print(f"  warn   {v:6d}  {k}")
    zw = sum(ZWJ in (r.get("question", "") + "".join(r.get("choices", []) or [])) for r in rows)
    print(f"records containing ZWJ (U+200D): {zw} ({zw / max(1, len(rows)):.1%})")


def cmd_clean(a):
    out, rej = [], []
    for r in load(a.inp):
        r = dict(r)
        if isinstance(r.get("answer"), str) and r["answer"].strip().isdigit():
            r["answer"] = int(r["answer"].strip())
        if isinstance(r.get("question"), str):
            r["question"] = clean_text(r["question"])
        if isinstance(r.get("choices"), list):
            ch = [clean_text(c) for c in r["choices"] if isinstance(c, str)]
            r["choices"] = strip_numbering(ch)
        errs, warns = problems_of(r)
        if "refers to a figure/diagram" in warns and not a.keep_figures:
            errs.append("needs a figure/diagram")
        if errs:
            rej.append({**r, "_reject": errs})
        else:
            out.append(r)
    dump(a.out, out)
    if a.rejects:
        dump(a.rejects, rej)
    print(f"kept {len(out)}, rejected {len(rej)}  -> {a.out}")
    for k, v in collections.Counter(e for x in rej for e in x["_reject"]).most_common():
        print(f"  rejected for: {k}: {v}")


def shingles(key, n=8):
    return {hash(key[i:i + n]) for i in range(max(1, len(key) - n + 1))}


def rec_key(r):
    return norm_key(r.get("question", "") + " " + " ".join(r.get("choices", [])))


class Index:
    def __init__(self):
        self.post = collections.defaultdict(list)
        self.n = 0

    def add(self, sh):
        for h in sh:
            self.post[h].append(self.n)
        self.n += 1

    def best_containment(self, sh):
        if not sh:
            return 0.0
        cnt = collections.Counter()
        for h in sh:
            for d in self.post.get(h, ()):
                cnt[d] += 1
        return max(cnt.values()) / len(sh) if cnt else 0.0


def cmd_dedupe(a):
    rows = load(a.inp)
    # 1) same question text: identical answer -> keep first, conflicting answers -> drop all (likely label noise)
    groups = collections.defaultdict(list)
    for r in rows:
        groups[norm_key(r["question"])].append(r)
    kept, conflicts = [], []
    for k, g in groups.items():
        ans = {norm_key(x["choices"][x["answer"] - 1]) for x in g}
        if len(ans) > 1:
            conflicts += g
        else:
            kept.append(g[0])
    dups = len(rows) - len(kept) - len(conflicts)
    # 2) near-duplicates inside the set
    idx, uniq, near = Index(), [], 0
    for r in kept:
        sh = shingles(rec_key(r))
        if idx.best_containment(sh) >= a.near_threshold:
            near += 1
            continue
        idx.add(sh)
        uniq.append(r)
    # 3) overlap with blocklist (e.g. public test questions): drop, never copied
    overlap = 0
    if a.blocklist:
        bfull, bq = Index(), Index()
        for p in a.blocklist:
            for b in load(p):
                if isinstance(b, dict) and b.get("question"):
                    bfull.add(shingles(norm_key(b["question"] + " " + " ".join(b.get("choices", []) or []))))
                    bq.add(shingles(norm_key(b["question"])))
        final = []
        for r in uniq:
            qk = norm_key(r["question"])
            # drop if the whole item matches, OR the question text alone matches (robust to differently formatted options)
            hit = bfull.best_containment(shingles(rec_key(r))) >= a.block_threshold or \
                (len(qk) >= 25 and bq.best_containment(shingles(qk)) >= 0.8)
            if hit:
                overlap += 1
            else:
                final.append(r)
        uniq = final
    dump(a.out, uniq)
    if conflicts and a.conflicts:
        dump(a.conflicts, conflicts)
    print(f"in {len(rows)} -> out {len(uniq)} | exact dup {dups} | answer conflicts {len(conflicts)} | "
          f"near-dup {near} | overlap with blocklist {overlap}")


def cmd_split(a):
    rows = load(a.inp)
    tr, va = [], []
    groups = set()
    for r in rows:
        g = str(getpath(r, a.group_key) or r.get("id", "")).strip() or "none"
        groups.add(g)
        h = int(hashlib.sha1(f"{a.seed}:{g}".encode()).hexdigest()[:8], 16) / 2 ** 32
        (va if h < a.val_frac else tr).append(r)
    dump(a.train, tr)
    dump(a.val, va)
    print(f"{len(groups)} groups ({a.group_key}) -> train {len(tr)}, val {len(va)}")
    if len(groups) < 10:
        print("WARNING: few groups; the split may be lopsided. Add metadata.source to every record.")


def cmd_stats(a):
    rows = load(a.inp)
    print(f"{len(rows)} records")
    for title, f in [("subject", lambda r: r.get("subject")), ("category", lambda r: r.get("category")),
                     ("difficulty", lambda r: getpath(r, "metadata.difficulty")),
                     ("exam_type", lambda r: getpath(r, "metadata.exam_type")),
                     ("year", lambda r: getpath(r, "metadata.year")),
                     ("n_choices", lambda r: len(r.get("choices", []))),
                     ("answer position", lambda r: r.get("answer")),
                     ("extraction", lambda r: getpath(r, "metadata.extraction_confidence") or getpath(r, "metadata.extraction"))]:
        c = collections.Counter(f(r) for r in rows)
        print(f"\n{title}:")
        for k, v in sorted(c.items(), key=lambda kv: -kv[1])[:20]:
            print(f"  {str(k):28s} {v:6d}  {v / len(rows):6.1%}")


def render(r, tpl, rng):
    ch = r["choices"]
    letters = "ABCDEF"
    subj = r.get("subject") or ""
    if tpl == "bare":
        body = r["question"] + "\n" + "\n".join(f"{i + 1}. {c}" for i, c in enumerate(ch)) + "\nපිළිතුර:"
        return body, f" {r['answer']}"
    if tpl == "si_instruction":
        head = SI_INSTR + "\n\n" + (f"විෂය: {subj}\n" if subj else "")
        body = head + f"ප්‍රශ්නය: {r['question']}\n" + "\n".join(f"{i + 1}. {c}" for i, c in enumerate(ch)) + "\nපිළිතුර:"
        return body, f" {r['answer']}"
    if tpl == "en_instruction":
        head = EN_INSTR + "\n\n" + (f"Subject: {subj}\n" if subj else "")
        body = head + f"Question: {r['question']}\n" + "\n".join(f"{i + 1}. {c}" for i, c in enumerate(ch)) + "\nAnswer:"
        return body, f" {r['answer']}"
    if tpl == "letters":
        body = f"Question: {r['question']}\n" + "\n".join(f"{letters[i]}. {c}" for i, c in enumerate(ch)) + "\nAnswer:"
        return body, f" {letters[r['answer'] - 1]}"
    raise ValueError(tpl)


TEMPLATE_WEIGHTS = {"si_instruction": 0.35, "bare": 0.25, "en_instruction": 0.25, "letters": 0.15}


def shuffle_safe(r):
    return not any(any_match(POSITIONAL, c) for c in r["choices"])


def cmd_build_sft(a):
    rng = random.Random(a.seed)
    names = a.templates or list(TEMPLATE_WEIGHTS)
    weights = [TEMPLATE_WEIGHTS[n] for n in names]
    rows, out, n_shuf_skip = load(a.inp), [], 0
    for r in rows:
        variants = [r]
        if a.shuffle_copies > 0:
            if shuffle_safe(r):
                for _ in range(a.shuffle_copies):
                    perm = list(range(len(r["choices"])))
                    rng.shuffle(perm)
                    new_choices = [r["choices"][i] for i in perm]
                    new_ans = perm.index(r["answer"] - 1) + 1
                    variants.append({**r, "choices": new_choices, "answer": new_ans})
            else:
                n_shuf_skip += 1
        for v in variants:
            v = dict(v)
            if rng.random() < a.strip_zwj_prob:
                v["question"] = strip_zwj(v["question"])
                v["choices"] = [strip_zwj(c) for c in v["choices"]]
                v["subject"] = strip_zwj(v.get("subject") or "")
            tpl = rng.choices(names, weights)[0]
            prompt, completion = render(v, tpl, rng)
            if a.format == "chat":
                rec = {"prompt": [{"role": "user", "content": prompt}],
                       "completion": [{"role": "assistant", "content": completion.strip()}]}
            else:
                lab = "ABCDEF" if tpl == "letters" else "123456"
                rec = {"prompt": prompt, "completion": completion,
                       "candidates": [f" {lab[i]}" for i in range(len(v["choices"]))],
                       "subject": r.get("subject")}
            out.append(rec)
    rng.shuffle(out)
    dump(a.out, out)
    pos = collections.Counter(
        (x["completion"] if isinstance(x["completion"], str) else x["completion"][0]["content"]).strip()
        for x in out)
    print(f"{len(rows)} questions -> {len(out)} training examples ({a.format}); "
          f"not shuffled (positional options): {n_shuf_skip}")
    print("answer label distribution:", dict(sorted(pos.items())))


def main():
    # Windows consoles/redirects default to cp1252, which cannot encode Sinhala: force UTF-8
    for stream in (sys.stdout, sys.stderr):
        try:
            stream.reconfigure(encoding="utf-8")
        except Exception:
            pass
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)
    p = sub.add_parser("example"); p.add_argument("out"); p.set_defaults(f=cmd_example)
    p = sub.add_parser("convert", help="PDF-extraction format -> MCQ schema")
    p.add_argument("inp"); p.add_argument("out")
    p.add_argument("--min_confidence", choices=["low", "medium", "high"], default="low")
    p.add_argument("--skip_needs_review", action="store_true")
    p.add_argument("--keep_shared_stem", action="store_true")
    p.add_argument("--rejects", default="", help="save the rows dropped for Latin garbage to this file")
    p.add_argument("--latin_allow", nargs="*", default=[], help="Latin tokens that are legitimate (e.g. CPU Excel RAM)")
    p.add_argument("--keep_latin", action="store_true",
                   help="do not drop records with Latin tokens (use for science/maths/English papers)")
    p.set_defaults(f=cmd_convert)
    p = sub.add_parser("sample", help="print random records for hand-checking"); p.add_argument("inp")
    p.add_argument("-n", type=int, default=20); p.add_argument("--seed", type=int, default=1)
    p.add_argument("--out", default="", help="write to this file as UTF-8 (use this instead of '>' on Windows)")
    p.set_defaults(f=cmd_sample)
    p = sub.add_parser("validate"); p.add_argument("inp"); p.set_defaults(f=cmd_validate)
    p = sub.add_parser("clean"); p.add_argument("inp"); p.add_argument("out")
    p.add_argument("--rejects", default=""); p.add_argument("--keep_figures", action="store_true"); p.set_defaults(f=cmd_clean)
    p = sub.add_parser("dedupe"); p.add_argument("inp"); p.add_argument("out")
    p.add_argument("--blocklist", nargs="*", default=[], help="files of questions to remove overlaps with")
    p.add_argument("--near_threshold", type=float, default=0.9)
    p.add_argument("--block_threshold", type=float, default=0.5)
    p.add_argument("--conflicts", default="conflicts.jsonl"); p.set_defaults(f=cmd_dedupe)
    p = sub.add_parser("split"); p.add_argument("inp"); p.add_argument("--train", required=True)
    p.add_argument("--val", required=True); p.add_argument("--val_frac", type=float, default=0.1)
    p.add_argument("--group_key", default="metadata.source"); p.add_argument("--seed", type=int, default=3407)
    p.set_defaults(f=cmd_split)
    p = sub.add_parser("stats"); p.add_argument("inp"); p.set_defaults(f=cmd_stats)
    p = sub.add_parser("build-sft"); p.add_argument("inp"); p.add_argument("out")
    p.add_argument("--format", choices=["plain", "chat"], default="plain")
    p.add_argument("--templates", nargs="*", choices=list(TEMPLATE_WEIGHTS), default=None)
    p.add_argument("--shuffle_copies", type=int, default=1)
    p.add_argument("--strip_zwj_prob", type=float, default=0.5)
    p.add_argument("--seed", type=int, default=3407); p.set_defaults(f=cmd_build_sft)
    a = ap.parse_args()
    a.f(a)


if __name__ == "__main__":
    main()