"""
Compute validation-study metrics from the three labeler outputs.

Per DESIGN.md §5: pairwise Cohen's kappa, per-frame F1 of detector
against majority of (curator, LLM-judge), macro-F1, confusion matrix,
per-stratum breakdown, inter-labeler disagreement patterns.

Writes results/results.json and prints a human-readable summary.

Scope note: the detector (Labeler C) can only emit labels for the 11
text-side FVS frames with suggest_frames rules. Comparison metrics
restrict to those 11 frames; meta-side frames (FVS-003/004/005/006/013/
018/019/020) are preserved in A and B labels for reference but not
included in A-vs-C or B-vs-C or majority-vs-C agreement.
"""

from __future__ import annotations

import json
from collections import Counter, defaultdict
from pathlib import Path
from statistics import mean

REPO_ROOT = Path(__file__).resolve().parents[2]
STUDY_DIR = REPO_ROOT / "fvs_eval" / "validation_study"
CORPUS_DIR = STUDY_DIR / "corpus"
LABELS_DIR = STUDY_DIR / "labels"
RESULTS_DIR = STUDY_DIR / "results"
RESULTS_DIR.mkdir(exist_ok=True)


FRAME_IDS_ALL = [f"FVS-{i:03d}" for i in range(1, 21)]

# Text-side frames the detector can emit (the 11-frame subset)
DETECTABLE_FRAMES = [
    "FVS-001", "FVS-002", "FVS-007", "FVS-008", "FVS-009",
    "FVS-010", "FVS-011", "FVS-012", "FVS-014", "FVS-015", "FVS-016",
]


def _load(path: Path) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def _binary_labels(labels_obj: dict, doc_id: str, frame_ids: list) -> list:
    """Return binary [0/1] list of labels for one doc over frame_ids."""
    doc_labels = labels_obj["labels"].get(doc_id, {})
    out = []
    for fid in frame_ids:
        entry = doc_labels.get(fid, {})
        out.append(1 if entry.get("exhibits") else 0)
    return out


def _cohen_kappa(a: list, b: list) -> float:
    """Cohen's kappa on binary labels."""
    n = len(a)
    if n == 0:
        return 0.0
    # Observed agreement
    agree = sum(1 for x, y in zip(a, b) if x == y) / n
    # Expected agreement by chance
    pa1 = sum(a) / n
    pb1 = sum(b) / n
    pe = pa1 * pb1 + (1 - pa1) * (1 - pb1)
    if pe == 1.0:
        return 1.0 if agree == 1.0 else 0.0
    return (agree - pe) / (1 - pe)


def _f1(pred: list, ref: list) -> tuple:
    """F1, precision, recall on binary labels. pred vs ref."""
    tp = sum(1 for p, r in zip(pred, ref) if p == 1 and r == 1)
    fp = sum(1 for p, r in zip(pred, ref) if p == 1 and r == 0)
    fn = sum(1 for p, r in zip(pred, ref) if p == 0 and r == 1)
    tn = sum(1 for p, r in zip(pred, ref) if p == 0 and r == 0)
    prec = tp / (tp + fp) if (tp + fp) > 0 else 0.0
    rec = tp / (tp + fn) if (tp + fn) > 0 else 0.0
    f1 = 2 * prec * rec / (prec + rec) if (prec + rec) > 0 else 0.0
    return f1, prec, rec, tp, fp, fn, tn


def _majority(a: list, b: list) -> list:
    """Majority of 2 labels: agree (both 1 or both 0) carries; disagree
    treated as 1 (permissive union) OR as 0 (strict intersection).

    We compute BOTH:
      - majority_union: 1 if either A or B said exhibits
      - majority_intersection: 1 only if both said exhibits
    """
    union = [1 if (x or y) else 0 for x, y in zip(a, b)]
    intersection = [1 if (x and y) else 0 for x, y in zip(a, b)]
    return union, intersection


