#!/usr/bin/env python3
"""
Render README.md for the MSQ bucket from results/metrics.json + repr_probe.json.

The prose lives in code/README_template.md; every number in it is injected
from the measured results, so the README can never drift from the data.
"""
from __future__ import annotations

import json
import pathlib
import re

ROOT = pathlib.Path(__file__).resolve().parent.parent
TEMPLATE = ROOT / "code" / "README_template.md"
OUT = ROOT / "README.md"

MODEL_LABEL = {
    "minicpm5-2b": "MiniCPM5-2B (primary)",
    "minicpm5-2b-base": "MiniCPM5-2B-Base (no SFT/RL)",
    "minicpm5-2b-sft": "MiniCPM5-2B-SFT (no RL)",
    "qwen3-4b": "Qwen3-4B-Instruct-2507",
    "smollm3-3b": "SmolLM3-3B",
    "phi4-mini": "Phi-4-mini-instruct",
    "tiny-test": "tiny-random (smoke test)",
}


def pct(x, nd=0):
    if x is None:
        return "n/a"
    return f"{x * 100:.{nd}f}%"


def num(x, nd=3):
    if x is None:
        return "n/a"
    return f"{x:.{nd}f}"


def build_qualitative(root: pathlib.Path, prim: str) -> str:
    """Pull real generations and real logit distributions from the primary model.

    Three vignettes, each auto-selected as the most striking example of its
    kind, so the section is never hand-picked to flatter the result.
    """
    p = root / "results" / "raw" / f"{prim}.jsonl"
    if not p.exists():
        return "_(raw generations not yet available)_"
    recs = [json.loads(l) for l in p.open(encoding="utf-8") if l.strip()]
    items = {i["id"]: i for i in (json.loads(l) for l in
                                 (root / "dataset" / "msq49.jsonl").open())}

    def get(probe, iid):
        for r in recs:
            if r["probe"] == probe and r["item_id"] == iid:
                return r
        return None

    def clip(s, n=700):
        s = " ".join((s or "").split())
        return (s[:n] + " …") if len(s) > n else s

    def dist_lines(r, only=None):
        out = []
        for i, (o, pr) in enumerate(zip(r["options"], r["probs"])):
            if only is not None and i != only:
                continue
            mark = "  ← chosen" if i == r["argmax_index"] else ""
            out.append(f"   **{chr(65+i)}**  p={pr:.3f}   {clip(o, 105)}{mark}")
        return "\n".join(out)

    out = []

    # ---- vignette 1: the instruction changed nothing -----------------------
    cands = []
    for r in recs:
        if r["probe"] != "mcq_neutral" or r["cls"] != "A":
            continue
        s = get("mcq_specified", r["item_id"])
        if not s or len(s["probs"]) != len(r["probs"]):
            continue
        tv = 0.5 * sum(abs(a - b) for a, b in zip(r["probs"], s["probs"]))
        if r["argmax_index"] == s["argmax_index"]:
            cands.append((tv, r, s))
    if cands:
        cands.sort()
        tv, neu, spec = cands[0]
        it = items[neu["item_id"]]
        out.append(f"**(a) The instruction changed nothing.** Item {neu['item_id']} — "
                   f"_{it['question']}_")
        out.append("")
        out.append(f"Ground truth: the schools are divided — "
                   f"{it['majority_schools']} vs {it['minority_school']}.")
        out.append("")
        out.append("*Nothing named:*")
        out.append("")
        out.append(dist_lines(neu))
        out.append("")
        out.append(f"*\"According to the {it['minority_school']} school …\":*")
        out.append("")
        out.append(dist_lines(spec))
        out.append("")
        out.append(f"Total-variation distance between the two distributions: "
                   f"**{tv:.3f}**. The model picked the same option with the same "
                   f"confidence, before and after being told which school it wanted.")
        out.append("")

    # ---- vignette 2: it volunteers the dependency and then ignores it ------
    cands = []
    for r in recs:
        if r["probe"] != "mcq_neutral" or r["cls"] != "C":
            continue
        s = get("mcq_specified", r["item_id"])
        if not s:
            continue
        if s["argmax_index"] == s.get("divided_index"):
            cands.append((s["max_prob"], r, s))
    if cands:
        cands.sort(key=lambda x: -x[0])
        _, neu, spec = cands[0]
        it = items[neu["item_id"]]
        out.append(f"**(b) It notices the dependency, and then ignores it.** Item "
                   f"{neu['item_id']} — _{it['question']}_")
        out.append("")
        out.append("*Nothing named:*")
        out.append("")
        out.append(dist_lines(neu))
        out.append("")
        out.append(f"*\"The user lives in {it['minority_school']}\":*")
        out.append("")
        out.append(dist_lines(spec))
        out.append("")
        out.append("The correct option is printed in the list, matching the stated "
                   "jurisdiction exactly. The model still answers with the generic "
                   "one.")
        out.append("")

    # ---- vignette 3: it invents a school ------------------------------------
    cands = []
    known = {}
    for it in items.values():
        known[it["id"]] = {s.lower() for p in it["positions"] for s in p["schools"]}
    for r in recs:
        if r["probe"] != "ft_two_schools" or r["cls"] != "A":
            continue
        body = (r["answer"] or "").lower()
        mentioned = {s for s in known[r["item_id"]] if s.split()[0] in body}
        # count capitalised "School" style tokens that are not in the ground truth
        invented = [w for w in ("contemporary", "modern", "progressive", "liberal",
                                "reformist", "rationalist", "modernist")
                    if w in body]
        if invented:
            cands.append((len(invented), r))
    if cands:
        cands.sort(key=lambda x: -x[0])
        _, r = cands[0]
        it = items[r["item_id"]]
        out.append(f"**(c) Asked to enumerate the competing positions, it invents "
                   f"one.** Item {r['item_id']} — _{it['question']}_")
        out.append("")
        out.append(f"Ground truth: {it['majority_schools']} vs {it['minority_school']}. "
                   f"Prompt: *\"Explain the different scholarly positions … name the "
                   f"school of thought that holds each.\"*")
        out.append("")
        out.append("> " + clip(r["answer"], 900))
        out.append("")

    # ---- vignette 4: free text, plain vs named ------------------------------
    for iid in ("A11", "A01", "A02"):
        fd, fs = get("ft_default", iid), get("ft_specified", iid)
        if fd and fs:
            it = items[iid]
            out.append(f"**(d) Free text, item {iid}.** _{it['question']}_")
            out.append("")
            out.append(f"*Plain question:*")
            out.append("")
            out.append("> " + clip(fd["answer"], 620))
            out.append("")
            out.append(f"*\"The user follows the {it['minority_school']} school\":*")
            out.append("")
            out.append("> " + clip(fs["answer"], 620))
            out.append("")
            break

    return "\n".join(out)


