code/metrics.py

"""Evaluation metrics (pure Python, standard library only). Spec: docs/classifier-plan.md §5.1.

Labels are option strings; a distribution is a dict option -> probability in taxonomy option order. Conventions:
- **Masked / null gold** is skipped by every per-field metric (``None`` in the gold list).
- **F1 of a class** = 2TP / (2TP + FP + FN); precision is None without predictions, recall None without support.
- **Macro-F1** = unweighted mean F1 over the field's options, leaving out options with no support and no
  prediction (F1 undefined) and the options in ``exclude`` (the taxonomy's ``macro_f1_exclude``; nsb-v2:
  emotion_fine ``mixed``, still in ``per_class``); a missed or only falsely predicted option counts as 0.
- **Cohen's kappa** = (p_o - p_e) / (1 - p_e) over paired non-null labels; 1.0 when p_e == 1 and p_o == 1.
- **Top-label ECE**: 15 equal-width bins on the top-label probability, bin = min(floor(p x 15), 14);
  sum over non-empty bins of (bin size / n) x |accuracy - mean confidence|. Class-wise ECE: the same on p(class).
- **Log-loss**: mean -ln(max(p(gold), 1e-12)).
- **Weighted shares** (§2.3): predicted share = sum(p(class) / inclusion_prob) / sum(1 / inclusion_prob) against
  the teacher share with the same weights; a missing inclusion_prob counts as 1.
- **Bootstrap**: author-cluster percentile bootstrap: sorted clusters, ``random.Random(seed)``, each replicate draws
  as many clusters as there are with replacement; ``fn(ids)`` returning None is skipped; interval
  [v[floor(a/2 (m-1))], v[floor((1-a/2)(m-1))]] over the m sorted kept values.
"""
from __future__ import annotations

import math
import random
from collections import Counter, defaultdict
from typing import Callable, Hashable, Iterable, Mapping, Sequence

VERSION = "nsb-metrics-v1"
FLOOR = 1e-12
BINS = 15
REPS = 1000
LEVEL = 0.95
SEED = "nsb-eval-20261017"


def top(dist: Mapping[str, float]) -> str:
    """Highest-probability option; ties go to the earlier option."""
    best = None
    for k, v in dist.items():
        if best is None or v > dist[best]:
            best = k
    return best


def _pairs(gold: Sequence, other: Sequence) -> list[tuple]:
    gold, other = list(gold), list(other)
    if len(gold) != len(other):
        raise ValueError(f"lengths differ: {len(gold)}, {len(other)}")
    return [(g, o) for g, o in zip(gold, other) if g is not None and o is not None]


def _ratio(a: float, b: float) -> float | None:
    return a / b if b else None


def _bin(p: float, bins: int) -> int:
    return min(max(int(math.floor(p * bins)), 0), bins - 1)


def _ece(points: list[tuple[float, float]], bins: int) -> float | None:
    if not points:
        return None
    total: dict[int, list[float]] = defaultdict(lambda: [0, 0.0, 0.0])
    for conf, hit in points:
        b = total[_bin(conf, bins)]
        b[0] += 1
        b[1] += conf
        b[2] += hit
    n = len(points)
    return sum(c / n * abs(acc / c - s / c) for c, s, acc in total.values())


# ---------------------------------------------------------------------- per field
def accuracy(gold, pred) -> float | None:
    pairs = _pairs(gold, pred)
    return _ratio(sum(1 for g, p in pairs if g == p), len(pairs))


def per_class(gold, pred, classes) -> dict:
    """{class: {support, predicted, tp, precision, recall, f1}} for every given class."""
    classes = list(classes)
    known = set(classes)
    hits, support, predicted = Counter(), Counter(), Counter()
    for g, p in _pairs(gold, pred):
        if g not in known or p not in known:
            raise ValueError(f"label outside the classes: {g!r} / {p!r}")
        support[g] += 1
        predicted[p] += 1
        if g == p:
            hits[g] += 1
    out = {}
    for c in classes:
        tp, s, n = hits[c], support[c], predicted[c]
        out[c] = {"support": s, "predicted": n, "tp": tp, "precision": _ratio(tp, n), "recall": _ratio(tp, s),
                  "f1": _ratio(2 * tp, n + s)}
    return out


def _macro(pc: Mapping[str, dict], exclude: Iterable[str] = ()) -> float | None:
    skip = set(exclude)
    scores = [v["f1"] for c, v in pc.items() if v["f1"] is not None and c not in skip]
    return sum(scores) / len(scores) if scores else None


def macro_f1(gold, pred, classes, exclude: Iterable[str] = ()) -> float | None:
    return _macro(per_class(gold, pred, classes), exclude)


def kappa(a, b) -> float | None:
    """Cohen's kappa of two label lists (pairs with a None on either side are skipped)."""
    pairs = _pairs(a, b)
    n = len(pairs)
    if not n:
        return None
    po = sum(1 for x, y in pairs if x == y) / n
    ca, cb = Counter(x for x, _ in pairs), Counter(y for _, y in pairs)
    pe = sum(ca[k] * cb[k] for k in ca) / (n * n)
    if pe >= 1.0:
        return 1.0 if po >= 1.0 else 0.0
    return (po - pe) / (1.0 - pe)


