aprender-core 0.71.1

//! Model calibration for confidence estimation.
//!
//! Calibration ensures model confidence matches actual accuracy.
//!
//! # Methods
//! - Temperature Scaling: Single parameter post-hoc calibration
//! - Platt Scaling: Logistic regression on logits
//! - Expected Calibration Error (ECE): Calibration metric

use crate::primitives::Vector;

/// Temperature scaling calibrator.
/// Divides logits by learned temperature T before softmax.
#[derive(Debug, Clone)]
pub struct TemperatureScaling {
    temperature: f32,
}

impl Default for TemperatureScaling {
    fn default() -> Self {
        Self::new()
    }
}

impl TemperatureScaling {
    #[must_use]
    pub fn new() -> Self {
        Self { temperature: 1.0 }
    }

    /// Fit temperature on validation set using grid search.
    pub fn fit(&mut self, logits: &[Vector<f32>], labels: &[usize]) {
        let mut best_temp = 1.0;
        let mut best_nll = f32::INFINITY;

        for t in [0.5, 0.75, 1.0, 1.25, 1.5, 1.75, 2.0, 2.5, 3.0] {
            let nll = Self::compute_nll(logits, labels, t);
            if nll < best_nll {
                best_nll = nll;
                best_temp = t;
            }
        }
        self.temperature = best_temp;
    }

    fn compute_nll(logits: &[Vector<f32>], labels: &[usize], temp: f32) -> f32 {
        let mut total_nll = 0.0;
        for (logit, &label) in logits.iter().zip(labels.iter()) {
            let scaled: Vec<f32> = logit.as_slice().iter().map(|&x| x / temp).collect();
            let probs = softmax(&scaled);
            total_nll -= probs[label].max(1e-10).ln();
        }
        total_nll / logits.len() as f32
    }

    /// Calibrate logits by dividing by temperature.
    #[must_use]
    pub fn calibrate(&self, logits: &Vector<f32>) -> Vector<f32> {
        let scaled: Vec<f32> = logits
            .as_slice()
            .iter()
            .map(|&x| x / self.temperature)
            .collect();
        Vector::from_slice(&scaled)
    }

    /// Get calibrated probabilities.
    #[must_use]
    pub fn predict_proba(&self, logits: &Vector<f32>) -> Vector<f32> {
        let calibrated = self.calibrate(logits);
        Vector::from_slice(&softmax(calibrated.as_slice()))
    }

    #[must_use]
    pub fn temperature(&self) -> f32 {
        self.temperature
    }
}

/// Platt scaling: fits logistic regression on logits.
#[derive(Debug, Clone)]
pub struct PlattScaling {
    a: f32,
    b: f32,
}

impl Default for PlattScaling {
    fn default() -> Self {
        Self::new()
    }
}

impl PlattScaling {
    #[must_use]
    pub fn new() -> Self {
        Self { a: 1.0, b: 0.0 }
    }

    /// Fit Platt scaling parameters using gradient descent.
    pub fn fit(&mut self, logits: &[f32], labels: &[bool]) {
        let mut a = 1.0_f32;
        let mut b = 0.0_f32;
        let lr = 0.01;

        for _ in 0..1000 {
            let mut grad_a = 0.0;
            let mut grad_b = 0.0;

            for (&logit, &label) in logits.iter().zip(labels.iter()) {
                let p = sigmoid(a * logit + b);
                let y = if label { 1.0 } else { 0.0 };
                let diff = p - y;
                grad_a += diff * logit;
                grad_b += diff;
            }

            a -= lr * grad_a / logits.len() as f32;
            b -= lr * grad_b / logits.len() as f32;
        }

        self.a = a;
        self.b = b;
    }

    /// Get calibrated probability.
    #[must_use]
    pub fn predict_proba(&self, logit: f32) -> f32 {
        sigmoid(self.a * logit + self.b)
    }

    #[must_use]
    pub fn params(&self) -> (f32, f32) {
        (self.a, self.b)
    }
}

