experiments/required_analyses/common.py

from __future__ import annotations

import csv
import json
import math
import random
from collections import Counter, defaultdict
from pathlib import Path
from statistics import mean, pstdev
from typing import Any, Dict, Iterable, List, Sequence, Tuple


Row = Dict[str, Any]


def ensure_dir(path: Path) -> Path:
    path.mkdir(parents=True, exist_ok=True)
    return path


def read_csv(path: Path) -> List[Row]:
    with path.open("r", encoding="utf-8-sig", newline="") as f:
        return list(csv.DictReader(f))


def read_jsonl(path: Path) -> List[Row]:
    rows: List[Row] = []
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                rows.append(json.loads(line))
    return rows


def write_csv(path: Path, rows: Sequence[Row], fieldnames: Sequence[str] | None = None) -> None:
    ensure_dir(path.parent)
    if fieldnames is None:
        fieldnames = []
        for row in rows:
            for key in row:
                if key not in fieldnames:
                    fieldnames.append(key)
    with path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def write_json(path: Path, payload: Any) -> None:
    ensure_dir(path.parent)
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def fnum(value: Any, default: float = math.nan) -> float:
    try:
        if value is None or value == "":
            return default
        return float(value)
    except (TypeError, ValueError):
        return default


def safe_mean(values: Iterable[Any]) -> float:
    vals = [fnum(v) for v in values]
    vals = [v for v in vals if not math.isnan(v)]
    return mean(vals) if vals else math.nan


def safe_std(values: Iterable[Any]) -> float:
    vals = [fnum(v) for v in values]
    vals = [v for v in vals if not math.isnan(v)]
    return pstdev(vals) if len(vals) > 1 else 0.0


def bootstrap_ci(values: Sequence[float], n: int = 10000, seed: int = 0) -> Tuple[float, float]:
    vals = [v for v in values if not math.isnan(v)]
    if not vals:
        return math.nan, math.nan
    rng = random.Random(seed)
    boot = []
    for _ in range(n):
        sample = [vals[rng.randrange(len(vals))] for _ in vals]
        boot.append(mean(sample))
    boot.sort()
    return boot[int(0.025 * (len(boot) - 1))], boot[int(0.975 * (len(boot) - 1))]


def pearson(xs: Sequence[float], ys: Sequence[float]) -> float:
    pairs = [(x, y) for x, y in zip(xs, ys) if not math.isnan(x) and not math.isnan(y)]
    if len(pairs) < 2:
        return math.nan
    xvals, yvals = zip(*pairs)
    mx, my = mean(xvals), mean(yvals)
    num = sum((x - mx) * (y - my) for x, y in pairs)
    den_x = math.sqrt(sum((x - mx) ** 2 for x in xvals))
    den_y = math.sqrt(sum((y - my) ** 2 for y in yvals))
    return num / (den_x * den_y) if den_x and den_y else math.nan


def rank(values: Sequence[float]) -> List[float]:
    order = sorted((v, i) for i, v in enumerate(values))
    out = [0.0] * len(values)
    i = 0
    while i < len(order):
        j = i
        while j + 1 < len(order) and order[j + 1][0] == order[i][0]:
            j += 1
        r = (i + j + 2) / 2
        for _, idx in order[i : j + 1]:
            out[idx] = r
        i = j + 1
    return out


def spearman(xs: Sequence[float], ys: Sequence[float]) -> float:
    pairs = [(x, y) for x, y in zip(xs, ys) if not math.isnan(x) and not math.isnan(y)]
    if len(pairs) < 2:
        return math.nan
    xvals, yvals = zip(*pairs)
    return pearson(rank(xvals), rank(yvals))


def mae(xs: Sequence[float], ys: Sequence[float]) -> float:
    pairs = [(x, y) for x, y in zip(xs, ys) if not math.isnan(x) and not math.isnan(y)]
    return mean(abs(x - y) for x, y in pairs) if pairs else math.nan


def prf(y_true: Sequence[str], y_pred: Sequence[str]) -> Row:
    labels = sorted(set(y_true) | set(y_pred))
    total = len(y_true)
    acc = sum(1 for t, p in zip(y_true, y_pred) if t == p) / total if total else math.nan
    scores = []
    for label in labels:
        tp = sum(1 for t, p in zip(y_true, y_pred) if t == label and p == label)
        fp = sum(1 for t, p in zip(y_true, y_pred) if t != label and p == label)
        fn = sum(1 for t, p in zip(y_true, y_pred) if t == label and p != label)
        prec = tp / (tp + fp) if tp + fp else 0.0
        rec = tp / (tp + fn) if tp + fn else 0.0
        f1 = 2 * prec * rec / (prec + rec) if prec + rec else 0.0
        scores.append((prec, rec, f1))
    return {
        "accuracy": acc,
        "macro_precision": mean(s[0] for s in scores) if scores else math.nan,
        "macro_recall": mean(s[1] for s in scores) if scores else math.nan,
        "macro_f1": mean(s[2] for s in scores) if scores else math.nan,
        "n": total,
    }


