"""Calibration: temperature scaling and split-conformal abstention."""

from __future__ import annotations

import math
from dataclasses import dataclass

import numpy as np
import torch
import torch.nn.functional as F


def fit_temperature(logits: np.ndarray, targets: np.ndarray) -> float:
    """Fit one temperature by minimizing NLL on held-out logits."""
    x = torch.from_numpy(logits.astype(np.float32))
    y = torch.from_numpy(targets.astype(np.float32))
    log_t = torch.zeros((), requires_grad=True)
    opt = torch.optim.LBFGS([log_t], lr=0.1, max_iter=100)

    def closure():
        opt.zero_grad()
        t = log_t.exp().clamp(0.02, 50.0)
        loss = -(y * F.log_softmax(x / t, dim=-1)).sum(-1).mean()
        loss.backward()
        return loss

    opt.step(closure)
    return float(log_t.exp().clamp(0.02, 50.0).item())


def multiclass_ece(probs: np.ndarray, labels: np.ndarray, n_bins: int = 10) -> float:
    conf = probs.max(-1)
    pred = probs.argmax(-1)
    correct = (pred == labels).astype(np.float64)
    bins = np.linspace(0.0, 1.0, n_bins + 1)
    ece = 0.0
    for lo, hi in zip(bins[:-1], bins[1:]):
        m = (conf > lo) & (conf <= hi)
        if m.any():
            ece += m.mean() * abs(correct[m].mean() - conf[m].mean())
    return float(ece)


def brier(probs: np.ndarray, labels: np.ndarray) -> float:
    onehot = np.zeros_like(probs)
    onehot[np.arange(len(labels)), labels] = 1.0
    return float(((probs - onehot) ** 2).sum(-1).mean())


@dataclass
class ConformalGate:
    """Selective predictor with an empirically risk-controlled threshold.

    The gate accepts a prediction only when its top probability is at least
    ``threshold``. The threshold is the lowest confidence cut whose accepted
    set on held-out calibration data has a 95% Wilson upper bound on error at
    or below ``alpha``. That gives the widest coverage consistent with the
    requested risk on the calibration distribution.
    """

    threshold: float
    alpha: float
    coverage: float
    risk: float
    n: int

    @staticmethod
    def _wilson_upper(errors: int, total: int, z: float = 1.96) -> float:
        if total <= 0:
            return 1.0
        p = errors / total
        denom = 1.0 + z * z / total
        centre = (p + z * z / (2 * total)) / denom
        half = z * ((p * (1 - p) / total + z * z / (4 * total * total)) ** 0.5) / denom
        return min(1.0, centre + half)

    @classmethod
    def fit(
        cls,
        probs: np.ndarray,
        labels: np.ndarray,
        alpha: float = 0.05,
    ) -> "ConformalGate":
        pred = probs.argmax(-1)
        return cls.fit_scores(
            probs.max(-1).astype(np.float64),
            (pred == labels),
            alpha=alpha,
        )

    @classmethod
    def fit_scores(
        cls,
        scores: np.ndarray,
        correct: np.ndarray,
        alpha: float = 0.05,
    ) -> "ConformalGate":
        conf = np.asarray(scores, dtype=np.float64)
        wrong = (~np.asarray(correct, dtype=bool)).astype(np.int64)
        n = len(conf)
        order = np.argsort(conf, kind="stable")
        conf_s = conf[order]
        wrong_s = wrong[order]
        suffix_errors = np.concatenate([np.cumsum(wrong_s[::-1])[::-1], [0]])
        best_t = 1.0
        best_cov = 0.0
        best_risk = 0.0
        best_n = 0
        min_accept = max(30, int(0.01 * n))
        for i in range(n):
            accepted = n - i
            if accepted < min_accept:
                break
            t = float(conf_s[i])
            if t <= 0.0:
                continue
            errors = int(suffix_errors[i])
            upper = cls._wilson_upper(errors, accepted)
            if upper <= alpha:
                best_t = t
                best_cov = accepted / n
                best_risk = errors / accepted
                best_n = accepted
                break
        return cls(
            threshold=best_t,
            alpha=alpha,
            coverage=float(best_cov),
            risk=float(best_risk),
            n=int(best_n),
        )

    def accepts(self, prob: float) -> bool:
        return float(prob) >= self.threshold

    def to_dict(self) -> dict:
        return {
            "threshold": self.threshold,
            "alpha": self.alpha,
            "coverage": self.coverage,
            "risk": self.risk,
            "n": self.n,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "ConformalGate":
        return cls(**{k: d[k] for k in ("threshold", "alpha", "coverage", "risk", "n")})


def selective_report(
    probs: np.ndarray, labels: np.ndarray, thresholds: list[float]
) -> list[dict]:
    conf = probs.max(-1)
    pred = probs.argmax(-1)
    wrong = pred != labels
    out = []
    for t in thresholds:
        accepted = conf >= t
        out.append(
            {
                "threshold": float(t),
                "coverage": float(accepted.mean()),
                "risk": float(wrong[accepted].mean()) if accepted.any() else None,
                "accepted": int(accepted.sum()),
            }
        )
    return out