"""Selective classification: predict when confident, abstain otherwise.

Rule: answer with argmax p(x) if conf(x) = max_k p_k(x) >= lam, else abstain
("send to a human").

* coverage(lam)       = P(conf >= lam)                (fraction answered)
* selective risk(lam) = P(wrong | conf >= lam)        (error among answered)

Lowering lam answers more but (for a sensible confidence score) makes more
mistakes; the *risk-coverage curve* traces this trade-off.

Choosing lam with a guarantee
-----------------------------
We want the selective risk below a target r* with high confidence. On n
held-out labelled points (never used for training or calibration fitting),
for a fixed lam let n_lam = #{conf >= lam} and k_lam = #errors among them.
Under i.i.d. sampling, conditional on n_lam, k_lam ~ Binomial(n_lam, R(lam)),
so the one-sided Clopper-Pearson bound

    U(lam) = BetaInv(1 - delta', k_lam + 1, n_lam - k_lam)

satisfies P(R(lam) <= U(lam)) >= 1 - delta' for that single lam.

We evaluate a *fixed, data-independent* grid of m thresholds. With
delta' = delta / m (Bonferroni), the bounds hold *simultaneously* for all grid
points with probability >= 1 - delta, so we may pick any of them after
looking - we take the lowest (= most coverage) lam with U(lam) <= r*. This is
the same idea as Geifman & El-Yaniv (2017, "Selective Classification for Deep
Neural Networks"), who use a binary search with a union bound, and as the
"Learn then Test" framework (Angelopoulos et al., 2021).

What IS guaranteed: with probability >= 1 - delta over the draw of the
held-out set, the chosen threshold's true selective risk on future data from
the SAME distribution is <= r*.

What is NOT guaranteed: the error on any particular input; the realised error
rate on a finite test batch (it fluctuates around R(lam)); the coverage
(fraction answered) - that is only estimated; anything under distribution
shift. If no grid threshold certifies r*, we return ``None``: with this much
data the target cannot be certified (collect more data or relax r*).
"""

from __future__ import annotations

import numpy as np
from scipy.stats import beta

# Fixed grid: dense near 1 where high-precision thresholds live. Chosen without
# looking at data, which is what makes the union bound valid.
DEFAULT_GRID = np.unique(
    np.concatenate([np.round(np.arange(0.0, 0.95, 0.025), 4), 1 - np.logspace(-1.3, -3, 15)])
)


def clopper_pearson_upper(k: int, n: int, delta: float) -> float:
    """One-sided (1 - delta) upper confidence bound on a binomial proportion."""
    if n == 0:
        return 1.0
    if k >= n:
        return 1.0
    return float(beta.ppf(1 - delta, k + 1, n - k))


def risk_coverage_curve(conf: np.ndarray, correct: np.ndarray):
    """Empirical (coverage, selective risk, threshold) for every possible cut.

    Sorted from most confident; point i answers the i+1 most confident inputs.
    """
    order = np.argsort(-conf, kind="stable")
    c = conf[order]
    err = 1.0 - correct[order].astype(float)
    n = len(c)
    counts = np.arange(1, n + 1)
    coverage = counts / n
    risk = np.cumsum(err) / counts
    return coverage, risk, c


def aurc(conf: np.ndarray, correct: np.ndarray) -> float:
    """Area under the risk-coverage curve (lower is better)."""
    cov, risk, _ = risk_coverage_curve(conf, correct)
    return float(np.mean(risk))


def certify_thresholds(
    conf: np.ndarray, correct: np.ndarray, delta: float = 0.05, grid: np.ndarray = DEFAULT_GRID
) -> list[dict]:
    """Bonferroni-corrected Clopper-Pearson upper bound on selective risk per grid point."""
    m = len(grid)
    rows = []
    for lam in grid:
        sel = conf >= lam
        n_sel = int(sel.sum())
        k = int((~correct[sel].astype(bool)).sum())
        rows.append(
            {
                "threshold": float(lam),
                "n_selected": n_sel,
                "coverage": n_sel / len(conf),
                "errors": k,
                "empirical_risk": k / n_sel if n_sel else float("nan"),
                "risk_upper_bound": clopper_pearson_upper(k, n_sel, delta / m),
            }
        )
    return rows


def choose_threshold(
    conf: np.ndarray,
    correct: np.ndarray,
    target_risk: float = 0.05,
    delta: float = 0.05,
    grid: np.ndarray = DEFAULT_GRID,
) -> dict:
    """Lowest grid threshold whose simultaneous upper bound is <= ``target_risk``.

    Returns a dict with ``threshold`` (None if nothing certifies) and the
    corresponding row from ``certify_thresholds``.
    """
    rows = certify_thresholds(conf, correct, delta, grid)
    ok = [r for r in rows if r["risk_upper_bound"] <= target_risk]
    best = min(ok, key=lambda r: r["threshold"]) if ok else None
    return {
        "target_risk": target_risk,
        "delta": delta,
        "grid_size": len(grid),
        "n_holdout": int(len(conf)),
        "threshold": None if best is None else best["threshold"],
        "certified": best,
        "table": rows,
    }