//! Judge calibration math.
//!
//!  hands every verdict a [`ConfidenceBand`], but a band is a
//! relabel of inter-juror agreement: it says nothing about whether the
//! jury's stated confidence tracks its actual accuracy, nor whether an
//! imperfect judge can gate a PR at all. This module fills that gap with
//! three deterministic, transport-free pieces:
//!
//! 1. **Per-verdict calibration** ([`calibration_error`], [`brier_score`]):
//!    how well stated confidence matches observed correctness on a labeled
//!    set. Lower is better. Motivated by overconfidence in LLM judges
//!    (arXiv:2508.06225), where confidence routinely overstates accuracy.
//! 2. **Small-trusted-set validation** ([`estimate_reliability`],
//!    [`rogan_gladen`], [`corrected_rate_ci`]): a handful of human labels
//!    measure a judge's sensitivity and specificity, then the Rogan-Gladen
//!    estimator corrects the judge's raw positive rate for that measured
//!    error. Motivated by the bias-plus-uncertainty (arXiv:2605.06939) and
//!    noisy-but-valid (arXiv:2601.20913) lines of work.
//! 3. **Routing** ([`route_low_confidence`]): escalate a [`ConfidenceBand`]
//!    of `Low`, accept otherwise. Reuses the  band rather than
//!    inventing a second threshold scheme.
//!
//! Everything here is pure statistics: no LLM calls, no I/O. The escalation
//! *target* (a larger juror, or a human review queue) is the caller's
//! concern; the human queue itself is enterprise. This module
//! only produces the escalate-or-accept decision.

use super::agreement::ConfidenceBand;

/// One judged verdict paired with the ground truth, the unit of a
/// calibration set.
///
/// `confidence` is the judge's stated confidence in `[0.0, 1.0]` (for a
/// jury, the consensus confidence). `correct` is whether that verdict
/// matched the trusted label. A calibration set is a slice of these.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct LabeledVerdict {
    /// Stated confidence in the verdict, clamped to `[0.0, 1.0]` before use.
    pub confidence: f64,
    /// Whether the verdict matched the trusted ground-truth label.
    pub correct: bool,
}

/// Number of equal-width bins used by [`calibration_error`].
///
/// Ten bins is the conventional choice for Expected Calibration Error and
/// keeps the metric stable on the small label sets a single developer can
/// realistically hand-label.
const ECE_BINS: usize = 10;

/// Expected Calibration Error (ECE) of a labeled verdict set.
///
/// Partitions stated confidence into `ECE_BINS` equal-width bins over
/// `[0, 1]`, and for each non-empty bin takes the absolute gap between the
/// mean confidence and the observed accuracy, weighted by the bin's share
/// of samples:
///
/// ```text
/// ECE = sum_b (n_b / N) * | mean_confidence_b - accuracy_b |
/// ```
///
/// Lower is better: 0.0 means stated confidence matches accuracy in every
/// bin (perfect calibration). Returns 0.0 for an empty set (no claim to
/// disagree with). See arXiv:2508.06225 for why judge confidence needs
/// this check rather than being trusted directly.
pub fn calibration_error(samples: &[LabeledVerdict]) -> f64 {
    ece_over(samples.iter().map(|s| (clamp01(s.confidence), s.correct)))
}

/// Expected Calibration Error over an iterator of `(predicted_prob, correct)`
/// pairs, the shared binning both [`calibration_error`] and
/// [`calibrated_calibration_error`] reduce to.
///
/// Factoring the binning here keeps the raw and calibrated ECE provably
/// identical except for the probability fed in, so the calibrated metric
/// cannot drift from the raw one through a copied-then-edited loop. The
/// predicted probabilities are expected to already be clamped to `[0, 1]`.
/// Returns 0.0 for an empty iterator (no claim to disagree with).
fn ece_over(samples: impl Iterator<Item = (f64, bool)>) -> f64 {
    let predictions: Vec<(f64, bool)> = samples.collect();
    if predictions.is_empty() {
        return 0.0;
    }
    let total = predictions.len() as f64;
    let mut ece = 0.0;
    for bin in 0..ECE_BINS {
        let lo = bin as f64 / ECE_BINS as f64;
        let hi = (bin + 1) as f64 / ECE_BINS as f64;
        let mut count = 0.0;
        let mut conf_sum = 0.0;
        let mut correct_sum = 0.0;
        for (prob, correct) in predictions.iter().filter(|(p, _)| in_bin(*p, lo, hi, bin)) {
            count += 1.0;
            conf_sum += clamp01(*prob);
            if *correct {
                correct_sum += 1.0;
            }
        }
        if count == 0.0 {
            continue;
        }
        let mean_conf = conf_sum / count;
        let accuracy = correct_sum / count;
        ece += (count / total) * (mean_conf - accuracy).abs();
    }
    ece
}