/// Expected Calibration Error (ECE) - measures calibration quality.
#[provable_contracts_macros::contract("calibration-v1", equation = "expected_calibration_error")]
#[must_use]
pub fn expected_calibration_error(predictions: &[f32], labels: &[bool], n_bins: usize) -> f32 {
    contract_pre_expected_calibration_error!(predictions);
    ece_over(
        predictions
            .iter()
            .zip(labels.iter())
            .map(|(&pred, &label)| (pred, label)),
        predictions.len(),
        n_bins,
    )
}

/// The equal-width binning and weighted `|acc - conf|` reduction shared by both ECE forms.
///
/// Extracted so the binning rule — in particular the `.min(n_bins - 1)` clamp, which is what
/// makes `conf == 1.0` representable rather than out of range — exists ONCE. The binary and
/// multiclass entry points differ only in how they derive `(confidence, was_correct)` from
/// their inputs; that mapping stays at the call site, where the D-07 distinction between
/// positive-class probability and max-class probability is visible. A bin-edge change or a
/// clamp fix now lands in one body instead of two that must be kept in sync by hand.
///
/// `n` is the sample count used as the reduction denominator, passed in because the callers
/// derive it differently (slice length vs. validated row count).
fn ece_over(pairs: impl Iterator<Item = (f32, bool)>, n: usize, n_bins: usize) -> f32 {
    let mut bin_sums = vec![0.0_f32; n_bins];
    let mut bin_correct = vec![0.0_f32; n_bins];
    let mut bin_counts = vec![0usize; n_bins];

    for (conf, was_correct) in pairs {
        let bin = ((conf * n_bins as f32) as usize).min(n_bins - 1);
        bin_sums[bin] += conf;
        bin_correct[bin] += if was_correct { 1.0 } else { 0.0 };
        bin_counts[bin] += 1;
    }

    let n = n as f32;
    let mut ece = 0.0;

    for i in 0..n_bins {
        if bin_counts[i] > 0 {
            let avg_conf = bin_sums[i] / bin_counts[i] as f32;
            let avg_acc = bin_correct[i] / bin_counts[i] as f32;
            ece += (bin_counts[i] as f32 / n) * (avg_conf - avg_acc).abs();
        }
    }
    ece
}

/// Validate the shared multiclass preconditions and return the sample count.
///
/// Structural violations are hard `assert!`s rather than `debug_assert!`s: a label out of
/// range or a ragged probability matrix would otherwise index into the wrong row and
/// return a plausible-looking number, and these two metrics feed published claims (D-07).
/// The finiteness/non-emptiness preconditions declared in `calibration-v1` are checked by
/// the caller through the contract macro, exactly as the binary pair does it.
fn multiclass_rows(probabilities: &[f32], n_classes: usize, labels: &[usize]) -> usize {
    assert!(
        !probabilities.is_empty(),
        "multiclass calibration: probabilities are empty"
    );
    assert!(
        n_classes >= 2,
        "multiclass calibration: n_classes must be >= 2, got {n_classes}"
    );
    assert!(
        probabilities.len() % n_classes == 0,
        "multiclass calibration: {} probabilities do not divide into rows of {n_classes}",
        probabilities.len()
    );
    let n_samples = probabilities.len() / n_classes;
    assert!(
        labels.len() == n_samples,
        "multiclass calibration: {} labels for {n_samples} rows",
        labels.len()
    );
    for (i, &label) in labels.iter().enumerate() {
        assert!(
            label < n_classes,
            "multiclass calibration: label {label} at row {i} is outside 0..{n_classes}"
        );
    }
    for i in 0..n_samples {
        let sum: f32 = probabilities[i * n_classes..(i + 1) * n_classes]
            .iter()
            .sum();
        assert!(
            (sum - 1.0).abs() < 1e-3,
            "multiclass calibration: row {i} sums to {sum}, not 1"
        );
    }
    n_samples
}

