#!/usr/bin/env python3
"""Score the LLM classifier against human gold-standard codings.

Reads:
  - data/gold_standard.jsonl (one JSON object per line, written from the
    Classify admin UI; each row is one reviewer's coding of one bill)
  - data/classifications.json (from data_updating_scripts/classify_bills.py)

Writes:
  - data/eval/<date>_metrics.json — three axes of metrics:
    (1) is_ai_related — binary axis: was the bill actually about AI?
        Reported as precision/recall/F1 for the "AI" positive class plus
        Cohen's κ across reviewer pairs.
    (2) regulatory_targets — per-label + macro-F1 + κ. Computed only on
        bills where the human majority verdict is is_ai_related = True.
    (3) governance_mechanisms — same treatment as (2).

The is_ai_related gate matters: comparing labels on a bill that both
reviewers marked non-AI is meaningless (both are trivially empty).
Bills flagged non-AI drop out of the label-agreement math but stay in
the binary AI-vs-not-AI accuracy.

Usage:
  python scripts/evaluate_classifier.py
  python scripts/evaluate_classifier.py --gold data/gold_standard.jsonl \
      --classifications data/classifications.json --out data/eval

The script does not touch the network. Safe to re-run on every codebook edit.
"""

from __future__ import annotations

import argparse
import json
import statistics
from collections import defaultdict
from datetime import date
from pathlib import Path
from typing import Any


def load_gold(path: Path) -> dict[str, list[dict[str, Any]]]:
    """{bill_key: [coding, coding, ...]} — each coding is one reviewer's row."""
    out: dict[str, list[dict[str, Any]]] = defaultdict(list)
    if not path.exists():
        return out
    with path.open("r", encoding="utf-8") as fp:
        for line in fp:
            line = line.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError:
                continue
            key = row.get("bill_key")
            if not key:
                continue
            # Backward compat: accept both new ("reviewer") and old ("coder")
            # field names so older gold_standard.jsonl files still work.
            if "reviewer" not in row and "coder" in row:
                row["reviewer"] = row["coder"]
            out[key].append(row)
    return out


def load_predictions(path: Path) -> dict[str, dict[str, Any]]:
    if not path.exists():
        return {}
    with path.open("r", encoding="utf-8") as fp:
        data = json.load(fp)
    return data if isinstance(data, dict) else {}


def _prf(tp: int, fp: int, fn: int) -> tuple[float, float, float]:
    p = tp / (tp + fp) if (tp + fp) else 0.0
    r = tp / (tp + fn) if (tp + fn) else 0.0
    f = (2 * p * r / (p + r)) if (p + r) else 0.0
    return p, r, f


def _gold_ai_verdict(codings: list[dict[str, Any]]) -> bool | None:
    """Majority vote of reviewers on is_ai_related. Ignores rows where the
    field is null (legacy entries). Returns None if no reviewers voted."""
    votes = [c.get("is_ai_related") for c in codings if isinstance(c.get("is_ai_related"), bool)]
    if not votes:
        return None
    yes = sum(1 for v in votes if v)
    # Majority (ties → True, since a coder who says AI is asserting positive info)
    return yes >= (len(votes) + 1) // 2


def _pred_ai_verdict(entry: dict[str, Any]) -> bool | None:
    """Extract is_ai_related from a classifier entry, falling back to the
    labels-present heuristic for legacy classifications that predate the field."""
    raw = entry.get("is_ai_related")
    if isinstance(raw, bool):
        return raw
    return bool(entry.get("regulatory_targets") or entry.get("governance_mechanisms"))


def is_ai_related_metrics(
    gold_by_bill: dict[str, list[dict[str, Any]]],
    predictions: dict[str, dict[str, Any]],
) -> dict[str, Any]:
    """Binary axis: was the bill actually AI-related? Positive class = AI."""
    tp = fp = fn = tn = 0
    counted = 0
    for bill_key, codings in gold_by_bill.items():
        if bill_key not in predictions:
            continue
        gold = _gold_ai_verdict(codings)
        pred = _pred_ai_verdict(predictions[bill_key])
        if gold is None or pred is None:
            continue
        counted += 1
        if gold and pred:
            tp += 1
        elif pred and not gold:
            fp += 1
        elif gold and not pred:
            fn += 1
        else:
            tn += 1
    p, r, f = _prf(tp, fp, fn)
    accuracy = (tp + tn) / counted if counted else 0.0
    # Kappa across reviewer pairs on the binary axis.
    kappa = _binary_kappa(gold_by_bill, "is_ai_related")
    return {
        "positive_class": "is_ai_related=true",
        "precision": round(p, 3),
        "recall": round(r, 3),
        "f1": round(f, 3),
        "accuracy": round(accuracy, 3),
        "tp": tp, "fp": fp, "fn": fn, "tn": tn,
        "evaluable_bills": counted,
        "inter_rater_kappa": kappa,
    }


