#!/usr/bin/env python3
"""Inter-rater reliability for the IETF landscape survey.

Reads the two re-rating result files written by rerate-intercoder.py
(data/rerate/sonnet.jsonl, data/rerate/haiku.jsonl) plus the existing
production labels (ratings table), and reports:

  - Cohen's kappa on PRIMARY CATEGORY (nominal) for each rater pair
  - Quadratic-weighted kappa on each ordinal dimension (novelty, maturity,
    overlap, momentum, relevance)
  - Raw agreement %, and the category confusion (most-confused pairs)

Pairs compared: sonnet-rerate vs haiku-rerate (the controlled inter-coder run),
and each rerate vs the existing production labels (drift check).

Pure stdlib + numpy. No API calls. Read-only on the DB.
Writes data/reports/survey-kappa.md.

Usage: PYTHONPATH=src python3 scripts/survey-kappa.py
"""
from __future__ import annotations

import json
from collections import Counter
from pathlib import Path

import numpy as np

from ietf_analyzer.config import Config
from ietf_analyzer.db import Database

RERATE_DIR = Path("data/rerate")
OUT = Path("data/reports/survey-kappa.md")
DIMS = ["novelty", "maturity", "overlap", "momentum", "relevance"]
DIM_KEYS = {"novelty": ("n", "novelty"), "maturity": ("m", "maturity"),
            "overlap": ("o", "overlap"), "momentum": ("mo", "momentum"),
            "relevance": ("r", "relevance")}


def _strip_fence(text: str) -> str:
    t = text.strip()
    if t.startswith("```"):
        t = t.split("\n", 1)[1] if "\n" in t else t[3:]
        if t.rstrip().endswith("```"):
            t = t.rstrip()[:-3]
    return t.strip()


def _clamp(v):
    try:
        v = int(round(float(v)))
    except (TypeError, ValueError):
        return None
    return min(5, max(1, v))


def parse_rerate(alias: str) -> dict[str, dict]:
    """draft_name -> {primary_cat, dims...} from a rerate jsonl."""
    out = {}
    p = RERATE_DIR / f"{alias}.jsonl"
    if not p.exists():
        return out
    for line in p.read_text().splitlines():
        if not line.strip():
            continue
        rec = json.loads(line)
        if "raw" not in rec:
            continue
        try:
            data = json.loads(_strip_fence(rec["raw"]))
        except json.JSONDecodeError:
            continue
        cats = data.get("c", data.get("categories", []))
        primary = cats[0] if isinstance(cats, list) and cats else None
        entry = {"primary": primary}
        for dim, (k1, k2) in DIM_KEYS.items():
            entry[dim] = _clamp(data.get(k1, data.get(k2)))
        out[rec["draft_name"]] = entry
    return out


def load_prod(db: Database) -> dict[str, dict]:
    out = {}
    rows = db.conn.execute(
        """SELECT r.draft_name, r.categories, r.novelty, r.maturity,
                  r.overlap, r.momentum, r.relevance
           FROM drafts d JOIN ratings r ON d.name = r.draft_name
           WHERE d.source='ietf' AND (r.false_positive=0 OR r.false_positive IS NULL)"""
    ).fetchall()
    for name, cats_json, *dimvals in rows:
        try:
            cats = json.loads(cats_json) if cats_json else []
        except json.JSONDecodeError:
            cats = []
        e = {"primary": cats[0] if cats else None}
        for dim, v in zip(DIMS, dimvals):
            e[dim] = _clamp(v)
        out[name] = e
    return out


def cohen_kappa(a: list, b: list) -> tuple[float, float, int]:
    """Nominal Cohen's kappa. Returns (kappa, raw_agreement, n)."""
    labels = sorted(set(a) | set(b))
    idx = {l: i for i, l in enumerate(labels)}
    k = len(labels)
    m = np.zeros((k, k))
    for x, y in zip(a, b):
        m[idx[x], idx[y]] += 1
    n = m.sum()
    po = np.trace(m) / n
    pe = (m.sum(0) @ m.sum(1)) / (n * n)
    kappa = (po - pe) / (1 - pe) if (1 - pe) else 1.0
    return kappa, po, int(n)


