#!/usr/bin/env python3
"""
S9_recompute.py - Methodological audit recomputations for the umbrella review.

Computes, from raw rater files in _amstar2_toolkit/peer_review/{sergio,david,results}/:
  - Fleiss' kappa per AMSTAR-2 item and overall (3 raters x 33 SRs x 16 items),
    using the standard Fleiss (1971) formulation generalised to subjects with
    a fixed number of ratings.
  - Gwet's AC1 per item and overall (Gwet 2008), which corrects the Cohen-style
    chance-agreement term using the marginal category frequencies aggregated
    across raters - robust to the prevalence paradox documented by
    Feinstein & Cicchetti (1990).
  - Pairwise Cohen kappa per item (for parity with the existing report).
  - Corrected Covered Area (CCA) per Pieper et al. 2014: CCA = (N - r) / (rc - r)
    where N = total occurrences of primary studies across the included SRs,
    r = number of unique primary studies, c = number of SRs in the citation
    matrix. The script reads an optional matrix at
    06_supplementary/cca_matrix.csv (rows = primary-study DOIs/PMIDs,
    columns = SR keys, cells = 0/1). When the matrix is absent, the function
    returns the sentinel string "INPUT_MATRIX_REQUIRED" and CCA is reported
    as such in the output CSV.

Output: predictor_ypN0/paper_umbrella/06_supplementary/S9_recompute_results.csv
The script is idempotent and has zero non-stdlib dependencies.

References
----------
Fleiss JL. Measuring nominal scale agreement among many raters. Psychol Bull
    1971;76:378-382.
Gwet KL. Computing inter-rater reliability and its variance in the presence
    of high agreement. Br J Math Stat Psychol 2008;61:29-48.
Feinstein AR, Cicchetti DV. High agreement but low kappa: I. The problems of
    two paradoxes. J Clin Epidemiol 1990;43:543-549.
Pieper D, Antoine SL, Mathes T, Neugebauer EAM, Eikermann M. Systematic review
    finds overlapping reviews were not mentioned in every other overview.
    J Clin Epidemiol 2014;67:368-375.
"""

from __future__ import annotations

import csv
import json
import math
import sys
from pathlib import Path
from typing import Dict, List, Tuple

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------

THIS_FILE = Path(__file__).resolve()
PAPER_ROOT = THIS_FILE.parents[1]               # paper_umbrella/
PROJECT_ROOT = PAPER_ROOT.parents[1]            # ganglio sergioc/
TOOLKIT = PROJECT_ROOT / "_amstar2_toolkit"
PR_DIR = TOOLKIT / "peer_review"
RATER_DIRS = {
    "llm": TOOLKIT / "results",
    "sergio": PR_DIR / "sergio",
    "david": PR_DIR / "david",
}
OUT_CSV = PAPER_ROOT / "06_supplementary" / "S9_recompute_results.csv"
CCA_MATRIX = PAPER_ROOT / "06_supplementary" / "cca_matrix.csv"

ITEMS = [f"item{i}" for i in range(1, 17)]

# AMSTAR-2 categories that may appear in any rater file. The five canonical
# values are kept distinct (we do NOT collapse Yes/Partial Yes here so that
# the agreement metrics reflect the original 5-level rating scale).
CATEGORIES = ["Yes", "Partial Yes", "No", "Cannot tell", "No meta-analysis"]


# ---------------------------------------------------------------------------
# Data loading
# ---------------------------------------------------------------------------

def load_rater(rater_dir: Path) -> Dict[str, Dict[str, str]]:
    """Return {sr_key: {item: score}} for a single rater directory."""
    out: Dict[str, Dict[str, str]] = {}
    for f in sorted(rater_dir.glob("*.json")):
        try:
            data = json.loads(f.read_text(encoding="utf-8"))
        except Exception as e:                                  # pragma: no cover
            print(f"WARN  {f.name}: {e}", file=sys.stderr)
            continue
        items = data.get("items") or {}
        scores = {}
        for it in ITEMS:
            entry = items.get(it) or {}
            score = entry.get("score") if isinstance(entry, dict) else None
            scores[it] = (score or "Cannot tell").strip()
        out[f.stem] = scores
    return out