def _binary_kappa(
    gold_by_bill: dict[str, list[dict[str, Any]]],
    field: str,
) -> dict[str, Any] | None:
    """Cohen's κ on a binary field across the first two reviewers per bill."""
    pairs: list[tuple[bool, bool]] = []
    for codings in gold_by_bill.values():
        typed = [c.get(field) for c in codings if isinstance(c.get(field), bool)]
        if len(typed) < 2:
            continue
        pairs.append((typed[0], typed[1]))
    if not pairs:
        return None
    n = len(pairs)
    agree = sum(1 for a, b in pairs if a == b)
    po = agree / n
    p_a_yes = sum(1 for a, _ in pairs if a) / n
    p_b_yes = sum(1 for _, b in pairs if b) / n
    pe = p_a_yes * p_b_yes + (1 - p_a_yes) * (1 - p_b_yes)
    kappa = (po - pe) / (1 - pe) if (1 - pe) else 0.0
    return {
        "cohens_kappa": round(kappa, 3),
        "observed_agreement": round(po, 3),
        "chance_agreement": round(pe, 3),
        "matched_bills": n,
    }


def per_label_metrics(
    gold_by_bill: dict[str, list[dict[str, Any]]],
    predictions: dict[str, dict[str, Any]],
    axis_key: str,
    ai_only: bool = True,
) -> dict[str, dict[str, float]]:
    """For each label, tally TP/FP/FN comparing predictions against the
    majority-vote of reviewers. When ai_only=True, restricts to bills where
    the majority gold verdict is is_ai_related=True — comparing labels on
    non-AI bills is meaningless (both sides are trivially empty)."""
    per_label_counts: dict[str, dict[str, int]] = defaultdict(lambda: {"tp": 0, "fp": 0, "fn": 0})
    for bill_key, codings in gold_by_bill.items():
        if bill_key not in predictions:
            continue
        if ai_only and _gold_ai_verdict(codings) is False:
            continue
        pred_labels = set(predictions[bill_key].get(axis_key) or [])
        vote_counts: dict[str, int] = defaultdict(int)
        for c in codings:
            for lbl in c.get(axis_key) or []:
                vote_counts[lbl] += 1
        threshold = (len(codings) + 1) // 2  # majority (ties → yes)
        gold_labels = {lbl for lbl, n in vote_counts.items() if n >= threshold}
        for lbl in gold_labels | pred_labels:
            if lbl in gold_labels and lbl in pred_labels:
                per_label_counts[lbl]["tp"] += 1
            elif lbl in pred_labels:
                per_label_counts[lbl]["fp"] += 1
            else:
                per_label_counts[lbl]["fn"] += 1

    out: dict[str, dict[str, float]] = {}
    for lbl, cnt in per_label_counts.items():
        p, r, f = _prf(cnt["tp"], cnt["fp"], cnt["fn"])
        out[lbl] = {
            "precision": round(p, 3),
            "recall": round(r, 3),
            "f1": round(f, 3),
            "tp": cnt["tp"],
            "fp": cnt["fp"],
            "fn": cnt["fn"],
        }
    return out


def macro_f1(per_label: dict[str, dict[str, float]]) -> float:
    if not per_label:
        return 0.0
    return round(statistics.fmean(v["f1"] for v in per_label.values()), 3)