/// Brier score of a labeled verdict set.
///
/// Mean squared error between stated confidence and the 0/1 outcome:
///
/// ```text
/// Brier = (1/N) * sum ( confidence - outcome )^2
/// ```
///
/// Lower is better, ranging `[0.0, 1.0]`. It rewards both calibration and
/// sharpness, so it complements [`calibration_error`] rather than
/// duplicating it. Returns 0.0 for an empty set.
pub fn brier_score(samples: &[LabeledVerdict]) -> f64 {
    if samples.is_empty() {
        return 0.0;
    }
    let n = samples.len() as f64;
    let sum: f64 = samples
        .iter()
        .map(|s| {
            let outcome = if s.correct { 1.0 } else { 0.0 };
            (clamp01(s.confidence) - outcome).powi(2)
        })
        .sum();
    sum / n
}

/// A fitted 1-D logistic calibrator (the "linear probe"): sigmoid(w*conf + b).
///
/// Maps a judge's raw stated confidence onto a calibrated probability so an
/// over- or under-confident local judge can still gate CI honestly. Fit on a
/// labeled verdict set; see arXiv:2512.22245.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct LinearProbe {
    /// Slope on the raw confidence. Above 1.0 sharpens an under-confident
    /// judge; below 1.0 softens an over-confident one.
    pub weight: f64,
    /// Intercept that shifts the whole curve, correcting a judge that is
    /// systematically high or low regardless of the input.
    pub bias: f64,
}

/// Number of gradient-descent steps used to fit a [`LinearProbe`].
///
/// A fixed budget keeps the fit deterministic (the same labeled set always
/// yields the same probe) so the calibrated metric is reproducible in CI; a
/// few hundred steps is enough to converge on the 1-D log-loss surface.
const PROBE_FIT_ITERS: usize = 500;
/// Learning rate for the [`LinearProbe`] gradient descent.
const PROBE_FIT_LR: f64 = 0.1;

impl LinearProbe {
    /// Map a raw stated confidence through the fitted curve to a calibrated
    /// probability in `[0.0, 1.0]`.
    ///
    /// Applies `sigmoid(weight * raw + bias)` and clamps the result so a
    /// caller can feed it straight into a gate without re-checking the range.
    pub fn calibrate(&self, raw_confidence: f64) -> f64 {
        clamp01(sigmoid(self.weight * raw_confidence + self.bias))
    }
}

/// Fit a 1-D logistic [`LinearProbe`] by gradient descent on log-loss.
///
/// Treats each sample's `correct` flag as the 0/1 target and its `confidence`
/// as the single feature, then minimizes binary cross-entropy over a fixed
/// step budget so the result is deterministic. On an empty, single-sample, or
/// single-label set (where the logistic fit is undefined or unconstrained) it
/// returns the near-identity probe `{weight: 1.0, bias: 0.0}` so callers never
/// divide by zero or propagate a NaN. Any NaN that appears mid-descent also
/// falls back to the identity probe. See arXiv:2512.22245.
pub fn fit_linear_probe(samples: &[LabeledVerdict]) -> LinearProbe {
    let identity = LinearProbe {
        weight: 1.0,
        bias: 0.0,
    };
    if samples.len() < 2 {
        return identity;
    }
    let first = samples[0].correct;
    if samples.iter().all(|s| s.correct == first) {
        // A single-label set leaves the slope unconstrained; an honest fit
        // would run the bias to +/- infinity, so decline and stay identity.
        return identity;
    }
    let n = samples.len() as f64;
    let mut weight = 1.0;
    let mut bias = 0.0;
    for _ in 0..PROBE_FIT_ITERS {
        let mut grad_w = 0.0;
        let mut grad_b = 0.0;
        for s in samples {
            let x = clamp01(s.confidence);
            let target = if s.correct { 1.0 } else { 0.0 };
            let predicted = sigmoid(weight * x + bias);
            let error = predicted - target;
            grad_w += error * x;
            grad_b += error;
        }
        weight -= PROBE_FIT_LR * grad_w / n;
        bias -= PROBE_FIT_LR * grad_b / n;
        if !weight.is_finite() || !bias.is_finite() {
            return identity;
        }
    }
    LinearProbe { weight, bias }
}