def group_by(rows: Sequence[Row], key: str) -> Dict[str, List[Row]]:
    out: Dict[str, List[Row]] = defaultdict(list)
    for row in rows:
        out[str(row.get(key, ""))].append(row)
    return dict(out)


def krippendorff_alpha_interval(rows: Sequence[Row], item_key: str, rater_key: str, score_key: str) -> float:
    by_item: Dict[str, List[float]] = defaultdict(list)
    for row in rows:
        by_item[str(row[item_key])].append(fnum(row[score_key]))
    observed = []
    all_values = []
    for values in by_item.values():
        vals = [v for v in values if not math.isnan(v)]
        all_values.extend(vals)
        observed.extend((a - b) ** 2 for i, a in enumerate(vals) for b in vals[i + 1 :])
    if len(all_values) < 2 or not observed:
        return math.nan
    do = mean(observed)
    de = mean((a - b) ** 2 for i, a in enumerate(all_values) for b in all_values[i + 1 :])
    return 1 - do / de if de else math.nan


def fleiss_kappa(rows: Sequence[Row], item_key: str, label_key: str) -> float:
    by_item: Dict[str, List[str]] = defaultdict(list)
    for row in rows:
        by_item[str(row[item_key])].append(str(row[label_key]))
    categories = sorted({str(row[label_key]) for row in rows})
    if not categories:
        return math.nan
    n_raters = max(len(v) for v in by_item.values())
    p_items = []
    cat_totals = Counter()
    for labels in by_item.values():
        counts = Counter(labels)
        cat_totals.update(counts)
        denom = len(labels) * (len(labels) - 1)
        if denom:
            p_items.append((sum(c * c for c in counts.values()) - len(labels)) / denom)
    p_bar = mean(p_items) if p_items else math.nan
    total_labels = sum(cat_totals.values())
    p_e = sum((cat_totals[c] / total_labels) ** 2 for c in categories) if total_labels else math.nan
    return (p_bar - p_e) / (1 - p_e) if p_e != 1 else math.nan


def paired_bootstrap_delta(rows: Sequence[Row], baseline: str, system: str, metric: str, item_key: str = "item_id", n: int = 10000, seed: int = 0) -> Row:
    pairs: Dict[str, Dict[str, float]] = defaultdict(dict)
    for row in rows:
        pairs[str(row[item_key])][str(row["system"])] = fnum(row[metric])
    diffs = [v[system] - v[baseline] for v in pairs.values() if system in v and baseline in v]
    lo, hi = bootstrap_ci(diffs, n=n, seed=seed)
    delta = safe_mean(diffs)
    std = safe_std(diffs)
    effect = delta / std if std else math.inf
    p = sum(1 for d in diffs if d <= 0) / len(diffs) if diffs else math.nan
    p = min(1.0, 2 * min(p, 1 - p)) if not math.isnan(p) else math.nan
    return {
        "comparison": f"{system} - {baseline}",
        "metric": metric,
        "delta": delta,
        "ci_low": lo,
        "ci_high": hi,
        "cohens_dz": effect,
        "p_raw": p,
        "n_pairs": len(diffs),
    }


def holm(rows: Sequence[Row], p_key: str = "p_raw") -> List[Row]:
    out = [dict(row) for row in rows]
    out.sort(key=lambda row: fnum(row[p_key]))
    m = len(out)
    for i, row in enumerate(out):
        row["p_holm"] = min(1.0, fnum(row[p_key]) * (m - i))
    return out


def md_table(rows: Sequence[Row], headers: Sequence[str] | None = None) -> str:
    if headers is None:
        headers = list(rows[0].keys()) if rows else []
    lines = ["| " + " | ".join(headers) + " |", "| " + " | ".join("---" for _ in headers) + " |"]
    for row in rows:
        lines.append("| " + " | ".join(str(row.get(h, "")) for h in headers) + " |")
    return "\n".join(lines) + "\n"


def latex(rows: Sequence[Row], headers: Sequence[str]) -> str:
    out = ["\\begin{tabular}{" + "l" + "c" * (len(headers) - 1) + "}", "\\toprule"]
    out.append(" & ".join(headers) + " \\\\")
    out.append("\\midrule")
    for row in rows:
        out.append(" & ".join(str(row.get(h, "")) for h in headers) + " \\\\")
    out.append("\\bottomrule")
    out.append("\\end{tabular}")
    return "\n".join(out) + "\n"


def fmt(value: Any, digits: int = 3) -> str:
    x = fnum(value)
    return "--" if math.isnan(x) else f"{x:.{digits}f}"