plot_motivation.py

"""Motivation figure: per difficulty bin (K = #pre-debate-correct agents),
fraction of questions where the top-ranked agent under each signal
(min-LL / ppl / self-conf / SVR) is correct."""
import json
import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

sys.path.insert(0, str(Path(__file__).parent))
from data_loader import non_unanimous_qids

ROOT = Path(__file__).parent
FIG_DIR = ROOT / "figures"
FIG_DIR.mkdir(parents=True, exist_ok=True)

N_AGENTS = 6
BINS = [
    ("K=1",   (1,)),
    ("K=2",   (2,)),
    ("K=3",   (3,)),
    ("K=4-5", (4, 5)),
]
SIGNALS = ["min-LL", "ppl", "self-conf", "SVR"]

CONFIGS = [
    {"llm": "GPT-OSS-120B", "task": "math",
     "predebate":  ROOT / "runs/imobench/gpt-oss-120b/predebate_n6",
     "exhaust_d1": ROOT / "runs/imobench/gpt-oss-120b/exhaustive_d1_n6"},
    {"llm": "DeepSeek-V3.1", "task": "math",
     "predebate":  ROOT / "runs/imobench/deepseek-v3.1/predebate_n6",
     "exhaust_d1": ROOT / "runs/imobench/deepseek-v3.1/exhaustive_d1_n6"},
]


def _pre_canonical(ag):
    sr = ag.get("self_round") or {}
    return sr.get("parsed_math_answer") or sr.get("boxed_answer")


def _signal_scores_for_query(doc):
    """{signal -> {agent_id -> score}} for one question."""
    canon_gold = doc.get("parsed_ground_truth_answer") or doc.get("ground_truth_answer")
    scores = {sig: {} for sig in SIGNALS}
    pre_canon = {ag["agent_id"]: _pre_canonical(ag) for ag in doc["agents"]}
    for ag in doc["agents"]:
        A = ag["agent_id"]
        own_canon = pre_canon.get(A)
        if own_canon is None:
            continue  # exclude agents with no boxed answer from all signals
        sr = ag.get("self_round") or {}
        lp = sr.get("logprob_stats") or {}
        if lp.get("min_logprob") is not None:
            # min_logprob is negative; argmax picks the least-negative
            # (most-confident) agent. Use the raw value as the score.
            scores["min-LL"][A] = float(lp["min_logprob"])
        if lp.get("perplexity") is not None:
            scores["ppl"][A] = -float(lp["perplexity"])
        if sr.get("confidence") is not None:
            scores["self-conf"][A] = float(sr["confidence"])
        retain = flip = debates = 0
        for pr in ag.get("probes", []):
            p = pr["peer_id"]
            peer_canon = pre_canon.get(p)
            if peer_canon is None or peer_canon == own_canon:
                continue
            post = pr.get("post_debate_snapshot") or {}
            post_canon = post.get("parsed_math_answer") or post.get("boxed_answer")
            if post_canon is None:
                continue
            debates += 1
            if post_canon == own_canon:
                retain += 1
            elif post_canon == peer_canon:
                flip += 1
        if debates > 0:
            scores["SVR"][A] = (retain - flip) / debates
    return scores, canon_gold, pre_canon


def _split_tie_argmax(score_map, correct_map):
    if not score_map:
        return 0.0
    top = max(score_map.values())
    tied = [a for a, s in score_map.items() if s == top]
    return sum(int(correct_map.get(a, False)) for a in tied) / len(tied)


def _per_query_accs(doc):
    scores, canon_gold, pre_canon = _signal_scores_for_query(doc)
    correct_map = {a: (c is not None and c == canon_gold) for a, c in pre_canon.items()}
    K = sum(int(v) for v in correct_map.values())
    return {sig: _split_tie_argmax(scores[sig], correct_map) for sig in SIGNALS}, K