/// Top confidence and predicted class of one row: `(max_k p_k, argmax_k p_k)`.
///
/// Ties resolve to the LOWEST class index, matching `numpy.argmax` — the reference
/// fixtures' `uniform_uncertain` case depends on that tie-break.
fn top_label(row: &[f32]) -> (f32, usize) {
    let mut best_idx = 0;
    let mut best = row[0];
    for (k, &p) in row.iter().enumerate().skip(1) {
        if p > best {
            best = p;
            best_idx = k;
        }
    }
    (best, best_idx)
}

/// Top-label Expected Calibration Error for multiclass predictions (D-07).
///
/// The multiclass companion to [`expected_calibration_error`], bound to explicit ordered
/// labels: `probabilities` is row-major with `n_classes` entries per row, and `labels[i]`
/// is the true class index of row `i`.
///
/// ```text
/// conf_i = max_k p_ik
/// pred_i = argmax_k p_ik
/// bin(i) = min(floor(conf_i * n_bins), n_bins - 1)
/// ECE    = Σ_b (n_b / N) |acc_b - conf_b|
/// ```
///
/// Same equal-width binning and the same `.min(n_bins - 1)` clamp as the binary function,
/// so a saturated row (`conf = 1`, giving `conf * n_bins == n_bins`) lands in the top bin
/// rather than out of range.
///
/// This is a GENERALISATION of the binary metric, not a reduction of it: the binary form
/// bins by the positive class's probability, this one by the maximum over classes. The two
/// therefore disagree at K = 2, on purpose.
///
/// # Panics
///
/// Panics if `probabilities` is empty, `n_classes < 2`, `n_bins == 0`, the probability
/// count is not a multiple of `n_classes`, the label count does not match the row count,
/// any label is `>= n_classes`, or any row does not sum to 1 within 1e-3.
///
/// # Reference
///
/// Guo et al. (2017), *On Calibration of Modern Neural Networks*, §3. Reference values are
/// frozen in `scripts/setfit_fixtures/claims_stats/ece_top_label_cases.json`.
#[provable_contracts_macros::contract(
    "calibration-v1",
    equation = "expected_calibration_error_top_label"
)]
#[must_use]
pub fn expected_calibration_error_top_label(
    probabilities: &[f32],
    n_classes: usize,
    labels: &[usize],
    n_bins: usize,
) -> f32 {
    contract_pre_expected_calibration_error_top_label!(probabilities);
    assert!(n_bins > 0, "multiclass calibration: n_bins must be >= 1");
    let n_samples = multiclass_rows(probabilities, n_classes, labels);

    ece_over(
        labels.iter().enumerate().map(|(i, &label)| {
            let row = &probabilities[i * n_classes..(i + 1) * n_classes];
            let (conf, pred) = top_label(row);
            (conf, pred == label)
        }),
        n_samples,
        n_bins,
    )
}