def normalise(score: str) -> str:
    s = (score or "").strip()
    # tolerate slight casing or punctuation differences
    aliases = {
        "yes": "Yes",
        "partial yes": "Partial Yes",
        "partialyes": "Partial Yes",
        "no": "No",
        "cannot tell": "Cannot tell",
        "no meta-analysis": "No meta-analysis",
        "no metaanalysis": "No meta-analysis",
        "no meta analysis": "No meta-analysis",
        "n/a": "No meta-analysis",
        "na": "No meta-analysis",
    }
    return aliases.get(s.lower(), s if s in CATEGORIES else "Cannot tell")


# ---------------------------------------------------------------------------
# Agreement metrics
# ---------------------------------------------------------------------------

def fleiss_kappa(matrix: List[List[int]]) -> float:
    """
    Fleiss (1971) kappa.
    matrix: list of N rows (subjects). Each row has K columns (categories) and
    sums to n raters per subject (here, n=3).
    """
    N = len(matrix)
    if N == 0:
        return float("nan")
    n = sum(matrix[0])                                          # raters per subj
    if n < 2:
        return float("nan")
    K = len(matrix[0])

    # P_i: agreement for subject i
    P = []
    for row in matrix:
        s = sum(c * (c - 1) for c in row)
        P.append(s / (n * (n - 1)))
    P_bar = sum(P) / N

    # p_j: marginal proportion of category j
    totals = [0.0] * K
    for row in matrix:
        for j, c in enumerate(row):
            totals[j] += c
    grand = N * n
    p = [t / grand for t in totals]
    Pe_bar = sum(pj * pj for pj in p)

    if Pe_bar >= 1.0:
        return float("nan")
    return (P_bar - Pe_bar) / (1 - Pe_bar)


def gwet_ac1(matrix: List[List[int]]) -> float:
    """
    Gwet (2008) AC1 for nominal categories with n raters per subject.
    Uses the standard 'overall agreement' formulation:
        Pa = (1/N) sum_i (1/(n*(n-1))) sum_j r_ij*(r_ij-1)
        Pe = (1/(K-1)) sum_j p_j * (1 - p_j)
        AC1 = (Pa - Pe) / (1 - Pe)
    """
    N = len(matrix)
    if N == 0:
        return float("nan")
    n = sum(matrix[0])
    if n < 2:
        return float("nan")
    K = len(matrix[0])

    Pa = 0.0
    for row in matrix:
        Pa += sum(c * (c - 1) for c in row) / (n * (n - 1))
    Pa /= N

    totals = [0.0] * K
    for row in matrix:
        for j, c in enumerate(row):
            totals[j] += c
    grand = N * n
    p = [t / grand for t in totals]
    Pe = sum(pj * (1 - pj) for pj in p) / (K - 1)

    if Pe >= 1.0:
        return float("nan")
    return (Pa - Pe) / (1 - Pe)


def cohen_kappa(a: List[str], b: List[str], cats: List[str]) -> float:
    """Cohen kappa for two raters on the same N subjects."""
    n = len(a)
    if n == 0 or n != len(b):
        return float("nan")
    K = len(cats)
    idx = {c: i for i, c in enumerate(cats)}
    obs = [[0] * K for _ in range(K)]
    for x, y in zip(a, b):
        obs[idx[x]][idx[y]] += 1
    Po = sum(obs[i][i] for i in range(K)) / n
    rows = [sum(r) / n for r in obs]
    cols = [sum(obs[i][j] for i in range(K)) / n for j in range(K)]
    Pe = sum(rows[i] * cols[i] for i in range(K))
    if Pe >= 1.0:
        return float("nan")
    return (Po - Pe) / (1 - Pe)


# ---------------------------------------------------------------------------
# Build the rating matrix per item
# ---------------------------------------------------------------------------