def compute_per_llm(cfg):
    subset = non_unanimous_qids(cfg["predebate"], cfg["task"])
    accs, qK = {}, {}
    for qid in sorted(subset):
        f = cfg["exhaust_d1"] / f"{qid}.json"
        if not f.exists():
            continue
        doc = json.loads(f.read_text())
        a, k = _per_query_accs(doc)
        accs[qid] = a
        qK[qid] = k
    rows = []
    for label, kgroup in BINS:
        sub = [qid for qid in accs if qK[qid] in kgroup]
        n = len(sub)
        if n == 0:
            rows.append({"Bin": label, "n": 0, **{s: 0.0 for s in SIGNALS}})
            continue
        means = {sig: sum(accs[q][sig] for q in sub) / n for sig in SIGNALS}
        rows.append({"Bin": label, "n": n, **means})
    return rows


def main():
    per_llm = []
    for cfg in CONFIGS:
        print(f"computing {cfg['llm']}...")
        rows = compute_per_llm(cfg)
        for r in rows:
            sig_str = "  ".join(f"{s}={r[s]:.3f}" for s in SIGNALS)
            print(f"  {r['Bin']:<6} n={r['n']:>4}  {sig_str}")
        per_llm.append((cfg["llm"], rows))

    PRIOR_STYLES = [
        ("min-LL",    "min-LL", "o", "C0", "--"),
        ("ppl",       "PPL",    "s", "C1", "--"),
        ("self-conf", "conf",   "^", "C2", "--"),
    ]
    SVR_STYLE = ("SVR", "SVR", "D", "C3", "-")
    BIN_ORDER = ["K=4-5", "K=3", "K=2", "K=1"]
    BIN_LABELS = {"K=1": "1\n(hard)", "K=2": "2", "K=3": "3", "K=4-5": "4+\n(easy)"}

    fig, axs = plt.subplots(nrows=1, ncols=len(per_llm), figsize=(5, 3.5), sharey=True)
    for ax, (llm, rows) in zip(axs, per_llm):
        bin_to_row = {r["Bin"]: r for r in rows}
        ordered = [bin_to_row[b] for b in BIN_ORDER]
        x = list(range(len(ordered)))
        for col, label, marker, color, ls in PRIOR_STYLES:
            ax.plot(x, [r[col] * 100 for r in ordered], label=label, marker=marker,
                    markersize=4, linestyle=ls, color=color, markerfacecolor="white",
                    markeredgecolor=color, markeredgewidth=1.2)
        col, label, marker, color, ls = SVR_STYLE
        ax.plot(x, [r[col] * 100 for r in ordered], label=label, marker=marker,
                markersize=4, linestyle=ls, color=color, markerfacecolor="white",
                markeredgecolor=color, markeredgewidth=1.2)
        ax.set_xticks(x, labels=[BIN_LABELS[r["Bin"]] for r in ordered])
        ax.set_yticks(np.arange(0, 101, 25))
        ax.set_ylim(0, 100)
        ax.grid(True, linestyle="--", alpha=0.7)
        ax.tick_params(axis="y", which="major", labelsize=11)
        ax.tick_params(axis="x", which="major", labelsize=11)
        ax.spines["top"].set_visible(False)
        ax.spines["right"].set_visible(False)
        ax.set_title(llm, fontsize=12)

    axs[0].set_ylabel("% of Problems", fontsize=12)
    fig.supxlabel("Difficulty (# pre-debate correct agents)", fontsize=12, y=0.25)
    handles, labels = axs[0].get_legend_handles_labels()
    fig.tight_layout(rect=[0, 0.20, 1, 1], w_pad=0.8)
    fig.legend(handles, labels, loc="upper center",
               bbox_to_anchor=(0.5, 1.08), ncol=4, fontsize=12, frameon=False,
               borderpad=0, columnspacing=1.0, handletextpad=0.3)

    for ext in ("pdf", "png"):
        out = FIG_DIR / f"motivation.{ext}"
        fig.savefig(out, dpi=300, bbox_inches="tight", pad_inches=0)
        print(f"\nsaved {out}")


if __name__ == "__main__":
    main()