"""Summarise the EM panel: how often each method commits to a spillover claim, and how often
it gets the direction right, split by broad vs narrow twin.

Two quantities, kept separate on purpose, because collapsing them is what made the panel hard
to read:

  commit rate         fraction of elicitations the judge could score at all. The binary rubric
                      returns "unclear" unless the output explicitly addresses out-of-domain
                      behaviour, so this measures whether a method says anything falsifiable.
  conditional accuracy of the elicitations it did commit on, fraction with the right direction.

A method can look bad two very different ways: never committing (low commit rate, LoRAcle) or
committing confidently in the wrong direction (high commit rate, zero accuracy - AO and IA on
narrow twins). The published pooled number mixes those.

Error bars are organism-clustered bootstrap, matching the convention in
scripts/plot_method_bench_scorer.py: organisms are the resampling unit, not elicitations.
"""
import argparse, os, sys
from collections import defaultdict
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

SQL = """
SELECT p.lora_id, mr.method, sc.scorer_model,
       (s.scorer_score::text <> 'NaN') AS committed,
       CASE WHEN s.scorer_score::text <> 'NaN' THEN s.scorer_score ELSE NULL END AS score
FROM scores s
JOIN scorer sc ON sc.scorer_id = s.scorer_id
JOIN predictions p ON p.pred_id = s.pred_id
JOIN method_runs mr ON mr.method_run_id = p.method_run_id
JOIN qa_pairs q ON q.qa_id = p.qa_id
WHERE q.category = 'em'
  AND p.sub_prompt_kind LIKE 'narrow%%'
  AND sc.template_path LIKE '%%scorer_binary_yesno.txt'
  AND s.scorer_id = %s
"""


def latest_binary_scorer(cur):
    cur.execute("""SELECT sc.scorer_id, sc.scorer_model FROM scorer sc
                   WHERE sc.template_path LIKE '%%scorer_binary_yesno.txt'
                     AND EXISTS (SELECT 1 FROM scores s WHERE s.scorer_id = sc.scorer_id)
                   ORDER BY sc.created_at DESC LIMIT 1""")
    return cur.fetchone()


def twin_of(lora_id):
    lid = lora_id.lower()
    if "narrow" in lid:
        return "narrow"
    if "broad" in lid:
        return "broad"
    return "other"


def boot(per_org, n=2000, seed=0):
    """Organism-clustered bootstrap of a mean over per-organism (hits, n) pairs."""
    per_org = [(h, k) for h, k in per_org if k]
    if not per_org:
        return None, None
    rates = np.array([h / k for h, k in per_org], dtype=float)
    rng = np.random.default_rng(seed)
    draws = [rates[rng.integers(0, len(rates), len(rates))].mean() for _ in range(n)]
    return float(rates.mean()), float(np.std(draws))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--db", default=os.environ.get("METHOD_BENCH_DB_URL", ""))
    ap.add_argument("--scorer-id", type=int, default=None, help="default: latest binary scorer")
    ap.add_argument("--out", default="/workspace-vast/jbauer/em_organisms/eval/em_panel.png")
    args = ap.parse_args()
    if not args.db:
        raise SystemExit("METHOD_BENCH_DB_URL not set")

    import psycopg
    with psycopg.connect(args.db, connect_timeout=15) as conn:
        cur = conn.cursor()
        sid, model = (args.scorer_id, "(given)") if args.scorer_id else latest_binary_scorer(cur)
        cur.execute(SQL, (sid,))
        rows = cur.fetchall()
    print(f"[panel] scorer_id={sid} model={model}  rows={len(rows)}")
    if not rows:
        raise SystemExit("no scored EM rows yet")

    # (method, twin, organism) -> [n, n_committed, n_correct]
    cell = defaultdict(lambda: [0, 0, 0])
    for lora_id, method, _sm, committed, score in rows:
        c = cell[(method, twin_of(lora_id), lora_id)]
        c[0] += 1
        if committed:
            c[1] += 1
            c[2] += int(score == 1.0)

    methods = sorted({k[0] for k in cell})
    twins = ["broad", "narrow"]

    print(f"\n{'method':9s} {'twin':7s} {'orgs':>5s} {'elicit':>7s} {'commit':>7s} "
          f"{'commit%':>8s} {'correct':>8s} {'cond.acc':>9s}")
    print("-" * 70)
    agg = {}
    for meth in methods:
        for tw in twins:
            cells = [(k, v) for k, v in cell.items() if k[0] == meth and k[1] == tw]
            if not cells:
                continue
            n = sum(v[0] for _, v in cells)
            nc = sum(v[1] for _, v in cells)
            ncorr = sum(v[2] for _, v in cells)
            commit_rate, commit_se = boot([(v[1], v[0]) for _, v in cells])
            acc, acc_se = boot([(v[2], v[1]) for _, v in cells])
            agg[(meth, tw)] = (commit_rate, commit_se, acc, acc_se)
            accs = "     n/a" if acc is None else f"{100 * acc:7.1f}%"
            print(f"{meth:9s} {tw:7s} {len(cells):5d} {n:7d} {nc:7d} "
                  f"{100 * commit_rate:7.1f}% {ncorr:8d} {accs}")

    print("\nper-organism detail")
    for (meth, tw, lid), v in sorted(cell.items()):
        acc = f"{100 * v[2] / v[1]:.0f}%" if v[1] else "n/a"
        print(f"  {meth:9s} {tw:7s} {lid.replace('@latest', ''):46s} "
              f"commit {v[1]:3d}/{v[0]:3d}  correct {v[2]:3d}  acc {acc}")

    fig, axes = plt.subplots(1, 2, figsize=(11, 4.5))
    x = np.arange(len(methods))
    width = 0.36
    for ax, idx, title, ylab in [
            (axes[0], 0, "Does the method commit to a spillover claim?", "commit rate"),
            (axes[1], 2, "When it commits, is the direction right?", "conditional accuracy")]:
        for off, tw in ((-width / 2, "broad"), (width / 2, "narrow")):
            vals = [(agg.get((m, tw)) or (None,) * 4)[idx] for m in methods]
            errs = [(agg.get((m, tw)) or (None,) * 4)[idx + 1] for m in methods]
            ax.bar(x + off, [np.nan if v is None else 100 * v for v in vals], width,
                   yerr=[0 if e is None else 100 * e for e in errs], capsize=3, label=f"{tw} twin")
            # A bar at exactly 0 is invisible, and "0% correct" is the whole finding on the
            # narrow twins -- label it so the panel cannot be read as one bar at 100%.
            for xi, m in enumerate(methods):
                v = (agg.get((m, tw)) or (None,) * 4)[idx]
                if v is None:
                    ax.annotate("no commits", (xi + off, 2), ha="center", va="bottom",
                                rotation=90, fontsize="small", color="gray")
                elif v == 0:
                    ax.annotate("0%", (xi + off, 1), ha="center", va="bottom", fontsize="small")
        ax.set_xticks(x)
        ax.set_xticklabels(methods)
        ax.set_ylabel(f"{ylab} (%)")
        ax.set_title(title)
        ax.grid(axis="y", alpha=0.3)
        ax.set_ylim(0, 105)
        ax.legend()
    fig.tight_layout()
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(out, dpi=200, bbox_inches="tight")
    print(f"\n{out}")


if __name__ == "__main__":
    main()