def multilabel_kappa(
    gold_by_bill: dict[str, list[dict[str, Any]]],
    axis_key: str,
    all_labels: set[str],
    ai_only: bool = True,
) -> dict[str, Any] | None:
    """Multi-label κ across the first two reviewers, computed only on bills
    where both reviewers agreed the bill is AI-related (labels on non-AI
    bills are trivially empty and would inflate agreement)."""
    per_label_yes_a: dict[str, int] = defaultdict(int)
    per_label_yes_b: dict[str, int] = defaultdict(int)
    matched_bills = 0
    agree_total = 0
    total_cells = 0
    for bill_key, codings in gold_by_bill.items():
        if len(codings) < 2:
            continue
        if ai_only:
            # Both reviewers must have voted True on is_ai_related for the
            # label agreement to be meaningful.
            va = codings[0].get("is_ai_related")
            vb = codings[1].get("is_ai_related")
            if not (va is True and vb is True):
                continue
        matched_bills += 1
        a = set(codings[0].get(axis_key) or [])
        b = set(codings[1].get(axis_key) or [])
        for lbl in all_labels:
            in_a = lbl in a
            in_b = lbl in b
            per_label_yes_a[lbl] += 1 if in_a else 0
            per_label_yes_b[lbl] += 1 if in_b else 0
            if in_a == in_b:
                agree_total += 1
            total_cells += 1
    if not matched_bills:
        return None
    po = agree_total / total_cells if total_cells else 0.0
    pes: list[float] = []
    for lbl in all_labels:
        pa_y = per_label_yes_a[lbl] / matched_bills
        pb_y = per_label_yes_b[lbl] / matched_bills
        pes.append(pa_y * pb_y + (1 - pa_y) * (1 - pb_y))
    pe = statistics.fmean(pes) if pes else 0.0
    kappa = (po - pe) / (1 - pe) if (1 - pe) else 0.0
    return {
        "cohens_kappa": round(kappa, 3),
        "observed_agreement": round(po, 3),
        "chance_agreement": round(pe, 3),
        "matched_bills": matched_bills,
        "scope": "bills where both reviewers said is_ai_related=True",
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--gold", type=Path, default=Path("data/gold_standard.jsonl"))
    parser.add_argument("--classifications", type=Path, default=Path("data/classifications.json"))
    parser.add_argument("--out", type=Path, default=Path("data/eval"))
    args = parser.parse_args()

    gold_by_bill = load_gold(args.gold)
    predictions = load_predictions(args.classifications)

    print(f"Gold standard: {sum(len(v) for v in gold_by_bill.values())} codings across {len(gold_by_bill)} bills")
    print(f"LLM predictions: {len(predictions)} bills")
    overlap = [k for k in gold_by_bill if k in predictions]
    print(f"Overlap (evaluable): {len(overlap)} bills")

    if not overlap:
        print("No overlap — skipping metric computation.")
        return

    # Axis 0: is_ai_related (binary)
    ai_metrics = is_ai_related_metrics(gold_by_bill, predictions)

    # Axes 1 & 2: label-level (restricted to bills where gold says AI)
    targets_per_label = per_label_metrics(gold_by_bill, predictions, "regulatory_targets", ai_only=True)
    mechanisms_per_label = per_label_metrics(gold_by_bill, predictions, "governance_mechanisms", ai_only=True)

    # Full label vocab (from predictions + gold) — needed for kappa marginals.
    all_targets: set[str] = set()
    all_mechs: set[str] = set()
    for p in predictions.values():
        all_targets.update(p.get("regulatory_targets") or [])
        all_mechs.update(p.get("governance_mechanisms") or [])
    for codings in gold_by_bill.values():
        for c in codings:
            all_targets.update(c.get("regulatory_targets") or [])
            all_mechs.update(c.get("governance_mechanisms") or [])

    kappa_targets = multilabel_kappa(gold_by_bill, "regulatory_targets", all_targets, ai_only=True)
    kappa_mechs = multilabel_kappa(gold_by_bill, "governance_mechanisms", all_mechs, ai_only=True)

    result = {
        "date": date.today().isoformat(),
        "counts": {
            "gold_codings": sum(len(v) for v in gold_by_bill.values()),
            "gold_bills": len(gold_by_bill),
            "predicted_bills": len(predictions),
            "overlap_bills": len(overlap),
        },
        "is_ai_related": ai_metrics,
        "regulatory_targets": {
            "per_label": targets_per_label,
            "macro_f1": macro_f1(targets_per_label),
            "inter_rater_kappa": kappa_targets,
        },
        "governance_mechanisms": {
            "per_label": mechanisms_per_label,
            "macro_f1": macro_f1(mechanisms_per_label),
            "inter_rater_kappa": kappa_mechs,
        },
    }

    args.out.mkdir(parents=True, exist_ok=True)
    out_path = args.out / f"{date.today().isoformat()}_metrics.json"
    with out_path.open("w", encoding="utf-8") as fp:
        json.dump(result, fp, indent=2, ensure_ascii=False)
    print(f"\nis_ai_related F1:                  {ai_metrics['f1']} (accuracy {ai_metrics['accuracy']}, {ai_metrics['evaluable_bills']} bills)")
    print(f"Macro-F1 (Regulatory Targets):     {result['regulatory_targets']['macro_f1']}")
    print(f"Macro-F1 (Governance Mechanisms): {result['governance_mechanisms']['macro_f1']}")
    if ai_metrics.get("inter_rater_kappa"):
        print(f"Cohen's κ (is_ai_related):         {ai_metrics['inter_rater_kappa']['cohens_kappa']} ({ai_metrics['inter_rater_kappa']['matched_bills']} doubly-coded)")
    if kappa_targets:
        print(f"Cohen's κ (Regulatory Targets):    {kappa_targets['cohens_kappa']} ({kappa_targets['matched_bills']} AI-eligible doubly-coded bills)")
    if kappa_mechs:
        print(f"Cohen's κ (Governance Mechanisms): {kappa_mechs['cohens_kappa']}")
    print(f"\nMetrics written to {out_path}")


if __name__ == "__main__":
    main()