"""Split conformal prediction sets for classification.

Setting. We have a fitted classifier p(.|x), and n calibration points
(x_i, y_i) that the classifier never saw. Pick a *non-conformity score*
s(x, y) (large = "y looks wrong for x"). Let q be the ceil((n+1)(1-alpha))-th
smallest calibration score. The prediction set for a new x is

    C(x) = { y : s(x, y) <= q }.

Guarantee (Vovk, Gammerman & Shafer, "Algorithmic Learning in a Random World",
2005; split/inductive version e.g. Papadopoulos et al. 2002, Lei et al. 2018):
if calibration and test points are exchangeable (e.g. i.i.d. from the same
distribution) then

    P( Y_test in C(X_test) ) >= 1 - alpha,

*marginally* - averaged over the calibration draw and the test point. With
continuous scores (no ties; we randomise APS to ensure this) coverage is also
at most 1 - alpha + 1/(n+1). It holds for ANY model, however bad: a bad model
just produces bigger sets.

Scores implemented:

* LAC ("least ambiguous set-valued classifier", Sadinle, Lei & Wasserman 2019):
  s = 1 - p_y. Gives the smallest average sets, but coverage tends to be
  uneven across classes/difficulty (over-covers easy, under-covers hard).
* APS (adaptive prediction sets, Romano, Sesia & Candes, NeurIPS 2020):
  s = total probability of classes ranked above y, plus U * p_y with
  U ~ Uniform(0,1). Sets grow on hard inputs, shrink on easy ones; better
  conditional behaviour, larger sets on average. (RAPS, Angelopoulos et al.
  2021, adds a size penalty; not implemented here.)

Marginal coverage says nothing about any specific class. *Mondrian* (class-
conditional) conformal computes a separate threshold q_k from the calibration
points with label k and includes class k iff s(x, k) <= q_k. Then
P(Y in C(X) | Y = k) >= 1 - alpha for every class k - at the cost of needing
enough calibration examples per class (with n_k points the guarantee is
vacuous - q_k = inf - unless n_k >= 1/alpha - 1).
"""

from __future__ import annotations

import math

import numpy as np

SCORES = ("lac", "aps")


def conformal_quantile(scores: np.ndarray, alpha: float) -> float:
    """The finite-sample-corrected (1-alpha) quantile used by split conformal.

    Returns +inf when n is too small for the requested alpha
    (ceil((n+1)(1-alpha)) > n), i.e. "include everything".
    """
    n = len(scores)
    k = math.ceil((n + 1) * (1 - alpha))
    if n == 0 or k > n:
        return float("inf")
    return float(np.sort(scores)[k - 1])


def all_class_scores(probs: np.ndarray, score: str, u: np.ndarray | None = None) -> np.ndarray:
    """Non-conformity score of every (example, class) pair, shape (n, K).

    ``u`` holds one Uniform(0,1) draw per example (APS randomisation).
    """
    probs = np.asarray(probs, dtype=float)
    if score == "lac":
        return 1.0 - probs
    if score == "aps":
        if u is None:
            raise ValueError("APS needs the per-example uniform draws u")
        n, K = probs.shape
        order = np.argsort(-probs, axis=1, kind="stable")
        sorted_p = np.take_along_axis(probs, order, axis=1)
        cum_before = np.cumsum(sorted_p, axis=1) - sorted_p  # mass ranked strictly above
        s_sorted = cum_before + u[:, None] * sorted_p
        out = np.empty_like(probs)
        np.put_along_axis(out, order, s_sorted, axis=1)
        return out
    raise ValueError(f"unknown score {score!r}; choose from {SCORES}")


class SplitConformal:
    """Split conformal classifier on top of any probability vector.

    Parameters
    ----------
    alpha : target miscoverage (0.1 -> 90% coverage).
    score : "lac" or "aps".
    mondrian : if True, one threshold per class (class-conditional coverage).
    seed : RNG seed for the APS tie-breaking uniforms.
    """

    def __init__(self, alpha: float = 0.1, score: str = "aps", mondrian: bool = False, seed: int = 0):
        if not 0 < alpha < 1:
            raise ValueError("alpha must be in (0, 1)")
        self.alpha, self.score, self.mondrian, self.seed = alpha, score, mondrian, seed

    def fit(self, probs_cal: np.ndarray, y_cal: np.ndarray):
        rng = np.random.default_rng(self.seed)
        n, K = probs_cal.shape
        u = rng.uniform(size=n)
        s = all_class_scores(probs_cal, self.score, u)[np.arange(n), y_cal]
        self.n_cal_ = n
        self.n_classes_ = K
        if self.mondrian:
            self.qhat_ = np.array(
                [conformal_quantile(s[y_cal == k], self.alpha) for k in range(K)]
            )
        else:
            self.qhat_ = np.full(K, conformal_quantile(s, self.alpha))
        # Fresh stream for test-time randomisation, distinct from calibration.
        self._rng_test = np.random.default_rng(self.seed + 7919)
        return self

    def predict_sets(self, probs: np.ndarray, u: np.ndarray | None = None) -> np.ndarray:
        """Boolean matrix (n, K): True where the class is in the prediction set.

        Sets may be empty (the input looks unlike every class at this alpha);
        treat an empty set as a signal to abstain.
        """
        if u is None:
            u = self._rng_test.uniform(size=len(probs))
        S = all_class_scores(probs, self.score, u)
        return S <= self.qhat_[None, :]


def evaluate_sets(sets: np.ndarray, y: np.ndarray) -> dict:
    """Empirical coverage, set sizes, and per-class coverage."""
    n, K = sets.shape
    covered = sets[np.arange(n), y]
    sizes = sets.sum(1)
    per_class = {}
    for k in range(K):
        m = y == k
        per_class[k] = float(covered[m].mean()) if m.any() else float("nan")
    return {
        "coverage": float(covered.mean()),
        "avg_set_size": float(sizes.mean()),
        "frac_singleton": float(np.mean(sizes == 1)),
        "frac_empty": float(np.mean(sizes == 0)),
        "size_hist": np.bincount(sizes, minlength=K + 1).tolist(),
        "per_class_coverage": per_class,
        "worst_class_coverage": float(np.nanmin(list(per_class.values()))),
    }


def coverage_upper_bound(n_cal: int, alpha: float) -> float:
    """Upper end of the marginal coverage interval for continuous scores."""
    return min(1.0, 1 - alpha + 1.0 / (n_cal + 1))