/// Calibrated Expected Calibration Error: the same binning as
/// [`calibration_error`] but each confidence is first mapped through `probe`.
///
/// Reports how well a judge tracks accuracy *after* its raw confidence is run
/// through the fitted calibrator, so a CI gate can compare the post-calibration
/// ECE against the raw one and confirm the probe actually helped. Shares the
/// private `ece_over` binning with [`calibration_error`] so the two metrics
/// stay defined identically. See arXiv:2512.22245.
pub fn calibrated_calibration_error(samples: &[LabeledVerdict], probe: &LinearProbe) -> f64 {
    ece_over(
        samples
            .iter()
            .map(|s| (probe.calibrate(s.confidence), s.correct)),
    )
}

/// Logistic sigmoid mapping any real number to `(0.0, 1.0)`, the link function
/// behind the [`LinearProbe`] calibrator.
fn sigmoid(z: f64) -> f64 {
    1.0 / (1.0 + (-z).exp())
}

/// A judge's measured error rates against a trusted label set.
///
/// `sensitivity` is the true-positive rate (the judge passes things that
/// should pass); `specificity` is the true-negative rate (the judge fails
/// things that should fail). Both live in `[0.0, 1.0]`. A perfect judge is
/// `{1.0, 1.0}`; an uninformative coin-flip judge is `{0.5, 0.5}`.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct JudgeReliability {
    /// True-positive rate `tp / (tp + fn)`.
    pub sensitivity: f64,
    /// True-negative rate `tn / (tn + fp)`.
    pub specificity: f64,
}

/// Estimate judge reliability from a confusion matrix on the trusted set.
///
/// `tp`/`fn_` are counts where the truth is positive; `tn`/`fp` where the
/// truth is negative. Sensitivity and specificity each guard their
/// divide-by-zero: an empty positive class yields sensitivity 0.0 and an
/// empty negative class yields specificity 0.0, so the function never
/// panics on a lopsided trusted set.
pub fn estimate_reliability(tp: u32, fn_: u32, tn: u32, fp: u32) -> JudgeReliability {
    let positives = (tp + fn_) as f64;
    let negatives = (tn + fp) as f64;
    let sensitivity = if positives > 0.0 {
        tp as f64 / positives
    } else {
        0.0
    };
    let specificity = if negatives > 0.0 {
        tn as f64 / negatives
    } else {
        0.0
    };
    JudgeReliability {
        sensitivity,
        specificity,
    }
}

/// Rogan-Gladen bias correction of an observed positive rate.
///
/// An imperfect judge reports `observed_positive_rate` (the fraction of
/// items it passed). Given its measured `reliability`, the true positive
/// rate is recovered with:
///
/// ```text
/// true = (observed + specificity - 1) / (sensitivity + specificity - 1)
/// ```
///
/// The result is clamped to `[0, 1]`. The denominator is Youden's J; when
/// it is non-positive the judge carries no information (e.g. a coin flip
/// `{0.5, 0.5}` gives J = 0), so we return `observed_positive_rate`
/// unchanged rather than divide by zero or amplify noise. See
/// arXiv:2605.06939.
pub fn rogan_gladen(observed_positive_rate: f64, reliability: &JudgeReliability) -> f64 {
    let denom = reliability.sensitivity + reliability.specificity - 1.0;
    if denom <= 0.0 {
        return clamp01(observed_positive_rate);
    }
    let corrected = (observed_positive_rate + reliability.specificity - 1.0) / denom;
    clamp01(corrected)
}