def main() -> None:
    metrics = json.loads((ROOT / "results" / "metrics.json").read_text())
    reprp = ROOT / "results" / "repr_probe.json"
    reprd = json.loads(reprp.read_text()) if reprp.exists() else {}
    items = [json.loads(l) for l in (ROOT / "dataset" / "msq49.jsonl").open()]

    vals: dict[str, str] = {}
    order = [m for m in ("minicpm5-2b", "minicpm5-2b-base", "minicpm5-2b-sft",
                          "qwen3-4b", "smollm3-3b", "phi4-mini") if m in metrics]
    if not order:
        order = list(metrics)

    # headline table
    params = {"minicpm5-2b": "2.52B", "minicpm5-2b-base": "2.52B",
              "minicpm5-2b-sft": "2.52B", "qwen3-4b": "4.02B",
              "smollm3-3b": "3.08B", "phi4-mini": "3.82B"}
    lines = ["| Model | Params | context response (TV) | answer unchanged by context "
             "| PRR (fiqh) | MSQ (fiqh) | Δ on the correct answer | control accuracy |",
             "|---|---|---|---|---|---|---|---|"]
    for m in order:
        h = metrics[m]["headline"]
        cr = metrics[m].get("context_response", {}).get("ALL", {})
        lines.append(
            f"| {MODEL_LABEL.get(m, m)} | {params.get(m, '?')} | "
            f"{num(cr.get('mean_context_response_tv'), 3)} | "
            f"{pct(cr.get('argmax_unchanged_rate'))} | "
            f"{pct(h.get('PRR_A'))} | {pct(h.get('MSQ_A'))} | "
            f"{num(cr.get('median_delta_on_named_position'), 3)} | "
            f"{pct(h.get('acc_D'))} |")
    vals["headline_table"] = "\n".join(lines)

    prim = "minicpm5-2b" if "minicpm5-2b" in metrics else order[0]
    ph = metrics[prim]["headline"]
    ft = metrics[prim]["free_text"]
    mc = metrics[prim]["mcq"]

    vals["primary"] = MODEL_LABEL.get(prim, prim)
    vals["n_items"] = str(len(items))
    vals["n_models"] = str(len(order))
    contested = [i for i in items if i["cls"] in ("A", "B", "C", "E")]
    vals["n_contested"] = str(len(contested))
    import statistics as _st
    chance = []
    for it in contested:
        n = len(it["positions"]) + (1 if it["cls"] in ("A", "B") else 0)
        if it["cls"] == "C":
            n = len(it["positions"]) + len(it.get("jur_map", {}))
        chance.append(1.0 / n)
    vals["chance_rate"] = pct(round(_st.fmean(chance), 4))
    vals["prr_A"] = pct(ph.get("PRR_A"))
    vals["prr_B"] = pct(ph.get("PRR_B"))
    vals["prr_C"] = pct(ph.get("PRR_C"))
    vals["prr_E"] = pct(ph.get("PRR_E"))
    vals["msq_A"] = pct(ph.get("MSQ_A"))
    vals["msq_C"] = pct(ph.get("MSQ_C"))
    vals["msq_E"] = pct(ph.get("MSQ_E"))
    vals["prr_all"] = pct(ph.get("PRR"))
    vals["msq_all"] = pct(ph.get("MSQ"))
    vals["gap_all"] = pct(ph.get("blind_spot_gap"))
    vals["prr_mass"] = pct(ph.get("PRR_mass"))
    vals["acc_D"] = pct(ph.get("acc_D"))

    vals["disc_default"] = pct(ft.get("ft_default", {}).get("ALL", {}).get("disclosure_rate"))
    vals["disc_specified"] = pct(ft.get("ft_specified", {}).get("ALL", {}).get("disclosure_rate"))
    vals["disc_two"] = pct(ft.get("ft_two_schools", {}).get("ALL", {}).get("disclosure_rate"))
    vals["disc_technical"] = pct(ft.get("ft_technical", {}).get("ALL", {}).get("disclosure_rate"))
    vals["attr_default"] = pct(ft.get("ft_default", {}).get("ALL", {}).get("school_attribution_rate"))
    vals["attr_two"] = pct(ft.get("ft_two_schools", {}).get("ALL", {}).get("school_attribution_rate"))
    vals["attr_technical"] = pct(ft.get("ft_technical", {}).get("ALL", {}).get("school_attribution_rate"))
    vals["clar_default"] = pct(ft.get("ft_default", {}).get("ALL", {}).get("clarification_rate"))
    vals["clar_two"] = pct(ft.get("ft_two_schools", {}).get("ALL", {}).get("clarification_rate"))
    vals["hedge_default"] = pct(ft.get("ft_default", {}).get("ALL", {}).get("hedge_rate"))
    vals["hedge_two"] = pct(ft.get("ft_two_schools", {}).get("ALL", {}).get("hedge_rate"))
    vals["words_default"] = num(ft.get("ft_default", {}).get("ALL", {}).get("mean_words"), 1)
    vals["words_two"] = num(ft.get("ft_two_schools", {}).get("ALL", {}).get("mean_words"), 1)

    vals["conf_contested"] = pct(ph.get("conf_contested"))
    vals["conf_control"] = pct(ph.get("conf_control"))
    vals["change_rate"] = pct(metrics[prim]["answer_change"].get("changed_rate"))
    vals["similarity"] = num(metrics[prim]["answer_change"].get("mean_similarity"), 2)

    for cls, tag in (("A", "fiqh"), ("B", "controversy"), ("C", "jurisdiction"),
                     ("E", "python version"), ("D", "control")):
        vals[f"conf_{cls}"] = pct(mc.get("mcq_neutral", {}).get(cls, {}).get("mean_max_prob"))
        vals[f"prr_{cls}"] = pct(mc.get("mcq_neutral", {}).get(cls, {}).get("plurality_recognition_rate"))

    pc = metrics[prim].get("positions_contained", {})
    for probe, key in (("ft_default", "positions_default"),
                       ("ft_two_schools", "positions_two_schools"),
                       ("ft_pair", "positions_pair"),
                       ("ft_repair", "positions_repair"),
                       ("ft_technical", "positions_technical")):
        vals[key] = num(pc.get(probe, {}).get("ALL", {}).get("mean_positions_contained"), 2)
    vals["disc_repair"] = pct(ft.get("ft_repair", {}).get("ALL", {}).get("disclosure_rate"))
    vals["attr_repair"] = pct(ft.get("ft_repair", {}).get("ALL", {}).get("plural_attribution_rate"))

    ca = metrics[prim].get("content_adherence", {})
    vals["names_context_rate"] = pct(ca.get("ALL", {}).get("names_context_rate"))
    vals["content_adherence_rate"] = pct(ca.get("ALL", {}).get("content_adherence_rate"))
    vals["label_adoption_rate"] = pct(ca.get("ALL", {}).get("label_adoption_rate"))
    vals["split_ratio"] = num(ca.get("ALL", {}).get("mean_split_ratio"), 3)

    bm = metrics[prim].get("behaviour_matrix", {})
    vals["bm_both"] = pct(bm.get("ALL", {}).get("both"))
    vals["bm_default_only"] = pct(bm.get("ALL", {}).get("default_only"))
    vals["bm_switch_only"] = pct(bm.get("ALL", {}).get("switch_only"))
    vals["bm_neither"] = pct(bm.get("ALL", {}).get("neither"))

    cr = metrics[prim].get("context_response", {})
    for cls in ("A", "B", "C", "E", "ALL"):
        c = cr.get(cls, {})
        vals[f"cr_{cls}"] = num(c.get("mean_context_response_tv"), 3)
        vals[f"unchanged_{cls}"] = pct(c.get("argmax_unchanged_rate"))
        vals[f"delta_{cls}"] = num(c.get("median_delta_on_named_position"), 4)
        vals[f"n_unchanged_{cls}"] = str(c.get("n_unchanged_distributions", 0))
    so = metrics[prim].get("specified_outcome", {}).get("mcq_specified", {})
    for cls in ("A", "C", "E", "ALL"):
        s = so.get(cls, {})
        vals[f"so_named_{cls}"] = pct(s.get("named_school"))
        vals[f"so_divided_{cls}"] = pct(s.get("divided_anyway"))
        vals[f"so_other_{cls}"] = pct(s.get("other_position"))
    sn = metrics[prim].get("specified_outcome", {}).get("mcq_neutral", {})
    for cls in ("A", "B", "C", "E", "ALL"):
        vals[f"neut_div_{cls}"] = pct(sn.get(cls, {}).get("divided"))
        vals[f"neut_pos_{cls}"] = pct(sn.get(cls, {}).get("a_position"))

    # dataset-level facts
    lone = 0
    for it in items:
        if it["cls"] != "A":
            continue
        maj = set(it["majority_schools"])
        other = set()
        for p in it["positions"]:
            for s in p["schools"]:
                if s != it["minority_school"]:
                    other.add(s)
        if it["minority_school"] not in maj and len(other) == 3:
            lone += 1
    vals["n_lone_outlier"] = str(lone)

    # tokenizer
    if reprd:
        f = reprd["fertility_by_script"]
        vp = reprd["vocab_script_profile"]
        tm = reprd["term_morphology"]
        vals["fert_en"] = num(f["en"]["tokens_per_char"], 4)
        vals["fert_ar"] = num(f["ar_plain"]["tokens_per_char"], 4)
        vals["fert_arv"] = num(f["ar_vocalized"]["tokens_per_char"], 4)
        vals["fert_ur"] = num(f["ur"]["tokens_per_char"], 4)
        vals["fert_zh"] = num(f["zh"]["tokens_per_char"], 4)
        vals["byte_en"] = pct(f["en"]["byte_level_rate"])
        vals["byte_zh"] = pct(f["zh"]["byte_level_rate"])
        vals["byte_ur"] = pct(f["ur"]["byte_level_rate"])
        vals["byte_ar"] = pct(f["ar_plain"]["byte_level_rate"])
        vals["byte_arv"] = pct(f["ar_vocalized"]["byte_level_rate"])
        vals["vocab_size"] = f"{vp['vocab_size']:,}"
        vals["n_byte_tokens"] = f"{vp['n_byte_level_tokens']:,}"
        vals["share_byte_level"] = pct(vp["share_byte_level"], 1)
        vals["n_latin_units"] = f"{vp['n_latin_text_units']:,}"
        vals["n_cjk_units"] = f"{vp['n_cjk_text_units']:,}"
        vals["n_arabic_units"] = str(vp["n_arabic_text_units"])
        vals["share_arabic_units"] = pct(vp["share_arabic_text_units"], 3)
        vals["max_arabic_chars"] = str(vp["max_arabic_chars_in_one_unit"])
        vals["mean_arabic_chars"] = num(vp["mean_arabic_chars_per_unit"], 2)
        vals["terms_plain"] = num(tm["mean_tokens_plain"], 2)
        vals["terms_voc"] = num(tm["mean_tokens_vocalized"], 2)
        vals["terms_ratio"] = num(tm["mean_diacritic_cost_ratio"], 2)
        vals["merge_rate"] = pct(reprd["diacritic_discrimination"]["merge_rate"])
        vals["ar_x_ferrite"] = num(f["ar_vocalized"]["tokens_per_char"] /
                                   max(f["en"]["tokens_per_char"], 1e-9), 1)

    # qualitative examples straight from the raw generations
    vals["qual_examples"] = build_qualitative(ROOT, prim)

    # per-model param counts
    for m, p in (("minicpm5-2b", "2.52B"), ("minicpm5-2b-base", "2.52B"),
                 ("minicpm5-2b-sft", "2.52B"), ("qwen3-4b", "4.02B"),
                 ("smollm3-3b", "3.08B"), ("phi4-mini", "3.82B")):
        vals[f"params_{m}"] = p

    txt = TEMPLATE.read_text(encoding="utf-8")

    def sub(m):
        key = m.group(1)
        if key not in vals:
            raise KeyError(f"template placeholder missing: {key}")
        return vals[key]

    txt = re.sub(r"\{\{(\w+)\}\}", sub, txt)
    OUT.write_text(txt, encoding="utf-8")
    print(f"wrote {OUT} ({len(txt)} chars)")


if __name__ == "__main__":
    main()