#!/usr/bin/env python3
"""
Aggregate raw MSQ probe output into the paper's metrics, tables and figures.

Headline metrics
----------------
PRR  Plurality Recognition Rate
     On contested items, how often does the model choose the "the scholars
     are divided" option when the options are presented without any school
     label? High PRR = the model sees the dispute.

MSQ  Madhhab Sensitivity Quotient
     On contested items, how often does the model choose the position of the
     *named* school when a school is named explicitly? High MSQ = the model
     can produce a non-default position on request.

BLIND-SPOT GAP  = PRR(divided) when nothing is named  vs  MSQ when it is.
     A large gap means the model holds the competing positions but does not
     surface them: it can represent the plurality, but will not volunteer it.

DISC  Free-text disclosure rate -- fraction of answers that signal that the
     question depends on an unstated context variable.

CONF  Confidence calibration -- mean probability the model puts on its own
     chosen option, contested items versus uncontested controls.
"""
from __future__ import annotations

import argparse
import difflib
import json
import pathlib
import re
import statistics
import sys
from collections import Counter, defaultdict

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
import msq  # noqa: E402

ROOT = pathlib.Path(__file__).resolve().parent.parent
RAW = ROOT / "results" / "raw"
OUTT = ROOT / "results"

STOP = set("""the a an of is are was were be been being to in on at for from by with
and or not no it its this that these those as if then than so such their they them
there here which who whom what when where how do does did can could may might will
would shall should must have has had you your i me my we our us he she his her
but also about into over under more most some any all one two per each other
according school position scholars view views ruling rule question answer""".split())


def content_tokens(s: str) -> list[str]:
    return [w for w in re.findall(r"[a-z0-9']+", (s or "").lower())
            if w not in STOP and len(w) > 2]


def content_sim(a: str, b: str) -> float:
    """Token F1 over content words -- how much of a ground-truth position the
    answer actually contains."""
    ta, tb = content_tokens(a), content_tokens(b)
    if not ta or not tb:
        return 0.0
    ca, cb = Counter(ta), Counter(tb)
    inter = sum((ca & cb).values())
    if inter == 0:
        return 0.0
    p, r = inter / len(ta), inter / len(tb)
    return 2 * p * r / (p + r)


def names(text: str, name: str) -> bool:
    """Does the answer name this school / version / jurisdiction?"""
    if not name:
        return False
    norm = re.sub(r"[^a-z0-9]", "", (text or "").lower())
    target = re.sub(r"[^a-z0-9]", "", name.lower())
    if not target:
        return False
    if target in norm:
        return True
    # 'shafi'i' -> 'shafii', 'python 2' -> 'python2'
    alt = target.replace("i", "").replace("python", "")
    if alt and alt in norm and len(alt) >= 4:
        return True
    return False


def load_raw(model: str) -> list[dict]:
    p = RAW / f"{model}.jsonl"
    if not p.exists():
        return []
    with p.open(encoding="utf-8") as fh:
        return [json.loads(l) for l in fh if l.strip()]


def mean(xs):
    xs = [x for x in xs if x is not None]
    return round(statistics.fmean(xs), 4) if xs else None


def rate(xs):
    xs = [x for x in xs if x is not None]
    return round(sum(xs) / len(xs), 4) if xs else None