/// Wald 95% confidence interval around the Rogan-Gladen corrected rate.
///
/// Treats the observed positive rate as a binomial proportion over
/// `trusted_set_size` items, builds the Wald interval on the observed
/// rate, then maps both endpoints through [`rogan_gladen`] so the band
/// reflects the same bias correction as the point estimate. Returns
/// `(low, high)` clamped to `[0, 1]`.
///
/// Returns `(corrected, corrected)` (a zero-width band) when
/// `trusted_set_size` is zero, since a binomial standard error is
/// undefined without samples. The Wald form is the simple choice the
/// ticket calls for; it is known to be loose near 0 and 1, which is
/// acceptable for a small-trusted-set sanity band. See arXiv:2601.20913.
pub fn corrected_rate_ci(
    observed_positive_rate: f64,
    reliability: &JudgeReliability,
    trusted_set_size: u32,
) -> (f64, f64) {
    let point = rogan_gladen(observed_positive_rate, reliability);
    if trusted_set_size == 0 {
        return (point, point);
    }
    let p = clamp01(observed_positive_rate);
    let n = trusted_set_size as f64;
    let se = (p * (1.0 - p) / n).sqrt();
    let margin = 1.96 * se;
    let obs_low = clamp01(p - margin);
    let obs_high = clamp01(p + margin);
    let a = rogan_gladen(obs_low, reliability);
    let b = rogan_gladen(obs_high, reliability);
    // Rogan-Gladen rises with the observed rate for an informative judge
    // but inverts for one with J < 0; order the endpoints defensively.
    (a.min(b), a.max(b))
}

/// Where a verdict goes after the confidence band is known.
///
/// The variants are deliberately coarse: this module decides only whether
/// a verdict is trustworthy enough to accept as-is or needs a second look.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Routing {
    /// The verdict stands; report it as a settled pass or fail.
    Accept,
    /// The verdict is low-confidence; hand it to the escalation target.
    Escalate,
}

/// Route a verdict by its  confidence band.
///
/// Escalates on [`ConfidenceBand::Low`] and accepts `High` or `Medium`.
/// This reuses the existing band (itself derived from the low-agreement
/// threshold) rather than inventing a new cutoff, so the routing decision
/// can never disagree with the band the reporter already shows. The
/// escalation *target*, a larger juror or a human review queue, is the
/// caller's or enterprise's concern.
pub fn route_low_confidence(band: ConfidenceBand) -> Routing {
    match band {
        ConfidenceBand::Low => Routing::Escalate,
        ConfidenceBand::High | ConfidenceBand::Medium => Routing::Accept,
    }
}

/// Clamp a value to the `[0.0, 1.0]` interval.
fn clamp01(v: f64) -> f64 {
    v.clamp(0.0, 1.0)
}