def compute_all():
    curator = _load(LABELS_DIR / "labels_curator.json")
    llm = _load(LABELS_DIR / "labels_llm.json")
    detector = _load(LABELS_DIR / "labels_detector.json")
    manifest = _load(CORPUS_DIR / "manifest.json")

    doc_ids = [d["doc_id"] for d in manifest["documents"]]
    strata = {d["doc_id"]: d["stratum"] for d in manifest["documents"]}

    # Per-document binary label vectors
    # Detectable frames only for A-vs-C, B-vs-C, majority-vs-C
    # All frames for A-vs-B
    per_doc = {}
    for doc_id in doc_ids:
        per_doc[doc_id] = {
            "A_all": _binary_labels(curator, doc_id, FRAME_IDS_ALL),
            "B_all": _binary_labels(llm, doc_id, FRAME_IDS_ALL),
            "A_det": _binary_labels(curator, doc_id, DETECTABLE_FRAMES),
            "B_det": _binary_labels(llm, doc_id, DETECTABLE_FRAMES),
            "C_det": _binary_labels(detector, doc_id, DETECTABLE_FRAMES),
        }

    # ─── Per-frame F1 of detector vs majority ───
    # For each frame, build A, B, C vectors across all docs
    per_frame = {}
    for i, fid in enumerate(DETECTABLE_FRAMES):
        a_vec = [per_doc[d]["A_det"][i] for d in doc_ids]
        b_vec = [per_doc[d]["B_det"][i] for d in doc_ids]
        c_vec = [per_doc[d]["C_det"][i] for d in doc_ids]
        maj_union, maj_inter = _majority(a_vec, b_vec)

        f1_ma_u, p_u, r_u, tp_u, fp_u, fn_u, tn_u = _f1(c_vec, maj_union)
        f1_ma_i, p_i, r_i, tp_i, fp_i, fn_i, tn_i = _f1(c_vec, maj_inter)
        f1_a, p_a, r_a, tp_a, fp_a, fn_a, tn_a = _f1(c_vec, a_vec)
        f1_b, p_b, r_b, tp_b, fp_b, fn_b, tn_b = _f1(c_vec, b_vec)
        k_ab = _cohen_kappa(a_vec, b_vec)
        k_ac = _cohen_kappa(a_vec, c_vec)
        k_bc = _cohen_kappa(b_vec, c_vec)

        per_frame[fid] = {
            "n_A_exhibits": sum(a_vec),
            "n_B_exhibits": sum(b_vec),
            "n_C_exhibits": sum(c_vec),
            "kappa_AB": round(k_ab, 3),
            "kappa_AC": round(k_ac, 3),
            "kappa_BC": round(k_bc, 3),
            "f1_C_vs_A": round(f1_a, 3),
            "f1_C_vs_B": round(f1_b, 3),
            "f1_C_vs_majority_union": round(f1_ma_u, 3),
            "f1_C_vs_majority_intersection": round(f1_ma_i, 3),
            "precision_C_vs_majority_union": round(p_u, 3),
            "recall_C_vs_majority_union": round(r_u, 3),
            "tp_fp_fn_tn_C_vs_majority_union": [tp_u, fp_u, fn_u, tn_u],
        }

    # ─── Aggregate macro metrics ───
    def _macro(key):
        vals = [f[key] for f in per_frame.values() if f[key] is not None]
        return round(mean(vals), 3) if vals else 0.0

    # Per-document kappa (across the 11 detectable frames)
    per_doc_metrics = {}
    for doc_id in doc_ids:
        a = per_doc[doc_id]["A_det"]
        b = per_doc[doc_id]["B_det"]
        c = per_doc[doc_id]["C_det"]
        maj_u, maj_i = _majority(a, b)
        per_doc_metrics[doc_id] = {
            "stratum": strata[doc_id],
            "n_A_exhibits": sum(a),
            "n_B_exhibits": sum(b),
            "n_C_exhibits": sum(c),
            "kappa_AB": round(_cohen_kappa(a, b), 3),
            "kappa_AC": round(_cohen_kappa(a, c), 3),
            "kappa_BC": round(_cohen_kappa(b, c), 3),
            "f1_C_vs_majority_union": round(_f1(c, maj_u)[0], 3),
            "f1_C_vs_majority_intersection": round(_f1(c, maj_i)[0], 3),
            "f1_C_vs_A": round(_f1(c, a)[0], 3),
            "f1_C_vs_B": round(_f1(c, b)[0], 3),
            "A_labels": [DETECTABLE_FRAMES[i] for i, v in enumerate(a) if v],
            "B_labels": [DETECTABLE_FRAMES[i] for i, v in enumerate(b) if v],
            "C_labels": [DETECTABLE_FRAMES[i] for i, v in enumerate(c) if v],
        }

    # Per-stratum macro-F1
    per_stratum = defaultdict(list)
    for doc_id, m in per_doc_metrics.items():
        per_stratum[m["stratum"]].append(m["f1_C_vs_majority_union"])
    per_stratum_macro = {s: round(mean(vs), 3) for s, vs in per_stratum.items()}

    # ─── A vs B agreement on ALL 20 frames (for reference) ───
    per_doc_AB_all = {}
    for doc_id in doc_ids:
        a_all = per_doc[doc_id]["A_all"]
        b_all = per_doc[doc_id]["B_all"]
        per_doc_AB_all[doc_id] = round(_cohen_kappa(a_all, b_all), 3)

    results = {
        "version": 1,
        "study": "FVS detector external validation v1",
        "corpus_size": len(doc_ids),
        "detectable_frames": DETECTABLE_FRAMES,
        "per_frame": per_frame,
        "per_document": per_doc_metrics,
        "per_stratum_macro_f1": per_stratum_macro,
        "per_doc_AB_kappa_all_frames": per_doc_AB_all,
        "aggregate": {
            "macro_kappa_AB_detectable_frames": _macro("kappa_AB"),
            "macro_kappa_AC_detectable_frames": _macro("kappa_AC"),
            "macro_kappa_BC_detectable_frames": _macro("kappa_BC"),
            "macro_f1_C_vs_A": _macro("f1_C_vs_A"),
            "macro_f1_C_vs_B": _macro("f1_C_vs_B"),
            "macro_f1_C_vs_majority_union": _macro("f1_C_vs_majority_union"),
            "macro_f1_C_vs_majority_intersection": _macro("f1_C_vs_majority_intersection"),
            "mean_kappa_AB_all_20_frames": round(mean(per_doc_AB_all.values()), 3),
        },
    }

    # ─── DESIGN.md §6 threshold assessment ───
    macro_f1_union = results["aggregate"]["macro_f1_C_vs_majority_union"]
    if macro_f1_union >= 0.7:
        verdict = "F1 >= 0.7: provisional evidence detector is at state-of-the-art. H1 strongly supported."
    elif macro_f1_union >= 0.6:
        verdict = "0.6 <= F1 < 0.7: detector aligned with labelers at usefully high levels. H1 supported."
    elif macro_f1_union >= 0.4:
        verdict = "0.4 <= F1 < 0.6: detector partially aligned. Product positioning and canon claims need scope revision."
    else:
        verdict = "F1 < 0.4: detector does not reliably measure what multi-source labeling identifies. H3 fires; project pivot required."
    results["verdict"] = {
        "macro_f1_C_vs_majority_union": macro_f1_union,
        "assessment": verdict,
    }

    # Write
    with open(RESULTS_DIR / "results.json", "w", encoding="utf-8") as f:
        json.dump(results, f, indent=2)

    # ─── Human-readable summary ───
    print("=" * 72)
    print("FVS External Validation Study v1: Results Summary")
    print("=" * 72)
    print(f"\nCorpus: {len(doc_ids)} documents (A={sum(1 for s in strata.values() if s=='A')}, B={sum(1 for s in strata.values() if s=='B')}, C={sum(1 for s in strata.values() if s=='C')})")
    print(f"Detectable frames scored: {len(DETECTABLE_FRAMES)} (text-side frames with detector rules)")

    print(f"\n--- Aggregate metrics ---")
    agg = results["aggregate"]
    print(f"  Macro-kappa A vs B (detectable frames): {agg['macro_kappa_AB_detectable_frames']:+.3f}")
    print(f"  Macro-kappa A vs C (curator vs detector): {agg['macro_kappa_AC_detectable_frames']:+.3f}")
    print(f"  Macro-kappa B vs C (LLM-judge vs detector): {agg['macro_kappa_BC_detectable_frames']:+.3f}")
    print(f"  Macro-F1 C vs A: {agg['macro_f1_C_vs_A']:.3f}")
    print(f"  Macro-F1 C vs B: {agg['macro_f1_C_vs_B']:.3f}")
    print(f"  Macro-F1 C vs majority-union (A|B): {agg['macro_f1_C_vs_majority_union']:.3f}  <- PRIMARY")
    print(f"  Macro-F1 C vs majority-intersection (A&B): {agg['macro_f1_C_vs_majority_intersection']:.3f}")
    print(f"  Mean kappa A-vs-B over ALL 20 frames: {agg['mean_kappa_AB_all_20_frames']:+.3f}")

    print(f"\n--- Per-stratum macro-F1 (C vs majority-union) ---")
    for s, v in sorted(per_stratum_macro.items()):
        print(f"  Stratum {s}: {v:.3f}  (n={sum(1 for x,st in strata.items() if st==s)} docs)")

    print(f"\n--- Per-frame results (detectable frames) ---")
    print(f"  {'Frame':<8}  {'nA':>3}  {'nB':>3}  {'nC':>3}  {'k_AB':>6}  {'k_AC':>6}  {'k_BC':>6}  {'F1_u':>5}  {'F1_i':>5}")
    for fid in DETECTABLE_FRAMES:
        f = per_frame[fid]
        print(f"  {fid:<8}  {f['n_A_exhibits']:>3}  {f['n_B_exhibits']:>3}  {f['n_C_exhibits']:>3}  "
              f"{f['kappa_AB']:>+6.3f}  {f['kappa_AC']:>+6.3f}  {f['kappa_BC']:>+6.3f}  "
              f"{f['f1_C_vs_majority_union']:>5.3f}  {f['f1_C_vs_majority_intersection']:>5.3f}")

    print(f"\n--- Per-document labels (side-by-side) ---")
    for doc_id in doc_ids:
        m = per_doc_metrics[doc_id]
        print(f"\n  {doc_id} ({m['stratum']}): k_AB={m['kappa_AB']:+.3f} k_AC={m['kappa_AC']:+.3f} F1_u={m['f1_C_vs_majority_union']:.3f}")
        print(f"    A: {m['A_labels']}")
        print(f"    B: {m['B_labels']}")
        print(f"    C: {m['C_labels']}")

    print(f"\n--- DESIGN.md §6 verdict ---")
    print(f"  Primary metric (Macro-F1 C vs majority-union): {macro_f1_union:.3f}")
    print(f"  {verdict}")

    print(f"\nResults written: {RESULTS_DIR / 'results.json'}")


if __name__ == "__main__":
    compute_all()