# ------------------------------------------------------------------ metrics --
def free_text_metrics(records: list[dict]) -> dict:
    rows_by_probe = defaultdict(list)
    for r in records:
        if r["probe"] in ("ft_default", "ft_specified", "ft_technical",
                          "ft_repair", "ft_two_schools", "ft_pair"):
            body, _ = msq.strip_think(r.get("answer") or "")
            schools = msq.mentions_schools(body)
            rows_by_probe[r["probe"]].append(dict(
                cls=r["cls"],
                disclosure=msq.has_disclosure(body),
                schools=len(schools),
                clarify=msq.has_clarification(body),
                hedge=msq.has_hedge(body),
                refusal=msq.has_refusal(body),
                words=len(body.split()),
                logprob=r.get("mean_logprob"),
            ))

    out = {}
    for probe, rows in rows_by_probe.items():
        per_cls = defaultdict(list)
        for row in rows:
            per_cls[row["cls"]].append(row)

        def agg_of(rs):
            return dict(
                n=len(rs),
                disclosure_rate=rate([r["disclosure"] for r in rs]),
                school_attribution_rate=rate([r["schools"] >= 1 for r in rs]),
                plural_attribution_rate=rate([r["schools"] >= 2 for r in rs]),
                clarification_rate=rate([r["clarify"] for r in rs]),
                hedge_rate=rate([r["hedge"] for r in rs]),
                refusal_rate=rate([r["refusal"] for r in rs]),
                mean_words=mean([r["words"] for r in rs]),
                mean_logprob=mean([r["logprob"] for r in rs]),
            )

        agg = {cls: agg_of(rs) for cls, rs in per_cls.items()}
        agg["ALL"] = agg_of(rows)
        out[probe] = agg
    return out


def mcq_metrics(records: list[dict]) -> dict:
    out = {}
    for probe in ("mcq_neutral", "mcq_specified"):
        rs = [r for r in records if r["probe"] == probe]
        if not rs:
            continue
        per_cls = defaultdict(list)
        for r in rs:
            gold = r["gold_index"]
            arg = r["argmax_index"]
            divided = r.get("divided_index")
            if divided is None:
                divided = max((i for i, o in enumerate(r["options"])
                               if o.lower().startswith(("the scholars are divided",
                                                        "this is a contested matter",
                                                        "it depends on the jurisdiction"))),
                              default=None)
            per_cls[r["cls"]].append(dict(
                gold=gold, arg=arg, divided=divided,
                max_prob=r["max_prob"],
                entropy=msq.entropy(r["probs"]),
                probs=r["probs"],
                picked_divided=(divided is not None and arg == divided),
                divided_mass=(r["probs"][divided] if divided is not None else None),
                correct=(arg == gold),
                schools_labeled=any("(" in o and ")" in o for o in r["options"]),
            ))

        agg = {}
        for cls, rows in per_cls.items():
            contested = cls in ("A", "B", "C", "E")
            agg[cls] = dict(
                n=len(rows),
                plurality_recognition_rate=rate([r["picked_divided"] for r in rows])
                if contested or cls == "C" else None,
                plurality_mass=mean([r["divided_mass"] for r in rows])
                if contested or cls == "C" else None,
                school_adherence_rate=rate([r["correct"] for r in rows])
                if probe == "mcq_specified" else None,
                accuracy=rate([r["correct"] for r in rows])
                if cls == "D" else None,
                mean_max_prob=mean([r["max_prob"] for r in rows]),
                mean_entropy=mean([r["entropy"] for r in rows]),
                argmax_hist=defaultdict(int, {
                    k: sum(1 for r in rows if r["arg"] == k)
                    for k in set(r["arg"] for r in rows)}),
            )
            agg[cls]["argmax_hist"] = dict(sorted(agg[cls]["argmax_hist"].items()))
        agg["ALL"] = dict(
            n=sum(len(v) for v in per_cls.values()),
            plurality_recognition_rate=rate(
                [r["picked_divided"] for rows in per_cls.values() for r in rows]),
            plurality_mass=mean(
                [r["divided_mass"] for rows in per_cls.values() for r in rows]),
            school_adherence_rate=rate(
                [r["correct"] for rows in per_cls.values() for r in rows])
            if probe == "mcq_specified" else None,
            mean_max_prob=mean([r["max_prob"] for rows in per_cls.values() for r in rows]),
            mean_entropy=mean([r["entropy"] for rows in per_cls.values() for r in rows]),
        )
        out[probe] = agg
    return out