def build_matrix(
    item: str,
    raters: Dict[str, Dict[str, Dict[str, str]]],
    sr_keys: List[str],
) -> Tuple[List[List[int]], Dict[str, List[str]]]:
    """
    For one AMSTAR-2 item, build the N x K count matrix used by Fleiss/Gwet,
    and also a {rater: [scores]} mapping aligned by sr_keys for pairwise
    Cohen kappa.
    """
    K = len(CATEGORIES)
    cat_index = {c: i for i, c in enumerate(CATEGORIES)}
    matrix = []
    aligned = {r: [] for r in raters}
    for sr in sr_keys:
        row = [0] * K
        for rater_name, data in raters.items():
            score = normalise(data.get(sr, {}).get(item, "Cannot tell"))
            row[cat_index[score]] += 1
            aligned[rater_name].append(score)
        matrix.append(row)
    return matrix, aligned


# ---------------------------------------------------------------------------
# CCA
# ---------------------------------------------------------------------------

def cca(citation_matrix_csv: Path, sr_subset: List[str] | None = None):
    """
    Compute CCA per Pieper 2014 from a primary-study x SR citation matrix.

    citation_matrix_csv: rows = primary studies, columns include 'study_id'
                         followed by SR keys; cells are 0/1.
    sr_subset: list of SR keys to restrict to (None = all columns).
    Returns (CCA_value or 'INPUT_MATRIX_REQUIRED', N, r, c) tuple.
    """
    if not citation_matrix_csv.exists():
        return "INPUT_MATRIX_REQUIRED", None, None, None

    with citation_matrix_csv.open(encoding="utf-8") as f:
        rdr = csv.reader(f)
        header = next(rdr)
        sr_cols = header[1:] if sr_subset is None else sr_subset
        col_idx = {h: i for i, h in enumerate(header)}
        N = 0
        r = 0
        for row in rdr:
            occ = sum(int(row[col_idx[c]]) for c in sr_cols if c in col_idx)
            if occ >= 1:
                r += 1
                N += occ
    c = len(sr_cols)
    if c <= 1 or r == 0 or (r * c - r) == 0:
        return float("nan"), N, r, c
    return (N - r) / (r * c - r), N, r, c


# ---------------------------------------------------------------------------
# Driver
# ---------------------------------------------------------------------------