def weighted_kappa(a: list, b: list, k: int = 5) -> tuple[float, float, int]:
    """Quadratic-weighted kappa for ordinal 1..k ratings."""
    pairs = [(x, y) for x, y in zip(a, b) if x is not None and y is not None]
    if not pairs:
        return float("nan"), float("nan"), 0
    a2, b2 = zip(*pairs)
    o = np.zeros((k, k))
    for x, y in pairs:
        o[x - 1, y - 1] += 1
    n = o.sum()
    w = np.zeros((k, k))
    for i in range(k):
        for j in range(k):
            w[i, j] = (i - j) ** 2 / (k - 1) ** 2
    ha = np.array([a2.count(v) for v in range(1, k + 1)], float)
    hb = np.array([b2.count(v) for v in range(1, k + 1)], float)
    e = np.outer(ha, hb) / n
    num = (w * o).sum()
    den = (w * e).sum()
    kappa = 1 - num / den if den else 1.0
    raw_agree = np.trace(o) / n
    return kappa, raw_agree, int(n)


def interpret(k: float) -> str:
    if k != k:
        return "n/a"
    if k < 0: return "worse than chance"
    if k < 0.20: return "slight"
    if k < 0.40: return "fair"
    if k < 0.60: return "moderate"
    if k < 0.80: return "substantial"
    return "almost perfect"


def compare(name: str, A: dict, B: dict, lines: list):
    shared = sorted(set(A) & set(B))
    lines.append(f"\n## {name}  (n shared = {len(shared)})\n")
    # primary category
    pa = [A[d]["primary"] for d in shared if A[d]["primary"] and B[d]["primary"]]
    pb = [B[d]["primary"] for d in shared if A[d]["primary"] and B[d]["primary"]]
    kappa, po, n = cohen_kappa(pa, pb)
    lines.append(f"**Primary category** (Cohen's κ): κ = {kappa:.3f} ({interpret(kappa)}), "
                 f"raw agreement {po:.1%}, n = {n}\n")
    # confusion: top disagreements
    dis = Counter()
    for d in shared:
        x, y = A[d]["primary"], B[d]["primary"]
        if x and y and x != y:
            dis[tuple(sorted((x, y)))] += 1
    if dis:
        lines.append("\nMost-confused category pairs:\n")
        lines.append("| A | B | count |\n|---|---|------:|\n")
        for (x, y), c in dis.most_common(8):
            lines.append(f"| {x} | {y} | {c} |\n")
    # ordinal dims
    lines.append("\n**Ordinal dimensions** (quadratic-weighted κ):\n\n")
    lines.append("| dimension | κ_w | raw agree | n |\n|---|---:|---:|---:|\n")
    for dim in DIMS:
        a = [A[d][dim] for d in shared]
        b = [B[d][dim] for d in shared]
        kw, ra, n = weighted_kappa(a, b)
        lines.append(f"| {dim} | {kw:.3f} ({interpret(kw)}) | {ra:.1%} | {n} |\n")


def main():
    db = Database(Config.load())
    sonnet = parse_rerate("sonnet")
    haiku = parse_rerate("haiku")
    prod = load_prod(db)
    print(f"parsed: sonnet={len(sonnet)} haiku={len(haiku)} prod={len(prod)}")
    if not sonnet or not haiku:
        print("rerate files incomplete — run rerate-intercoder.py --collect first")
        return

    lines = ["# Inter-rater reliability — IETF landscape survey\n",
             f"\nCorpus: clean IETF (n≈524). Sonnet={len(sonnet)}, Haiku={len(haiku)}, prod labels={len(prod)}.\n",
             "\nκ interpretation (Landis & Koch): <0.2 slight, 0.2–0.4 fair, "
             "0.4–0.6 moderate, 0.6–0.8 substantial, >0.8 almost perfect.\n"]
    compare("Sonnet (re-rate) vs Haiku (re-rate) — controlled inter-coder", sonnet, haiku, lines)
    compare("Sonnet (re-rate) vs Production labels — drift/stability", sonnet, prod, lines)
    compare("Haiku (re-rate) vs Production labels", haiku, prod, lines)

    OUT.write_text("".join(lines))
    print(f"wrote {OUT}")
    print("".join(lines))


if __name__ == "__main__":
    main()