"""Probability calibration: metrics and post-hoc recalibration maps.

A classifier is *calibrated* if, among inputs where it says "70% class k",
class k occurs ~70% of the time. We measure the top-label version:
P(Y = argmax p | max p = c) = c.

Expected Calibration Error (ECE) bins examples by confidence and averages
|accuracy - confidence| weighted by bin size. It is a biased, binning-dependent
estimate - use it to compare methods on the same data, not as an absolute truth.

Recalibration maps (all fitted on a held-out split, never the training data):

* Temperature scaling (Guo et al., 2017, "On Calibration of Modern Neural
  Networks"): divide logits by one scalar T > 0 fitted by NLL. One parameter,
  so it cannot overfit much and never changes the argmax.
* Platt scaling: per-class logistic regression on the class logit (one-vs-rest
  for multiclass, then renormalise).
* Isotonic regression: per-class monotone non-parametric map (one-vs-rest,
  renormalised). Flexible; needs more data; can change the argmax after
  renormalisation.
"""

from __future__ import annotations

from abc import ABC, abstractmethod

import numpy as np
from scipy.optimize import minimize_scalar
from scipy.special import log_softmax, softmax
from sklearn.isotonic import IsotonicRegression
from sklearn.linear_model import LogisticRegression

EPS = 1e-12


# ---------------------------------------------------------------- metrics ---
def reliability_bins(probs: np.ndarray, y: np.ndarray, n_bins: int = 15):
    """Equal-width confidence bins. Returns (bin_conf, bin_acc, bin_count, edges)."""
    conf = probs.max(1)
    correct = (probs.argmax(1) == y).astype(float)
    edges = np.linspace(0.0, 1.0, n_bins + 1)
    # Right-closed bins so that confidence 1.0 lands in the last bin.
    which = np.clip(np.searchsorted(edges, conf, side="left") - 1, 0, n_bins - 1)
    count = np.bincount(which, minlength=n_bins).astype(float)
    with np.errstate(invalid="ignore", divide="ignore"):
        bconf = np.bincount(which, weights=conf, minlength=n_bins) / count
        bacc = np.bincount(which, weights=correct, minlength=n_bins) / count
    return bconf, bacc, count, edges


def ece(probs: np.ndarray, y: np.ndarray, n_bins: int = 15) -> float:
    """Top-label expected calibration error (equal-width bins)."""
    bconf, bacc, count, _ = reliability_bins(probs, y, n_bins)
    m = count > 0
    return float(np.sum(count[m] * np.abs(bacc[m] - bconf[m])) / count.sum())


def nll(probs: np.ndarray, y: np.ndarray) -> float:
    """Mean negative log-likelihood of the true class."""
    return float(-np.mean(np.log(np.clip(probs[np.arange(len(y)), y], EPS, 1.0))))


def brier(probs: np.ndarray, y: np.ndarray) -> float:
    """Multiclass Brier score: mean squared error to the one-hot label."""
    onehot = np.zeros_like(probs)
    onehot[np.arange(len(y)), y] = 1.0
    return float(np.mean(np.sum((probs - onehot) ** 2, axis=1)))


def calibration_metrics(probs: np.ndarray, y: np.ndarray) -> dict:
    return {
        "accuracy": float(np.mean(probs.argmax(1) == y)),
        "ece": ece(probs, y),
        "nll": nll(probs, y),
        "brier": brier(probs, y),
        "mean_confidence": float(probs.max(1).mean()),
    }


# ------------------------------------------------------------ calibrators ---
def _logits(probs: np.ndarray) -> np.ndarray:
    """log p is a valid logit vector (softmax is shift-invariant)."""
    return np.log(np.clip(probs, EPS, 1.0))


class IdentityCalibrator:
    """No-op, so 'no calibration' fits the same interface."""

    name = "none"

    def fit(self, probs, y):
        return self

    def transform(self, probs):
        return np.asarray(probs, dtype=float)