def positions_contained(records: list[dict], threshold: float = 0.04) -> dict:
    """How many of the competing positions does an answer actually carry?

    Scored in the model's own representation space (cosine similarity of the
    answer to each ground-truth position, computed at run time). We report a
    threshold-free version -- the gap between the best-matching position and
    the second-best -- because a large gap means the answer has collapsed onto
    a single position, and a small one means it carries two.
    """
    by_probe = defaultdict(list)
    for r in records:
        if not r["probe"].startswith("ft_"):
            continue
        sims = r.get("position_sims")
        if not sims:
            continue
        s = sorted(sims, reverse=True)
        top1 = s[0]
        top2 = s[1] if len(s) > 1 else 0.0
        by_probe[r["probe"]].append(dict(
            cls=r["cls"], n_positions=sum(1 for x in sims if x >= top1 - threshold),
            gap=round(top1 - top2, 4), top1=round(top1, 4), top2=round(top2, 4)))
    out = {}
    for probe, rows in by_probe.items():
        per_cls = defaultdict(list)
        for row in rows:
            per_cls[row["cls"]].append(row)

        def agg(rs):
            return dict(n=len(rs),
                        mean_positions_contained=mean([r["n_positions"] for r in rs]),
                        mean_top1=mean([r["top1"] for r in rs]),
                        mean_top2=mean([r["top2"] for r in rs]),
                        mean_gap=mean([r["gap"] for r in rs]))
        out[probe] = {"ALL": agg(rows),
                      **{cls: agg(rs) for cls, rs in per_cls.items()}}
    return out


def behaviour_matrix(records: list[dict]) -> dict:
    """The decisive 2x2: does the model have the default right, and does it
    respond to the stated context?

        default correct + switch correct  -> knows both, discloses neither
        default correct + switch wrong    -> has one answer and ignores context
        default wrong   + switch correct  -> knows the minority, will not volunteer
        default wrong   + switch wrong    -> knows neither
    """
    items = {i["id"]: i for i in msq.load_items()}
    neu = {r["item_id"]: r for r in records if r["probe"] == "mcq_neutral"}
    spec = {r["item_id"]: r for r in records if r["probe"] == "mcq_specified"}
    rows = []
    for iid, it in items.items():
        if iid not in neu or iid not in spec:
            continue
        n, s = neu[iid], spec[iid]
        # which ground-truth position did the neutral answer pick?
        default_positions = [p for p in it["positions"]
                             if set(p["schools"]) & set(it["majority_schools"])]
        default_txts = {p["txt"] for p in default_positions}
        picked = n["options"][n["argmax_index"]] if n["argmax_index"] < len(n["options"]) else ""
        default_correct = (picked in default_txts) if default_txts else None
        switch_correct = (s["argmax_index"] == s["gold_index"])
        rows.append(dict(item_id=iid, cls=it["cls"],
                         default_correct=default_correct,
                         switch_correct=switch_correct))
    per_cls = defaultdict(list)
    for row in rows:
        per_cls[row["cls"]].append(row)

    def classify(rs):
        c = Counter()
        for r in rs:
            if r["default_correct"] is None:
                continue
            if r["default_correct"] and r["switch_correct"]:
                c["both"] += 1
            elif r["default_correct"] and not r["switch_correct"]:
                c["default_only"] += 1
            elif not r["default_correct"] and r["switch_correct"]:
                c["switch_only"] += 1
            else:
                c["neither"] += 1
        n = sum(c.values())
        return dict(n=n, **{k: round(v / n, 3) for k, v in c.items()} if n else {})

    out = {cls: classify(rs) for cls, rs in per_cls.items()}
    out["ALL"] = classify(rows)
    out["_rows"] = rows
    return out