/// Top-label ECE in f64 with exactly-rounded sums — the Laya gate's `ece_top_label`
/// (contracts/laya-finetune-gate-v1.yaml `numeric_agreement.ece_top_label`).
///
/// ```text
/// conf_i, pred_i = max_k p_ik, argmax_k p_ik        (first max on ties; conf widened f32 -> f64)
/// bin(i)         = min(floor(conf_i x n_bins), n_bins - 1)          in f64
/// acc_b          = fsum(correct_i for i in b) / n_b
/// conf_b         = fsum(conf_i for i in b) / n_b
/// ECE            = fsum over occupied bins, in bin order, of (n_b / N) x |acc_b - conf_b|
/// ```
///
/// Bit-identical to `scripts/laya_train/metrics.py` `ece_top_label`, which computes the same
/// expression with `math.fsum` on the f64-widened probabilities. It shares the input validation
/// and the top-label rule with [`expected_calibration_error_top_label`], which it does NOT
/// replace: that function's f32 results are frozen for its other callers.
///
/// # Panics
///
/// As [`expected_calibration_error_top_label`].
#[must_use]
pub fn expected_calibration_error_top_label_f64(
    probabilities: &[f32],
    n_classes: usize,
    labels: &[usize],
    n_bins: usize,
) -> f64 {
    contract_pre_expected_calibration_error_top_label!(probabilities);
    assert!(n_bins > 0, "multiclass calibration: n_bins must be >= 1");
    let n_samples = multiclass_rows(probabilities, n_classes, labels);
    let mut conf_in: Vec<Vec<f64>> = vec![Vec::new(); n_bins];
    let mut correct_in: Vec<usize> = vec![0; n_bins];
    for (i, &label) in labels.iter().enumerate() {
        let (conf, pred) = top_label(&probabilities[i * n_classes..(i + 1) * n_classes]);
        let conf = f64::from(conf);
        let bin = ((conf * n_bins as f64).floor() as usize).min(n_bins - 1);
        conf_in[bin].push(conf);
        if pred == label {
            correct_in[bin] += 1;
        }
    }
    let n = n_samples as f64;
    let terms = conf_in
        .iter()
        .zip(&correct_in)
        .filter(|(confs, _)| !confs.is_empty())
        .map(|(confs, &correct)| {
            let n_b = confs.len() as f64;
            // fsum of `correct` 1.0s is the integer count, exactly: the cast IS that sum.
            let acc = correct as f64 / n_b;
            let mean_conf = crate::metrics::fsum(confs.iter().copied()) / n_b;
            (n_b / n) * (acc - mean_conf).abs()
        })
        .collect::<Vec<f64>>();
    crate::metrics::fsum(terms)
}

/// Multiclass Brier score, original (Brier 1950) UNNORMALISED definition (D-07).
///
/// ```text
/// BS = (1/N) Σ_i Σ_k (p_ik - y_ik)²      with one-hot y
/// ```
///
/// # Normalization — stated once, here, so nobody has to re-derive it
///
/// The class sum is **not** divided by `K`. Consequences, both contracted in
/// `calibration-v1` and both asserted by the tests:
///
/// - The codomain is `[0, 2]` for every `K >= 2`, **not** `[0, 1]`. A confidently-wrong
///   one-hot row contributes `(0-1)² + (1-0)² = 2`.
/// - At `K = 2` with `p_i = [1-p, p]`, the class sum is
///   `(p-y)² + ((1-p)-(1-y))² = 2(p-y)²`, so this function returns **exactly twice**
///   [`brier_score`] — never the same value. An implementation that made the two equal
///   would have silently divided by `K`.
///
/// The one-vs-rest identity `BS = Σ_k binary_brier(labels == k, p_·k)` holds and is
/// verified against scikit-learn 1.9.0 in the fixture generator.
///
/// Lower is better. `BS = 0` exactly when every row is a correct one-hot prediction.
///
/// # Panics
///
/// Panics if `probabilities` is empty, `n_classes < 2`, the probability count is not a
/// multiple of `n_classes`, the label count does not match the row count, any label is
/// `>= n_classes`, or any row does not sum to 1 within 1e-3.
///
/// # Reference
///
/// Brier (1950), *Verification of Forecasts Expressed in Terms of Probability*. Reference
/// values are frozen in `scripts/setfit_fixtures/claims_stats/brier_multiclass_cases.json`.
#[provable_contracts_macros::contract("calibration-v1", equation = "brier_score_multiclass")]
#[must_use]
pub fn brier_score_multiclass(probabilities: &[f32], n_classes: usize, labels: &[usize]) -> f32 {
    contract_pre_brier_score_multiclass!(probabilities);
    let n_samples = multiclass_rows(probabilities, n_classes, labels);

    let mut total = 0.0_f32;
    for (i, &label) in labels.iter().enumerate() {
        let row = &probabilities[i * n_classes..(i + 1) * n_classes];
        for (k, &p) in row.iter().enumerate() {
            let y = if k == label { 1.0 } else { 0.0 };
            total += (p - y).powi(2);
        }
    }
    total / n_samples as f32
}