def log_loss(gold, dists) -> float | None:
    pairs = _pairs(gold, dists)
    if not pairs:
        return None
    return sum(-math.log(max(d[g], FLOOR)) for g, d in pairs) / len(pairs)


def ece(gold, dists, bins: int = BINS) -> float | None:
    """Top-label expected calibration error."""
    points = []
    for g, d in _pairs(gold, dists):
        best = top(d)
        points.append((d[best], 1.0 if best == g else 0.0))
    return _ece(points, bins)


def class_ece(gold, dists, cls: str, bins: int = BINS) -> float | None:
    return _ece([(d[cls], 1.0 if g == cls else 0.0) for g, d in _pairs(gold, dists)], bins)


def weighted_shares(gold, dists, cls: str, weights=None) -> dict:
    """Predicted share (mean p(cls)) and teacher share, both weighted (weights default 1)."""
    gold, dists = list(gold), list(dists)
    weights = [1.0] * len(gold) if weights is None else [1.0 if w is None else float(w) for w in weights]
    num_p = num_g = den = 0.0
    n = 0
    for g, d, w in zip(gold, dists, weights):
        if g is None:
            continue
        num_p += w * d[cls]
        num_g += w * (1.0 if g == cls else 0.0)
        den += w
        n += 1
    pred, obs = _ratio(num_p, den), _ratio(num_g, den)
    return {"n": n, "predicted_share": pred, "teacher_share": obs,
            "difference": None if pred is None else pred - obs}


def inverse_weights(inclusion_probs: Iterable[float | None]) -> list[float]:
    """1 / inclusion_prob (1.0 when missing or not positive)."""
    return [1.0 / p if p else 1.0 for p in inclusion_probs]


def key_class(gold, dists, cls: str, weights=None) -> dict:
    """Precision, recall, F1 (top label), class-wise ECE and weighted predicted vs teacher share of one class."""
    gold, dists = list(gold), list(dists)
    pred = [top(d) for d in dists]
    known = {x for x in gold + pred if x is not None} | {cls}
    out = dict(per_class(gold, pred, sorted(known))[cls], class_ece=class_ece(gold, dists, cls))
    out.update(weighted_shares(gold, dists, cls, weights))
    return out


def field_summary(gold, dists, classes, macro_exclude: Iterable[str] = ()) -> dict:
    """Accuracy, macro-F1 (without ``macro_exclude``), kappa vs teacher, log-loss, top-label ECE and per-class
    figures (every class) of one field."""
    gold, dists = list(gold), list(dists)
    pred = [top(d) for d in dists]
    pc = per_class(gold, pred, classes)
    out = {"n": sum(1 for g in gold if g is not None), "masked": sum(1 for g in gold if g is None),
           "accuracy": accuracy(gold, pred), "macro_f1": _macro(pc, macro_exclude),
           "kappa": kappa(gold, pred), "log_loss": log_loss(gold, dists), "ece": ece(gold, dists),
           "per_class": pc}
    if macro_exclude:
        out["macro_f1_excludes"] = sorted(macro_exclude)
        out["macro_f1_all"] = _macro(pc)
    return out


# ---------------------------------------------------------------------- series
def mae(xs: Sequence[float], ys: Sequence[float]) -> float | None:
    if len(xs) != len(ys):
        raise ValueError("lengths differ")
    return sum(abs(x - y) for x, y in zip(xs, ys)) / len(xs) if xs else None


def pearson(xs: Sequence[float], ys: Sequence[float]) -> float | None:
    """Pearson r; None with fewer than 2 points or a constant series."""
    n = len(xs)
    if n != len(ys):
        raise ValueError("lengths differ")
    if n < 2:
        return None
    mx, my = sum(xs) / n, sum(ys) / n
    sxy = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    sxx = sum((x - mx) ** 2 for x in xs)
    syy = sum((y - my) ** 2 for y in ys)
    if sxx <= 0 or syy <= 0:
        return None
    return sxy / math.sqrt(sxx * syy)


# ---------------------------------------------------------------------- bootstrap
def percentile_interval(values: Sequence[float], level: float = LEVEL) -> list[float] | None:
    if not values:
        return None
    values = sorted(values)
    low = round((1.0 - level) / 2, 12)
    high = round(1.0 - low, 12)
    m = len(values) - 1
    return [values[int(low * m)], values[int(high * m)]]


def bootstrap(ids: Iterable[Hashable], group_of: Mapping | Callable, fn: Callable[[list], float | None],
              reps: int = REPS, seed=SEED, level: float = LEVEL) -> dict:
    """Author-cluster percentile bootstrap of fn(ids). ``group_of``: dict or callable id -> cluster."""
    get = group_of.__getitem__ if isinstance(group_of, Mapping) else group_of
    groups: dict = defaultdict(list)
    for i in ids:
        groups[get(i)].append(i)
    keys, rng, values = sorted(groups), random.Random(str(seed)), []
    for _ in range(reps):
        draw = [i for k in (rng.choice(keys) for _ in keys) for i in groups[k]] if keys else []
        v = fn(draw)
        if v is not None:
            values.append(v)
    return {"interval": percentile_interval(values, level), "level": level, "reps": reps, "reps_used": len(values),
            "groups": len(keys), "seed": str(seed)}