/// True when `conf` falls in bin `[lo, hi)`, with the top bin closed at
/// `hi` so a confidence of exactly 1.0 lands in the last bin rather than
/// falling off the end.
fn in_bin(conf: f64, lo: f64, hi: f64, bin: usize) -> bool {
    let c = clamp01(conf);
    if bin == ECE_BINS - 1 {
        c >= lo && c <= hi
    } else {
        c >= lo && c < hi
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn approx(a: f64, b: f64, tol: f64) -> bool {
        (a - b).abs() < tol
    }

    fn v(confidence: f64, correct: bool) -> LabeledVerdict {
        LabeledVerdict {
            confidence,
            correct,
        }
    }

    #[test]
    fn calibration_error_is_zero_for_empty_set() {
        assert!(approx(calibration_error(&[]), 0.0, 1e-12));
    }

    #[test]
    fn calibration_error_is_zero_for_perfectly_calibrated_set() {
        // A bin at confidence 1.0 that is always correct, and a bin at
        // confidence 0.0 that is always wrong: accuracy matches confidence
        // in every populated bin, so ECE is 0.
        let samples = vec![v(1.0, true), v(1.0, true), v(0.0, false), v(0.0, false)];
        assert!(approx(calibration_error(&samples), 0.0, 1e-9));
    }

    #[test]
    fn calibration_error_flags_overconfidence() {
        // Judge states 1.0 confidence but is right only half the time: the
        // single populated bin has confidence 1.0, accuracy 0.5, gap 0.5.
        let samples = vec![v(1.0, true), v(1.0, false)];
        assert!(approx(calibration_error(&samples), 0.5, 1e-9));
    }

    #[test]
    fn calibration_error_weights_bins_by_population() {
        // Bin near 0.95 (top bin): three samples, two correct -> conf 0.95,
        // acc 2/3, gap ~0.2833, weight 3/4. Bin near 0.05: one sample,
        // wrong -> conf 0.05, acc 0, gap 0.05, weight 1/4.
        let samples = vec![v(0.95, true), v(0.95, true), v(0.95, false), v(0.05, false)];
        let expected = 0.75 * (0.95_f64 - 2.0 / 3.0).abs() + 0.25 * 0.05;
        assert!(approx(calibration_error(&samples), expected, 1e-9));
    }

    #[test]
    fn calibration_error_clamps_out_of_range_confidence() {
        // Confidence above 1.0 is clamped into the top bin and treated as
        // 1.0; an always-correct sample then shows zero gap.
        let samples = vec![v(1.5, true), v(1.5, true)];
        assert!(approx(calibration_error(&samples), 0.0, 1e-9));
    }

    #[test]
    fn brier_score_is_zero_for_empty_set() {
        assert!(approx(brier_score(&[]), 0.0, 1e-12));
    }

    #[test]
    fn brier_score_is_zero_for_perfect_confident_judge() {
        let samples = vec![v(1.0, true), v(0.0, false)];
        assert!(approx(brier_score(&samples), 0.0, 1e-12));
    }

    #[test]
    fn brier_score_is_one_for_perfectly_wrong_confident_judge() {
        let samples = vec![v(1.0, false), v(0.0, true)];
        assert!(approx(brier_score(&samples), 1.0, 1e-12));
    }

    #[test]
    fn brier_score_is_quarter_for_uniform_half_confidence() {
        // Every verdict at 0.5 confidence: (0.5)^2 = 0.25 regardless of
        // outcome.
        let samples = vec![v(0.5, true), v(0.5, false), v(0.5, true)];
        assert!(approx(brier_score(&samples), 0.25, 1e-12));
    }

    #[test]
    fn linear_probe_lowers_ece_of_an_overconfident_judge() {
        // 100 verdicts that all state 0.95 confidence but are correct only
        // ~60% of the time: a single bin with a 0.35 gap. The probe should
        // pull the calibrated confidence down toward the observed 0.6 and
        // shrink the gap.
        let mut samples = Vec::with_capacity(100);
        for i in 0..100 {
            samples.push(v(0.95, i < 60));
        }
        let probe = fit_linear_probe(&samples);
        let calibrated = calibrated_calibration_error(&samples, &probe);
        let raw = calibration_error(&samples);
        assert!(
            calibrated < raw,
            "calibrated ECE {calibrated} should beat raw ECE {raw}"
        );
    }

    #[test]
    fn linear_probe_calibrate_is_monotone_in_confidence() {
        // A separating set (low confidence wrong, high confidence right) fits
        // a positive slope, so a higher raw confidence maps to a calibrated
        // probability that is at least as high.
        let samples = vec![
            v(0.1, false),
            v(0.2, false),
            v(0.3, false),
            v(0.7, true),
            v(0.8, true),
            v(0.9, true),
        ];
        let probe = fit_linear_probe(&samples);
        assert!(probe.calibrate(0.9) >= probe.calibrate(0.1));
    }

    #[test]
    fn fit_linear_probe_is_safe_on_empty_set() {
        // No samples means the identity probe; calibrating still yields a
        // value in range rather than a NaN or a panic.
        let calibrated = fit_linear_probe(&[]).calibrate(0.5);
        assert!((0.0..=1.0).contains(&calibrated), "got {calibrated}");
    }

    #[test]
    fn estimate_reliability_computes_rates_from_confusion_matrix() {
        // tp=9 fn=1 -> sensitivity 0.9; tn=8 fp=2 -> specificity 0.8.
        let r = estimate_reliability(9, 1, 8, 2);
        assert!(approx(r.sensitivity, 0.9, 1e-12));
        assert!(approx(r.specificity, 0.8, 1e-12));
    }

    #[test]
    fn estimate_reliability_guards_empty_classes() {
        // No positive-truth items -> sensitivity 0.0, not a panic.
        let r = estimate_reliability(0, 0, 5, 5);
        assert!(approx(r.sensitivity, 0.0, 1e-12));
        assert!(approx(r.specificity, 0.5, 1e-12));
        // No negative-truth items -> specificity 0.0.
        let r2 = estimate_reliability(5, 5, 0, 0);
        assert!(approx(r2.specificity, 0.0, 1e-12));
    }

    #[test]
    fn rogan_gladen_corrects_observed_rate_with_known_values() {
        // (0.5 + 0.8 - 1) / (0.9 + 0.8 - 1) = 0.3 / 0.7 ~ 0.428571.
        let r = JudgeReliability {
            sensitivity: 0.9,
            specificity: 0.8,
        };
        assert!(approx(rogan_gladen(0.5, &r), 0.3 / 0.7, 1e-9));
    }

    #[test]
    fn rogan_gladen_is_identity_for_perfect_judge() {
        let r = JudgeReliability {
            sensitivity: 1.0,
            specificity: 1.0,
        };
        // denom = 1, corrected = observed.
        assert!(approx(rogan_gladen(0.42, &r), 0.42, 1e-12));
    }

    #[test]
    fn rogan_gladen_returns_observed_for_uninformative_judge() {
        // Coin-flip judge: J = 0.5 + 0.5 - 1 = 0, non-positive denominator,
        // so the observed rate passes through unchanged.
        let r = JudgeReliability {
            sensitivity: 0.5,
            specificity: 0.5,
        };
        assert!(approx(rogan_gladen(0.37, &r), 0.37, 1e-12));
    }

    #[test]
    fn rogan_gladen_returns_observed_for_anti_correlated_judge() {
        // J < 0: still non-positive, so we decline to "correct" and return
        // the observed rate.
        let r = JudgeReliability {
            sensitivity: 0.2,
            specificity: 0.3,
        };
        assert!(approx(rogan_gladen(0.6, &r), 0.6, 1e-12));
    }

    #[test]
    fn rogan_gladen_clamps_into_unit_interval() {
        // Low observed rate with high specificity can push the raw estimate
        // below 0; it must clamp to 0.
        let r = JudgeReliability {
            sensitivity: 0.9,
            specificity: 0.9,
        };
        let corrected = rogan_gladen(0.05, &r);
        assert!((0.0..=1.0).contains(&corrected));
        assert!(approx(corrected, 0.0, 1e-9));
    }

    #[test]
    fn corrected_rate_ci_is_zero_width_for_empty_trusted_set() {
        let r = JudgeReliability {
            sensitivity: 0.9,
            specificity: 0.8,
        };
        let (lo, hi) = corrected_rate_ci(0.5, &r, 0);
        let point = rogan_gladen(0.5, &r);
        assert!(approx(lo, point, 1e-12));
        assert!(approx(hi, point, 1e-12));
    }

    #[test]
    fn corrected_rate_ci_brackets_the_point_estimate() {
        let r = JudgeReliability {
            sensitivity: 0.9,
            specificity: 0.8,
        };
        let point = rogan_gladen(0.5, &r);
        let (lo, hi) = corrected_rate_ci(0.5, &r, 100);
        assert!(lo <= point && point <= hi);
        assert!(lo < hi);
        assert!((0.0..=1.0).contains(&lo));
        assert!((0.0..=1.0).contains(&hi));
    }

    #[test]
    fn corrected_rate_ci_narrows_with_more_samples() {
        let r = JudgeReliability {
            sensitivity: 0.9,
            specificity: 0.8,
        };
        let (lo_small, hi_small) = corrected_rate_ci(0.5, &r, 10);
        let (lo_large, hi_large) = corrected_rate_ci(0.5, &r, 1000);
        assert!((hi_large - lo_large) < (hi_small - lo_small));
    }

    #[test]
    fn route_low_confidence_escalates_only_low_band() {
        assert_eq!(route_low_confidence(ConfidenceBand::Low), Routing::Escalate);
        assert_eq!(
            route_low_confidence(ConfidenceBand::Medium),
            Routing::Accept
        );
        assert_eq!(route_low_confidence(ConfidenceBand::High), Routing::Accept);
    }
}