/// Maximum Calibration Error (MCE).
#[provable_contracts_macros::contract("calibration-v1", equation = "maximum_calibration_error")]
#[must_use]
pub fn maximum_calibration_error(predictions: &[f32], labels: &[bool], n_bins: usize) -> f32 {
    contract_pre_maximum_calibration_error!(predictions);
    let mut bin_sums = vec![0.0; n_bins];
    let mut bin_correct = vec![0.0; n_bins];
    let mut bin_counts = vec![0usize; n_bins];

    for (&pred, &label) in predictions.iter().zip(labels.iter()) {
        let bin = ((pred * n_bins as f32) as usize).min(n_bins - 1);
        bin_sums[bin] += pred;
        bin_correct[bin] += if label { 1.0 } else { 0.0 };
        bin_counts[bin] += 1;
    }

    let mut mce = 0.0_f32;
    for i in 0..n_bins {
        if bin_counts[i] > 0 {
            let avg_conf = bin_sums[i] / bin_counts[i] as f32;
            let avg_acc = bin_correct[i] / bin_counts[i] as f32;
            mce = mce.max((avg_conf - avg_acc).abs());
        }
    }
    mce
}

/// ONE PATH: Delegates to `nn::functional::softmax_1d` (UCBD §4).
fn softmax(logits: &[f32]) -> Vec<f32> {
    crate::nn::functional::softmax_1d(logits)
}

/// ONE PATH: Delegates to `nn::functional::sigmoid_scalar` (UCBD §4).
fn sigmoid(x: f32) -> f32 {
    crate::nn::functional::sigmoid_scalar(x)
}

/// Isotonic Regression calibrator.
///
/// Non-parametric calibration method that fits a monotonically increasing
/// step function to the data. Guaranteed to be monotonic.
///
/// # Reference
///
/// - Zadrozny, B., & Elkan, C. (2002). Transforming classifier scores into
///   accurate multiclass probability estimates.
#[derive(Debug, Clone)]
pub struct IsotonicRegression {
    /// Threshold values (sorted)
    pub thresholds: Vec<f32>,
    /// Calibrated values at each threshold
    pub values: Vec<f32>,
}

impl Default for IsotonicRegression {
    fn default() -> Self {
        Self::new()
    }
}

impl IsotonicRegression {
    /// Create a new isotonic regression calibrator.
    #[must_use]
    pub fn new() -> Self {
        Self {
            thresholds: Vec::new(),
            values: Vec::new(),
        }
    }

