from __future__ import annotations

import json
from itertools import combinations
from pathlib import Path

from __future__ import annotations

import json
from itertools import combinations
from pathlib import Path

import numpy as np
from sklearn.metrics import cohen_kappa_score
from statsmodels.stats.inter_rater import fleiss_kappa

ROOT = Path(__file__).resolve().parents[1]
ANN_DIR = ROOT / "annotations"
OUT = ANN_DIR / "_irr.json"

CAT_GSB = ["good", "same", "bad"]


def load_annotations() -> list[dict]:
    files = sorted(p for p in ANN_DIR.glob("*.json") if not p.name.startswith("_"))
    out = []
    for p in files:
        d = json.loads(p.read_text(encoding="utf-8"))
        d["_file"] = p.name
        out.append(d)
    return out


def arena_label(row: dict) -> str:
    """Canonical outcome: lex-smaller model wins / larger wins / tie."""
    a, b = sorted([row["left"], row["right"]])
    w = row.get("winner") or (
        row["left"] if row["choice"] == "left" else row["right"] if row["choice"] == "right" else "tie"
    )
    if w == "tie":
        return "tie"
    if w == a:
        return "smaller_wins"
    if w == b:
        return "larger_wins"
    raise ValueError(f"bad arena winner {w} for {a}/{b}")


def arena_item_key(row: dict) -> str:
    a, b = sorted([row["left"], row["right"]])
    return f"{row['prompt_id']}|{a}|{b}"


def gsb_item_key(row: dict) -> str:
    return f"{row['prompt_id']}|{row['candidate']}"


def build_rater_matrix(
    annos: list[dict],
    protocol: str,
    item_fn,
    label_fn,
    categories: list[str],
) -> tuple[list[str], np.ndarray, list[str]]:
    """Return (item_ids, labels[n_items, n_raters] as int codes, annotator names)."""
    raters = [d["annotator"] for d in annos]
    item_ids = sorted(
        {item_fn(row) for d in annos for row in d[protocol]},
        key=lambda x: x,
    )
    cat_index = {c: i for i, c in enumerate(categories)}
    mat = np.full((len(item_ids), len(raters)), -1, dtype=int)
    item_index = {iid: i for i, iid in enumerate(item_ids)}
    for r, d in enumerate(annos):
        for row in d[protocol]:
            iid = item_fn(row)
            lab = label_fn(row)
            if lab not in cat_index:
                raise ValueError(f"unknown label {lab} in {protocol}")
            mat[item_index[iid], r] = cat_index[lab]
    if (mat < 0).any():
        missing = int((mat < 0).sum())
        raise SystemExit(f"{protocol}: {missing} missing rater cells")
    return item_ids, mat, raters


def fleiss_from_labels(labels: np.ndarray, n_cats: int) -> float:
    """labels: n_items x n_raters of category codes 0..n_cats-1."""
    n_items, n_raters = labels.shape
    table = np.zeros((n_items, n_cats), dtype=float)
    for i in range(n_items):
        for r in range(n_raters):
            table[i, labels[i, r]] += 1.0
    return float(fleiss_kappa(table, method="fleiss"))


def mean_pairwise_cohen(labels: np.ndarray) -> float:
    n_raters = labels.shape[1]
    ks = []
    for i, j in combinations(range(n_raters), 2):
        ks.append(cohen_kappa_score(labels[:, i], labels[:, j]))
    return float(np.mean(ks))


def bootstrap_metric(
    labels: np.ndarray,
    fn,
    n_boot: int = 2000,
    seed: int = 42,
) -> tuple[float, float, float]:
    """Item-level bootstrap → point, 2.5%, 97.5%."""
    rng = np.random.default_rng(seed)
    n = labels.shape[0]
    point = fn(labels)
    boots = np.empty(n_boot, dtype=float)
    for b in range(n_boot):
        idx = rng.integers(0, n, size=n)
        boots[b] = fn(labels[idx])
    return point, float(np.percentile(boots, 2.5)), float(np.percentile(boots, 97.5))


def icc_2(Y: np.ndarray) -> dict[str, float]:
    """Shrout & Fleiss ICC(2,1) and ICC(2,k); Y is n_subjects × n_raters."""
    Y = np.asarray(Y, dtype=float)
    n, k = Y.shape
    if n < 2 or k < 2:
        raise ValueError("need ≥2 subjects and ≥2 raters for ICC")
    mean_s = Y.mean(axis=1, keepdims=True)
    mean_r = Y.mean(axis=0, keepdims=True)
    grand = Y.mean()
    ss_rows = k * ((mean_s - grand) ** 2).sum()
    ss_cols = n * ((mean_r - grand) ** 2).sum()
    ss_tot = ((Y - grand) ** 2).sum()
    ss_err = ss_tot - ss_rows - ss_cols
    ms_rows = ss_rows / (n - 1)
    ms_cols = ss_cols / (k - 1)
    ms_err = ss_err / ((n - 1) * (k - 1))
    icc21 = (ms_rows - ms_err) / (ms_rows + (k - 1) * ms_err + k * (ms_cols - ms_err) / n)
    icc2k = (ms_rows - ms_err) / (ms_rows + (ms_cols - ms_err) / n)
    return {
        "icc_2_1": float(icc21),
        "icc_2_k": float(icc2k),
        "n_subjects": n,
        "n_raters": k,
        "ms_rows": float(ms_rows),
        "ms_err": float(ms_err),
    }