class TemperatureScaler:
    """Single-parameter temperature scaling: p_T = softmax(log p / T)."""

    name = "temperature"

    def __init__(self):
        self.T_ = 1.0

    def fit(self, probs: np.ndarray, y: np.ndarray):
        z = _logits(probs)
        idx = np.arange(len(y))

        def loss(log_t):
            return -np.mean(log_softmax(z / np.exp(log_t), axis=1)[idx, y])

        # Search over log T so that T stays positive; [1/20, 20] is plenty.
        res = minimize_scalar(loss, bounds=(np.log(0.05), np.log(20.0)), method="bounded")
        self.T_ = float(np.exp(res.x))
        return self

    def transform(self, probs: np.ndarray) -> np.ndarray:
        return softmax(_logits(probs) / self.T_, axis=1)


class _OneVsRestCalibrator(ABC):
    """Shared logic for per-class 1-D calibration maps followed by renormalisation.

    Subclasses define how to build, fit and apply one 1-D map p_k -> p_k'.
    """

    @abstractmethod
    def _make(self):
        """Return a fresh, unfitted 1-D map."""

    @abstractmethod
    def _fit1(self, model, p, t):
        """Fit ``model`` mapping class-k probabilities ``p`` to 0/1 targets ``t``."""

    @abstractmethod
    def _apply1(self, model, p):
        """Apply a fitted map to class-k probabilities."""

    def fit(self, probs: np.ndarray, y: np.ndarray):
        K = probs.shape[1]
        self.binary_ = K == 2
        cols = [1] if self.binary_ else range(K)
        self.maps_ = {}
        for k in cols:
            target = (y == k).astype(float)
            if target.min() == target.max():  # class absent / only class: skip
                self.maps_[k] = None
                continue
            self.maps_[k] = self._fit1(self._make(), probs[:, k], target)
        return self

    def transform(self, probs: np.ndarray) -> np.ndarray:
        probs = np.asarray(probs, dtype=float)
        if self.binary_:
            p1 = probs[:, 1] if self.maps_[1] is None else self._apply1(self.maps_[1], probs[:, 1])
            p1 = np.clip(p1, 0.0, 1.0)
            return np.column_stack([1 - p1, p1])
        out = np.column_stack(
            [
                probs[:, k] if self.maps_[k] is None else self._apply1(self.maps_[k], probs[:, k])
                for k in range(probs.shape[1])
            ]
        )
        out = np.clip(out, 0.0, None)
        s = out.sum(1, keepdims=True)
        # Rows mapped to all-zeros fall back to the uncalibrated probabilities.
        return np.where(s > 0, out / np.where(s > 0, s, 1.0), probs)


class IsotonicCalibrator(_OneVsRestCalibrator):
    name = "isotonic"

    def _make(self):
        return IsotonicRegression(y_min=0.0, y_max=1.0, out_of_bounds="clip")

    def _fit1(self, model, p, t):
        return model.fit(p, t)

    def _apply1(self, model, p):
        return model.predict(p)


class PlattCalibrator(_OneVsRestCalibrator):
    name = "platt"

    def _make(self):
        return LogisticRegression(C=1e4, max_iter=1000)

    @staticmethod
    def _feat(p):
        p = np.clip(p, 1e-6, 1 - 1e-6)
        return np.log(p / (1 - p)).reshape(-1, 1)

    def _fit1(self, model, p, t):
        return model.fit(self._feat(p), t.astype(int))

    def _apply1(self, model, p):
        return model.predict_proba(self._feat(p))[:, 1]


CALIBRATORS = {
    "none": IdentityCalibrator,
    "temperature": TemperatureScaler,
    "platt": PlattCalibrator,
    "isotonic": IsotonicCalibrator,
}


def compare_calibrators(P_fit, y_fit, P_eval, y_eval, methods=tuple(CALIBRATORS)) -> list[dict]:
    """Fit each calibrator on (P_fit, y_fit) and report metrics on (P_eval, y_eval)."""
    rows = []
    for m in methods:
        cal = CALIBRATORS[m]().fit(P_fit, y_fit)
        row = {"method": m, **calibration_metrics(cal.transform(P_eval), y_eval)}
        if m == "temperature":
            row["T"] = cal.T_
        rows.append(row)
    return rows