"""Small, dependency-free statistics for evaluating a detector against human labels."""
from __future__ import annotations

import math
from collections.abc import Sequence
from dataclasses import dataclass


def wilson(successes: int, total: int, z: float = 1.96) -> tuple[float, float]:
    """Wilson score interval for a proportion; (0, 1) when there is no data."""
    if total <= 0:
        return 0.0, 1.0
    p = successes / total
    denom = 1 + z * z / total
    centre = (p + z * z / (2 * total)) / denom
    half = z * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total)) / denom
    return max(0.0, centre - half), min(1.0, centre + half)


@dataclass
class Confusion:
    """``positive`` means 'the item is unfaithful' and a detector 'flags' it."""

    tp: int = 0
    fp: int = 0
    fn: int = 0
    tn: int = 0

    @property
    def n(self) -> int:
        return self.tp + self.fp + self.fn + self.tn

    @property
    def recall(self) -> float:
        return self.tp / (self.tp + self.fn) if (self.tp + self.fn) else float("nan")

    @property
    def precision(self) -> float:
        return self.tp / (self.tp + self.fp) if (self.tp + self.fp) else float("nan")

    @property
    def false_flag_rate(self) -> float:
        return self.fp / (self.fp + self.tn) if (self.fp + self.tn) else float("nan")

    @property
    def f1(self) -> float:
        p, r = self.precision, self.recall
        return 2 * p * r / (p + r) if (p == p and r == r and (p + r) > 0) else float("nan")

    def as_dict(self) -> dict[str, float | int]:
        return {
            "tp": self.tp, "fp": self.fp, "fn": self.fn, "tn": self.tn, "n": self.n,
            "recall": self.recall, "precision": self.precision,
            "false_flag_rate": self.false_flag_rate, "f1": self.f1,
            "recall_ci": wilson(self.tp, self.tp + self.fn),
            "false_flag_ci": wilson(self.fp, self.fp + self.tn),
        }


def confusion(flagged: Sequence[bool], unfaithful: Sequence[bool]) -> Confusion:
    if len(flagged) != len(unfaithful):
        raise ValueError("flagged and unfaithful must have the same length")
    c = Confusion()
    for f, u in zip(flagged, unfaithful, strict=True):
        if f and u:
            c.tp += 1
        elif f and not u:
            c.fp += 1
        elif u:
            c.fn += 1
        else:
            c.tn += 1
    return c


def auc(scores: Sequence[float], unfaithful: Sequence[bool]) -> float:
    """ROC-AUC of 'higher score means unfaithful' (rank-sum form, ties count half)."""
    pos = [s for s, u in zip(scores, unfaithful, strict=True) if u]
    neg = [s for s, u in zip(scores, unfaithful, strict=True) if not u]
    if not pos or not neg:
        return float("nan")
    wins = sum((p > n) + 0.5 * (p == n) for p in pos for n in neg)
    return wins / (len(pos) * len(neg))


def cohen_kappa(a: Sequence[str], b: Sequence[str]) -> float:
    if len(a) != len(b) or not a:
        raise ValueError("need two label lists of equal, non-zero length")
    n = len(a)
    observed = sum(x == y for x, y in zip(a, b, strict=True)) / n
    labels = set(a) | set(b)
    expected = sum((a.count(k) / n) * (b.count(k) / n) for k in labels)
    return 1.0 if expected == 1 else (observed - expected) / (1 - expected)