    /// Fit isotonic regression on predictions and labels.
    ///
    /// Uses Pool Adjacent Violators (PAV) algorithm.
    pub fn fit(&mut self, predictions: &[f32], labels: &[bool]) {
        assert_eq!(predictions.len(), labels.len());

        // Sort by predictions
        let mut pairs: Vec<(f32, f32)> = predictions
            .iter()
            .zip(labels.iter())
            .map(|(&p, &l)| (p, if l { 1.0 } else { 0.0 }))
            .collect();
        pairs.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));

        if pairs.is_empty() {
            return;
        }

        // PAV algorithm: merge adjacent violators.
        // Track BOTH endpoints of each pooled block (x_min, x_max) so the fitted
        // function is FLAT across the whole x-range of a pooled block (sklearn
        // IsotonicRegression parity, PMAT-870). Recording only x_min would leave a
        // pooled block as a single knot and cause interior queries to interpolate
        // toward the next block instead of returning the constant pooled value.
        let mut blocks: Vec<(f32, f32, f32, usize)> = Vec::new(); // (x_min, x_max, sum, count)

        for (x, y) in pairs {
            blocks.push((x, x, y, 1));

            // While last two blocks violate monotonicity, merge them
            while blocks.len() >= 2 {
                let len = blocks.len();
                let avg_last = blocks[len - 1].2 / blocks[len - 1].3 as f32;
                let avg_prev = blocks[len - 2].2 / blocks[len - 2].3 as f32;

                if avg_last < avg_prev {
                    // Merge: keep x_min from prev, x_max from last, sum totals.
                    // Safe: we check blocks.len() >= 2 in the while condition
                    let last = blocks
                        .pop()
                        .expect("blocks guaranteed non-empty by while condition");
                    let prev = blocks
                        .pop()
                        .expect("blocks guaranteed to have 2+ elements by while condition");
                    blocks.push((prev.0, last.1, prev.2 + last.2, prev.3 + last.3));
                } else {
                    break;
                }
            }
        }

        // Extract thresholds and values. For a pooled block spanning [x_min, x_max]
        // with x_max > x_min, emit TWO knots at the same value so predict()
        // interpolates to the constant across the block interior; degenerate
        // single-point blocks emit one knot.
        self.thresholds.clear();
        self.values.clear();

        for (x_min, x_max, sum, count) in blocks {
            let val = sum / count as f32;
            self.thresholds.push(x_min);
            self.values.push(val);
            if x_max > x_min {
                self.thresholds.push(x_max);
                self.values.push(val);
            }
        }
    }

    /// Get calibrated probability for a prediction.
    #[must_use]
    pub fn predict(&self, prediction: f32) -> f32 {
        if self.thresholds.is_empty() {
            return prediction;
        }

        // Binary search to find appropriate block
        let idx = self.thresholds.partition_point(|&t| t <= prediction);

        if idx == 0 {
            self.values[0]
        } else if idx >= self.values.len() {
            self.values[self.values.len() - 1]
        } else {
            // Linear interpolation between blocks
            let t0 = self.thresholds[idx - 1];
            let t1 = self.thresholds[idx];
            let v0 = self.values[idx - 1];
            let v1 = self.values[idx];

            if (t1 - t0).abs() < 1e-10 {
                v0
            } else {
                let alpha = (prediction - t0) / (t1 - t0);
                v0 + alpha * (v1 - v0)
            }
        }
    }

    /// Get calibrated probabilities for multiple predictions.
    #[must_use]
    pub fn predict_batch(&self, predictions: &[f32]) -> Vec<f32> {
        predictions.iter().map(|&p| self.predict(p)).collect()
    }
}

/// Generate reliability diagram data.
///
/// Returns (`mean_predicted_prob`, `fraction_positive`) for each bin.
#[must_use]
pub fn reliability_diagram(predictions: &[f32], labels: &[bool], n_bins: usize) -> Vec<(f32, f32)> {
    contract_pre_reliability_diagram!(predictions);
    let mut bins: Vec<(f32, f32, usize)> = vec![(0.0, 0.0, 0); n_bins]; // (sum_pred, sum_pos, count)

    for (&pred, &label) in predictions.iter().zip(labels.iter()) {
        let bin = ((pred * n_bins as f32) as usize).min(n_bins - 1);
        bins[bin].0 += pred;
        bins[bin].1 += if label { 1.0 } else { 0.0 };
        bins[bin].2 += 1;
    }

    bins.iter()
        .enumerate()
        .map(|(i, &(sum_pred, sum_pos, count))| {
            if count > 0 {
                (sum_pred / count as f32, sum_pos / count as f32)
            } else {
                // Empty bin: use bin center, 0.0 accuracy
                let center = (i as f32 + 0.5) / n_bins as f32;
                (center, 0.0)
            }
        })
        .collect()
}

/// Brier score - measures calibration quality.
///
/// ```text
/// Brier = (1/n) * Σ(p_i - y_i)²
/// ```
///
/// Lower is better. Perfect calibration = 0.
#[must_use]
pub fn brier_score(predictions: &[f32], labels: &[bool]) -> f32 {
    let n = predictions.len() as f32;
    predictions
        .iter()
        .zip(labels.iter())
        .map(|(&p, &l)| {
            let y = if l { 1.0 } else { 0.0 };
            (p - y).powi(2)
        })
        .sum::<f32>()
        / n
}

#[cfg(test)]
#[path = "calibration_tests.rs"]
mod tests;

#[cfg(test)]
#[path = "calibration_tests_contract.rs"]
mod tests_cal_contract;