def mos_matrix(annos: list[dict]) -> tuple[np.ndarray, list[str]]:
    """Subjects = (prompt, model, dim) with all raters non-null."""
    raters = [d["annotator"] for d in annos]
    # collect scores keyed by subject → rater
    buckets: dict[str, dict[str, float]] = {}
    for d in annos:
        rname = d["annotator"]
        for row in d["mos"]:
            for dim, score in row["scores"].items():
                if score is None:
                    continue
                key = f"{row['prompt_id']}|{row['model']}|{dim}"
                buckets.setdefault(key, {})[rname] = float(score)
    complete = [k for k, v in buckets.items() if len(v) == len(raters)]
    complete.sort()
    Y = np.array([[buckets[k][r] for r in raters] for k in complete], dtype=float)
    return Y, complete


def interpret_kappa(k: float) -> str:
    # Landis & Koch (1977) rough bands
    if k < 0:
        return "低于机会水平"
    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 interpret_icc(v: float) -> str:
    # Koo & Li (2016) rough bands for ICC
    if v < 0.50:
        return "信度较差（poor）"
    if v < 0.75:
        return "中等信度（moderate）"
    if v < 0.90:
        return "良好信度（good）"
    return "优秀信度（excellent）"


def main() -> None:
    annos = load_annotations()
    n_raters = len(annos)

    # —— Arena ——
    arena_cats = ["smaller_wins", "larger_wins", "tie"]
    _, arena_lab, _ = build_rater_matrix(
        annos, "arena", arena_item_key, arena_label, arena_cats
    )
    arena_fleiss, arena_lo, arena_hi = bootstrap_metric(
        arena_lab, lambda L: fleiss_from_labels(L, len(arena_cats))
    )
    arena_cohen, cohen_lo, cohen_hi = bootstrap_metric(arena_lab, mean_pairwise_cohen)

    # —— GSB ——
    _, gsb_lab, _ = build_rater_matrix(
        annos, "gsb", gsb_item_key, lambda r: r["choice"], CAT_GSB
    )
    gsb_fleiss, gsb_lo, gsb_hi = bootstrap_metric(
        gsb_lab, lambda L: fleiss_from_labels(L, len(CAT_GSB))
    )
    gsb_cohen, gsb_c_lo, gsb_c_hi = bootstrap_metric(gsb_lab, mean_pairwise_cohen)

    # —— MOS ICC ——
    Y, subjects = mos_matrix(annos)
    mos_point = icc_2(Y)

    def boot_icc(kind: str):
        rng = np.random.default_rng(42)
        n = Y.shape[0]
        boots = []
        for _ in range(2000):
            idx = rng.integers(0, n, size=n)
            boots.append(icc_2(Y[idx])[kind])
        arr = np.array(boots, dtype=float)
        return float(np.percentile(arr, 2.5)), float(np.percentile(arr, 97.5))

    icc21_lo, icc21_hi = boot_icc("icc_2_1")
    icc2k_lo, icc2k_hi = boot_icc("icc_2_k")

    result = {
        "schema": "magicframe_irr_v1",
        "n_annotators": n_raters,
        "annotators": [d["annotator"] for d in annos],
        "n_bootstrap": 2000,
        "ci": "percentile_95",
        "arena": {
            "n_items": int(arena_lab.shape[0]),
            "categories": arena_cats,
            "fleiss_kappa": round(arena_fleiss, 4),
            "fleiss_ci_low": round(arena_lo, 4),
            "fleiss_ci_high": round(arena_hi, 4),
            "mean_pairwise_cohen_kappa": round(arena_cohen, 4),
            "cohen_ci_low": round(cohen_lo, 4),
            "cohen_ci_high": round(cohen_hi, 4),
            "interpretation": interpret_kappa(arena_fleiss),
        },
        "gsb": {
            "n_items": int(gsb_lab.shape[0]),
            "categories": CAT_GSB,
            "fleiss_kappa": round(gsb_fleiss, 4),
            "fleiss_ci_low": round(gsb_lo, 4),
            "fleiss_ci_high": round(gsb_hi, 4),
            "mean_pairwise_cohen_kappa": round(gsb_cohen, 4),
            "cohen_ci_low": round(gsb_c_lo, 4),
            "cohen_ci_high": round(gsb_c_hi, 4),
            "interpretation": interpret_kappa(gsb_fleiss),
        },
        "mos": {
            "n_subjects": mos_point["n_subjects"],
            "n_raters": mos_point["n_raters"],
            "subject_def": "(prompt_id, model, dim) with all raters non-null",
            "icc_2_1": round(mos_point["icc_2_1"], 4),
            "icc_2_1_ci_low": round(icc21_lo, 4),
            "icc_2_1_ci_high": round(icc21_hi, 4),
            "icc_2_k": round(mos_point["icc_2_k"], 4),
            "icc_2_k_ci_low": round(icc2k_lo, 4),
            "icc_2_k_ci_high": round(icc2k_hi, 4),
            "k": n_raters,
            "interpretation_single": interpret_icc(mos_point["icc_2_1"]),
            "interpretation_avg": interpret_icc(mos_point["icc_2_k"]),
        },
    }

    OUT.write_text(json.dumps(result, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result, ensure_ascii=False, indent=2))
    print(f"\nWrote {OUT}")


if __name__ == "__main__":
    main()