def main() -> int:
    raters = {name: load_rater(d) for name, d in RATER_DIRS.items()
              if d.exists()}
    if len(raters) < 2:
        print("ERROR: need at least 2 rater directories", file=sys.stderr)
        return 1

    sr_keys = sorted(set.intersection(*(set(d.keys()) for d in raters.values())))
    if not sr_keys:
        print("ERROR: no SR keys common to all raters", file=sys.stderr)
        return 1

    rows = []
    # per-item metrics
    overall_matrix: List[List[int]] = []
    overall_pairs = {("llm", "sergio"): [[], []],
                     ("llm", "david"): [[], []],
                     ("sergio", "david"): [[], []]}

    for it in ITEMS:
        m, aligned = build_matrix(it, raters, sr_keys)
        fk = fleiss_kappa(m)
        ac1 = gwet_ac1(m)
        c_ls = cohen_kappa(aligned["llm"], aligned["sergio"], CATEGORIES)
        c_ld = cohen_kappa(aligned["llm"], aligned["david"], CATEGORIES)
        c_sd = cohen_kappa(aligned["sergio"], aligned["david"], CATEGORIES)
        rows.append({
            "item": it,
            "fleiss_kappa": f"{fk:.3f}",
            "gwet_ac1": f"{ac1:.3f}",
            "cohen_llm_sergio": f"{c_ls:.3f}",
            "cohen_llm_david": f"{c_ld:.3f}",
            "cohen_sergio_david": f"{c_sd:.3f}",
            "n_subjects": len(sr_keys),
        })
        overall_matrix.extend(m)
        for (a, b), (la, lb) in overall_pairs.items():
            la.extend(aligned[a]); lb.extend(aligned[b])

    fk_overall = fleiss_kappa(overall_matrix)
    ac1_overall = gwet_ac1(overall_matrix)
    c_ls_overall = cohen_kappa(*overall_pairs[("llm", "sergio")], CATEGORIES)
    c_ld_overall = cohen_kappa(*overall_pairs[("llm", "david")], CATEGORIES)
    c_sd_overall = cohen_kappa(*overall_pairs[("sergio", "david")], CATEGORIES)

    rows.append({
        "item": "OVERALL",
        "fleiss_kappa": f"{fk_overall:.3f}",
        "gwet_ac1": f"{ac1_overall:.3f}",
        "cohen_llm_sergio": f"{c_ls_overall:.3f}",
        "cohen_llm_david": f"{c_ld_overall:.3f}",
        "cohen_sergio_david": f"{c_sd_overall:.3f}",
        "n_subjects": len(sr_keys) * len(ITEMS),
    })

    # CCA on three subsets, when matrix available
    triage_path = TOOLKIT / "eligibility_triage.json"
    eligible: List[str] = []
    borderline: List[str] = []
    full23: List[str] = []
    if triage_path.exists():
        triage = json.loads(triage_path.read_text(encoding="utf-8"))
        for ev in triage.get("evaluations", []):
            v = ev["verdict"]
            if v == "ELIGIBLE":
                eligible.append(ev["sr_key"]); full23.append(ev["sr_key"])
            elif v == "BORDERLINE":
                borderline.append(ev["sr_key"]); full23.append(ev["sr_key"])

    cca_full, N_full, r_full, c_full = cca(CCA_MATRIX, full23 or None)
    cca_elig, N_e, r_e, c_e = cca(CCA_MATRIX, eligible or None)
    cca_bord, N_b, r_b, c_b = cca(CCA_MATRIX, borderline or None)

    def fmt_cca(v, N, r, c):
        if isinstance(v, str):
            return {"cca": v, "N": "", "r": "", "c": str(c) if c else ""}
        return {"cca": f"{v*100:.2f}%", "N": str(N), "r": str(r), "c": str(c)}

    rows.append({"item": "CCA_full_23",
                 **{"fleiss_kappa": "", "gwet_ac1": "",
                    "cohen_llm_sergio": "", "cohen_llm_david": "",
                    "cohen_sergio_david": "", "n_subjects": ""},
                 **fmt_cca(cca_full, N_full, r_full, len(full23))})
    rows.append({"item": "CCA_eligible_13",
                 **{"fleiss_kappa": "", "gwet_ac1": "",
                    "cohen_llm_sergio": "", "cohen_llm_david": "",
                    "cohen_sergio_david": "", "n_subjects": ""},
                 **fmt_cca(cca_elig, N_e, r_e, len(eligible))})
    rows.append({"item": "CCA_borderline_10",
                 **{"fleiss_kappa": "", "gwet_ac1": "",
                    "cohen_llm_sergio": "", "cohen_llm_david": "",
                    "cohen_sergio_david": "", "n_subjects": ""},
                 **fmt_cca(cca_bord, N_b, r_b, len(borderline))})

    # write
    OUT_CSV.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = ["item", "fleiss_kappa", "gwet_ac1",
                  "cohen_llm_sergio", "cohen_llm_david", "cohen_sergio_david",
                  "n_subjects", "cca", "N", "r", "c"]
    with OUT_CSV.open("w", encoding="utf-8", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        for row in rows:
            w.writerow({k: row.get(k, "") for k in fieldnames})

    # console summary
    print(f"raters loaded:        {list(raters.keys())}")
    print(f"SRs in common:        {len(sr_keys)}")
    print(f"Fleiss kappa overall: {fk_overall:.3f}")
    print(f"Gwet AC1 overall:     {ac1_overall:.3f}")
    print(f"Cohen overall LLM-Sergio:    {c_ls_overall:.3f}")
    print(f"Cohen overall LLM-David:     {c_ld_overall:.3f}")
    print(f"Cohen overall Sergio-David:  {c_sd_overall:.3f}")
    print(f"CCA (full 23 SRs):    {cca_full}")
    print(f"CCA (13 eligible):    {cca_elig}")
    print(f"CCA (10 borderline):  {cca_bord}")
    print(f"output written:       {OUT_CSV}")
    return 0


if __name__ == "__main__":
    sys.exit(main())