def context_response(records: list[dict]) -> dict:
    """How much does naming the context actually move the answer?

    The neutral and specified conditions present byte-identical options in the
    same order, so the two probability distributions are directly comparable.
    We report the total-variation distance between them, and the change in the
    probability mass on the position of the context that was named.

    TV = 0 means the instruction had *no effect whatsoever* on the answer.
    """
    items = {i["id"]: i for i in msq.load_items()}
    neu = {r["item_id"]: r for r in records if r["probe"] == "mcq_neutral"}
    spec = {r["item_id"]: r for r in records if r["probe"] == "mcq_specified"}
    rows = []
    for iid, it in items.items():
        if iid not in neu or iid not in spec:
            continue
        n, s = neu[iid], spec[iid]
        pn, ps = n["probs"], s["probs"]
        if len(pn) != len(ps):
            continue
        tv = 0.5 * sum(abs(a - b) for a, b in zip(pn, ps))
        # index of the named context's position among the options
        named_idx = None
        opts = msq._options_shared(it)
        for i, o in enumerate(opts):
            if it.get("minority_school") and it["minority_school"] in o["schools"]:
                named_idx = i
                break
        if named_idx is None:
            named_idx = None
        rows.append(dict(
            item_id=iid, cls=it["cls"],
            tv=round(tv, 4),
            argmax_unchanged=(n["argmax_index"] == s["argmax_index"]),
            p_named_before=round(pn[named_idx], 4) if named_idx is not None else None,
            p_named_after=round(ps[named_idx], 4) if named_idx is not None else None,
            delta=round(ps[named_idx] - pn[named_idx], 4) if named_idx is not None else None,
        ))
    per_cls = defaultdict(list)
    for row in rows:
        per_cls[row["cls"]].append(row)

    def agg(rs):
        ds = [r["delta"] for r in rs if r["delta"] is not None]
        return dict(
            n=len(rs),
            mean_context_response_tv=mean([r["tv"] for r in rs]),
            argmax_unchanged_rate=rate([r["argmax_unchanged"] for r in rs]),
            n_unchanged_distributions=sum(1 for r in rs if r["tv"] < 0.01),
            mean_delta_on_named_position=mean(ds),
            median_delta_on_named_position=(
                round(sorted(ds)[len(ds) // 2], 4) if ds else None),
        )

    out = {cls: agg(rs) for cls, rs in per_cls.items()}
    out["ALL"] = agg(rows)
    out["_rows"] = rows
    return out


def content_adherence(records: list[dict]) -> dict:
    """Does the model follow the *content* of the school it was told to follow?

    Uses the cosine similarity of the answer to each ground-truth position in
    the model's own representation space (computed at run time), because the
    model paraphrases heavily and a lexical overlap measure is far too brittle.
    """
    items = {i["id"]: i for i in msq.load_items()}
    rows = []
    for r in records:
        if r["probe"] != "ft_specified":
            continue
        it = items.get(r["item_id"])
        if not it or not it.get("minority_school"):
            continue
        sims = r.get("position_sims")
        if not sims or len(sims) != len(it["positions"]):
            continue
        spec_i = next((i for i, p in enumerate(it["positions"])
                       if it["minority_school"] in p["schools"]), None)
        if spec_i is None:
            continue
        spec_sim = sims[spec_i]
        other = max([s for i, s in enumerate(sims) if i != spec_i], default=0.0)
        default_i = next((i for i, p in enumerate(it["positions"])
                          if set(p["schools"]) & set(it["majority_schools"])), None)
        default_sim = sims[default_i] if default_i is not None else None
        body, _ = msq.strip_think(r.get("answer") or "")
        rows.append(dict(
            item_id=r["item_id"], cls=r["cls"],
            names_context=names(body, it["minority_school"]),
            sim_specified=round(spec_sim, 4),
            sim_best_other=round(other, 4),
            sim_default=round(default_sim, 4) if default_sim is not None else None,
            split_ratio=(round(default_sim / (spec_sim + default_sim), 4)
                         if default_sim is not None else None),
            adheres=spec_sim > other,
            label_adoption=(names(body, it["minority_school"]) and spec_sim <= other),
        ))
    per_cls = defaultdict(list)
    for row in rows:
        per_cls[row["cls"]].append(row)

    def agg(rs):
        return dict(
            n=len(rs),
            names_context_rate=rate([r["names_context"] for r in rs]),
            content_adherence_rate=rate([r["adheres"] for r in rs]),
            label_adoption_rate=rate([r["label_adoption"] for r in rs]),
            mean_sim_specified=mean([r["sim_specified"] for r in rs]),
            mean_sim_best_other=mean([r["sim_best_other"] for r in rs]),
            mean_split_ratio=mean([r["split_ratio"] for r in rs]),
        )

    out = {cls: agg(rs) for cls, rs in per_cls.items()}
    out["ALL"] = agg(rows)
    out["_rows"] = rows
    return out


def specified_outcome(records: list[dict]) -> dict:
    """Decompose what the model picks in the *specified* condition.

    Three possible moves: it picked the position of the school it was told to
    follow; it picked the 'the scholars are divided' option anyway; or it
    picked some other position.
    """
    out = {}
    for probe in ("mcq_neutral", "mcq_specified"):
        rs = [r for r in records if r["probe"] == probe]
        if not rs:
            continue
        per_cls = defaultdict(list)
        for r in rs:
            if r["cls"] == "D":
                continue
            div = r.get("divided_index")
            arg = r["argmax_index"]
            if probe == "mcq_specified":
                bucket = ("named_school" if arg == r["gold_index"]
                          else "divided_anyway" if div is not None and arg == div
                          else "other_position")
            else:
                bucket = ("divided" if div is not None and arg == div
                          else "a_position")
            per_cls[r["cls"]].append(bucket)
        res = {}
        for cls, buckets in per_cls.items():
            c = Counter(buckets)
            n = sum(c.values())
            res[cls] = dict(n=n, **{k: round(v / n, 3) for k, v in c.items()})
        allb = [b for bs in per_cls.values() for b in bs]
        c = Counter(allb)
        n = sum(c.values())
        res["ALL"] = dict(n=n, **{k: round(v / n, 3) for k, v in c.items()})
        out[probe] = res
    return out


def answer_change(records: list[dict]) -> dict:
    """How much does naming the school actually change the free-text answer?"""
    by_item = defaultdict(dict)
    for r in records:
        if r["probe"] in ("ft_default", "ft_specified"):
            body, _ = msq.strip_think(r["answer"] or "")
            by_item[r["item_id"]][r["probe"]] = body
    sims, changed = [], []
    for iid, d in by_item.items():
        if "ft_default" in d and "ft_specified" in d:
            s = difflib.SequenceMatcher(None, d["ft_default"], d["ft_specified"]).ratio()
            sims.append(s)
            changed.append(s < 0.60)
    return dict(n=len(sims), mean_similarity=mean(sims),
                changed_rate=rate(changed))


def model_report(model: str) -> dict:
    recs = load_raw(model)
    if not recs:
        return {}
    rep = {"model": model, "n_records": len(recs)}
    rep["mcq"] = mcq_metrics(recs)
    rep["free_text"] = free_text_metrics(recs)
    rep["answer_change"] = answer_change(recs)
    rep["content_adherence"] = content_adherence(recs)
    rep["positions_contained"] = positions_contained(recs)
    rep["behaviour_matrix"] = behaviour_matrix(recs)
    rep["specified_outcome"] = specified_outcome(recs)
    rep["context_response"] = context_response(recs)
    n = rep["mcq"].get("mcq_neutral", {}).get("ALL", {})
    s = rep["mcq"].get("mcq_specified", {}).get("ALL", {})
    rep["headline"] = {
        "PRR": n.get("plurality_recognition_rate"),
        "PRR_mass": n.get("plurality_mass"),
        "MSQ": s.get("school_adherence_rate"),
        "PRR_A": rep["mcq"].get("mcq_neutral", {}).get("A", {}).get("plurality_recognition_rate"),
        "PRR_B": rep["mcq"].get("mcq_neutral", {}).get("B", {}).get("plurality_recognition_rate"),
        "PRR_C": rep["mcq"].get("mcq_neutral", {}).get("C", {}).get("plurality_recognition_rate"),
        "PRR_E": rep["mcq"].get("mcq_neutral", {}).get("E", {}).get("plurality_recognition_rate"),
        "MSQ_A": rep["mcq"].get("mcq_specified", {}).get("A", {}).get("school_adherence_rate"),
        "MSQ_C": rep["mcq"].get("mcq_specified", {}).get("C", {}).get("school_adherence_rate"),
        "MSQ_E": rep["mcq"].get("mcq_specified", {}).get("E", {}).get("school_adherence_rate"),
        "acc_D": rep["mcq"].get("mcq_neutral", {}).get("D", {}).get("accuracy"),
        "conf_contested": n.get("mean_max_prob"),
        "conf_control": rep["mcq"].get("mcq_neutral", {}).get("D", {}).get("mean_max_prob"),
        "disc_default": rep["free_text"].get("ft_default", {}).get("ALL", {}).get("disclosure_rate"),
        "disc_technical": rep["free_text"].get("ft_technical", {}).get("ALL", {}).get("disclosure_rate"),
        "plural_attr_two_schools": rep["free_text"].get("ft_two_schools", {}).get("ALL", {}).get("plural_attribution_rate"),
        "names_context_rate": rep["content_adherence"].get("ALL", {}).get("names_context_rate"),
        "content_adherence_rate": rep["content_adherence"].get("ALL", {}).get("content_adherence_rate"),
        "label_adoption_rate": rep["content_adherence"].get("ALL", {}).get("label_adoption_rate"),
        "positions_default": rep["positions_contained"].get("ft_default", {}).get("ALL", {}).get("mean_positions_contained"),
        "positions_two_schools": rep["positions_contained"].get("ft_two_schools", {}).get("ALL", {}).get("mean_positions_contained"),
        "positions_pair": rep["positions_contained"].get("ft_pair", {}).get("ALL", {}).get("mean_positions_contained"),
        "positions_repair": rep["positions_contained"].get("ft_repair", {}).get("ALL", {}).get("mean_positions_contained"),
        "disc_repair": rep["free_text"].get("ft_repair", {}).get("ALL", {}).get("disclosure_rate"),
        "bm_both": rep["behaviour_matrix"].get("ALL", {}).get("both"),
        "bm_default_only": rep["behaviour_matrix"].get("ALL", {}).get("default_only"),
        "bm_switch_only": rep["behaviour_matrix"].get("ALL", {}).get("switch_only"),
        "bm_neither": rep["behaviour_matrix"].get("ALL", {}).get("neither"),
    }
    h = rep["headline"]
    if h["PRR"] is not None and h["MSQ"] is not None:
        h["blind_spot_gap"] = round(h["MSQ"] - h["PRR"], 4)
    return rep


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--models", nargs="*", default=None)
    args = ap.parse_args()
    models = args.models or sorted(p.stem for p in RAW.glob("*.jsonl"))
    reports = {}
    for m in models:
        r = model_report(m)
        if r:
            reports[m] = r
            print(f"{m}: PRR={r['headline'].get('PRR')} MSQ={r['headline'].get('MSQ')} "
                  f"gap={r['headline'].get('blind_spot_gap')} "
                  f"disc={r['headline'].get('disc_default')}")
    OUTT.mkdir(parents=True, exist_ok=True)
    (OUTT / "metrics.json").write_text(json.dumps(reports, indent=2, ensure_ascii=False))
    print(f"\nwrote {OUTT / 'metrics.json'}")


if __name__ == "__main__":
    main()