//! Per-class threshold calibration: what operating point does a class support?
//!
//! The first run of `layer-eval` graded all seven classes against one global
//! threshold ([`DEFAULT_AUTO_THRESHOLD`], 0.85) and read the failure as the layer
//! being weak. Most of it was the threshold being wrong for six of the seven
//! classes.
//!
//! A score is the share of the k-nearest neighbourhood's kernel weight carrying a
//! tag, so 0.85 asks for roughly 13 of 15 neighbours to agree. Whether a class
//! can reach that is bounded by how many of its own members fall inside a fixed
//! `k`, which is a property of class size and local density rather than of how
//! distinguishable the sound is. Measured: hi-hat (109 files) and snare (188) sit
//! in the same top-1 band, 71.6% against 76.6%, and their recall at 0.85 differs
//! by a factor of nine. The score is well behaved inside a class and not
//! comparable across classes.
//!
//! So the question a ship gate should ask is not "what does this class do at
//! 0.85". It is "what is the most permissive threshold at which this class still
//! meets the precision we require, and what recall does it buy there". That
//! threshold is exactly a `tag_policy` row, which the layer format can already
//! carry and the export currently declines to ship.
//!
//! [`DEFAULT_AUTO_THRESHOLD`]: audiofiles_core::analysis::exemplar::DEFAULT_AUTO_THRESHOLD

/// Per-class counts at one score threshold.
#[derive(Default, Clone, Copy)]
pub(crate) struct Counts {
    pub(crate) tp: usize,
    pub(crate) fp: usize,
    pub(crate) fn_: usize,
}

impl Counts {
    /// Of what fired, how much was right. `None` when nothing fired, which is a
    /// different fact from firing and being wrong, and the two must not print
    /// the same.
    pub(crate) fn precision(self) -> Option<f64> {
        let predicted = self.tp + self.fp;
        (predicted > 0).then(|| self.tp as f64 / predicted as f64)
    }

    /// Of what should have fired, how much did.
    pub(crate) fn recall(self) -> Option<f64> {
        let actual = self.tp + self.fn_;
        (actual > 0).then(|| self.tp as f64 / actual as f64)
    }

    pub(crate) fn fired(self) -> usize {
        self.tp + self.fp
    }

    pub(crate) fn actual(self) -> usize {
        self.tp + self.fn_
    }

    pub(crate) fn add(&mut self, other: Self) {
        self.tp += other.tp;
        self.fp += other.fp;
        self.fn_ += other.fn_;
    }
}

/// One test sample's evidence for one class: what the layer scored it, whether it
/// really is that class, and which fold produced it.
///
/// The fold travels with the point because calibrating a threshold on the same
/// predictions it is then scored against is threshold-fitting on the test set.
/// See [`out_of_fold`].
#[derive(Clone, Copy)]
pub(crate) struct Point {
    pub(crate) score: f64,
    pub(crate) actual: bool,
    pub(crate) fold: usize,
}

/// A threshold and what shipping it would deliver.
#[derive(Clone, Copy)]
pub(crate) struct OperatingPoint {
    pub(crate) threshold: f64,
    pub(crate) counts: Counts,
}

/// z for a one-sided 95% lower bound.
const WILSON_Z: f64 = 1.645;

/// Lower bound of the Wilson score interval for `successes / trials`.
///
/// Why a bound rather than the ratio: the first calibrated run picked, for each
/// class, the threshold whose *observed* precision just cleared 95%, and held-out
/// precision then came in at 78.6% to 94.9%. Every class undershot. That is not
/// bad luck, it is what selecting the most permissive point that clears a bar
/// does: the point that clears it by the narrowest margin is the one most likely
/// to have cleared it by chance, and picking the extreme of a noisy set is
/// selection bias with a known direction.
///
/// The bound removes the ad-hoc part of the fix. Instead of "meet 95% and also
/// have at least N predictions", a class must be 95%-confident of being above the
/// bar, which asks for more evidence from a small sample and less from a large
/// one, on the same scale. 19 of 20 correct reads as 76% here; 190 of 200 reads
/// as 92%.
fn wilson_lower_bound(successes: usize, trials: usize) -> f64 {
    if trials == 0 {
        return 0.0;
    }
    let n = trials as f64;
    let p = successes as f64 / n;
    let z2 = WILSON_Z * WILSON_Z;
    let denom = 1.0 + z2 / n;
    let center = p + z2 / (2.0 * n);
    let margin = WILSON_Z * (p * (1.0 - p) / n + z2 / (4.0 * n * n)).sqrt();
    ((center - margin) / denom).max(0.0)
}

/// Counts at a fixed threshold. Mirrors the runtime rule: a tag applies when its
/// score is at or above the threshold.
pub(crate) fn counts_at(points: &[Point], threshold: f64) -> Counts {
    let mut c = Counts::default();
    for p in points {
        match (p.score >= threshold, p.actual) {
            (true, true) => c.tp += 1,
            (true, false) => c.fp += 1,
            (false, true) => c.fn_ += 1,
            (false, false) => {}
        }
    }
    c
}

/// The most permissive threshold at which this class is 95%-confident of holding
/// `target_precision`, with at least `min_support` predictions behind it.
///
/// Most permissive means lowest, because a lower threshold is more recall, and
/// recall is what we are buying once precision is fixed. Walks the score-sorted
/// points once, which visits every threshold that can produce a distinct split.
///
/// The bar is [`wilson_lower_bound`], not the observed ratio: see there for why
/// the observed ratio systematically overstates what a chosen threshold delivers
/// on data it did not see.
///
/// Two further guards:
///
/// - `min_support`. A floor under the bound, so a class cannot ship a policy off
///   a handful of predictions even when the arithmetic allows it.
/// - A threshold of zero is refused. It would fire on every sample including the
///   ones the index scored nothing for, which is not a classifier.
///
/// Precision is not monotone in the threshold, so a lower qualifying threshold
/// may sit below a stretch that does not qualify. That is fine and deliberate:
/// what ships is a single threshold, and the numbers reported are the ones
/// measured at it.
pub(crate) fn operating_point(
    points: &[Point],
    target_precision: f64,
    min_support: usize,
) -> Option<OperatingPoint> {
    let positives = points.iter().filter(|p| p.actual).count();
    if positives == 0 {
        return None;
    }

    let mut sorted: Vec<&Point> = points.iter().collect();
    sorted.sort_by(|a, b| b.score.total_cmp(&a.score));

    let mut tp = 0usize;
    let mut fp = 0usize;
    let mut best: Option<OperatingPoint> = None;

    let mut i = 0;
    while i < sorted.len() {
        // Every point sharing this score has to be taken with it: no threshold
        // can separate two samples that scored the same.
        let score = sorted[i].score;
        #[allow(
            clippy::float_cmp,
            reason = "exact equality is the point: a threshold cannot split two \
                      samples that scored bit-identically, so the group boundary \
                      has to be exact rather than within a tolerance"
        )]
        while i < sorted.len() && sorted[i].score == score {
            if sorted[i].actual {
                tp += 1;
            } else {
                fp += 1;
            }
            i += 1;
        }
        if score <= 0.0 {
            break;
        }
        let fired = tp + fp;
        if fired < min_support {
            continue;
        }
        if wilson_lower_bound(tp, fired) >= target_precision {
            // Descending walk, so each qualifying point is more permissive than
            // the last. Keep overwriting and the survivor is the lowest.
            best = Some(OperatingPoint {
                threshold: score,
                counts: Counts {
                    tp,
                    fp,
                    fn_: positives - tp,
                },
            });
        }
    }
    best
}

/// Calibrate and evaluate without letting a class pick its threshold from the
/// samples it is then graded on.
///
/// For each fold: choose the threshold from every *other* fold's points, then
/// count this fold's points at it. Summing across folds gives per-class numbers
/// that no threshold was fitted to. A fold whose calibration finds no qualifying
/// threshold contributes its positives as misses, because a policy that cannot be
/// derived is a policy that does not ship and a tag that never fires.
///
/// Returns the summed counts and the thresholds chosen, one per fold that found
/// one. The spread of those thresholds is worth reading: a class whose threshold
/// swings between folds is one whose operating point is not a stable property of
/// the class.
pub(crate) fn out_of_fold(
    points: &[Point],
    folds: usize,
    target_precision: f64,
    min_support: usize,
) -> (Counts, Vec<f64>) {
    let mut total = Counts::default();
    let mut thresholds = Vec::new();

    for fold in 0..folds {
        let calibration: Vec<Point> = points.iter().filter(|p| p.fold != fold).copied().collect();
        let held_out: Vec<Point> = points.iter().filter(|p| p.fold == fold).copied().collect();
        if held_out.is_empty() {
            continue;
        }

        // The calibration set is smaller than the whole, so scale the support
        // floor with it. Otherwise a class that clears `min_support` overall
        // fails to calibrate on 4/5 of the data for arithmetic reasons.
        let scaled_support = (min_support * (folds - 1)).div_ceil(folds).max(1);
        match operating_point(&calibration, target_precision, scaled_support) {
            Some(op) => {
                total.add(counts_at(&held_out, op.threshold));
                thresholds.push(op.threshold);
            }
            None => {
                // No threshold: nothing fires, so every positive here is a miss.
                total.fn_ += held_out.iter().filter(|p| p.actual).count();
            }
        }
    }
    (total, thresholds)
}

/// Mean, and the spread, of the per-fold thresholds.
pub(crate) fn threshold_spread(thresholds: &[f64]) -> Option<(f64, f64, f64)> {
    if thresholds.is_empty() {
        return None;
    }
    let mean = thresholds.iter().sum::<f64>() / thresholds.len() as f64;
    let min = thresholds.iter().copied().fold(f64::INFINITY, f64::min);
    let max = thresholds.iter().copied().fold(f64::NEG_INFINITY, f64::max);
    Some((mean, min, max))
}

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

    fn pts(spec: &[(f64, bool)]) -> Vec<Point> {
        spec.iter()
            .enumerate()
            .map(|(i, &(score, actual))| Point {
                score,
                actual,
                fold: i % 5,
            })
            .collect()
    }

    #[test]
    fn counts_at_matches_the_runtime_at_or_above_rule() {
        let p = pts(&[(0.9, true), (0.5, true), (0.5, false), (0.1, true)]);
        let c = counts_at(&p, 0.5);
        assert_eq!((c.tp, c.fp, c.fn_), (2, 1, 1), "0.5 fires at exactly 0.5");
    }

    #[test]
    fn operating_point_takes_the_most_permissive_qualifying_threshold() {
        // Perfectly ordered: every positive above every negative. Asserted as a
        // property rather than a magic threshold, because what "most permissive"
        // resolves to depends on the bound, the target and the sample size, and
        // a hardcoded number would be testing this fixture's arithmetic.
        let spec: Vec<(f64, bool)> = (0..100)
            .map(|i| (0.9 - f64::from(i) * 0.001, true))
            .chain((0..100).map(|i| (0.2 - f64::from(i) * 0.001, false)))
            .collect();
        let p = pts(&spec);
        let op = operating_point(&p, 0.95, 1).unwrap();

        // What it reports is what that threshold actually does.
        let at = counts_at(&p, op.threshold);
        assert_eq!((at.tp, at.fp), (op.counts.tp, op.counts.fp));
        assert!(wilson_lower_bound(at.tp, at.fired()) >= 0.95);

        // And nothing lower qualifies, which is what makes it the most permissive.
        for lower in p
            .iter()
            .map(|x| x.score)
            .filter(|s| *s < op.threshold && *s > 0.0)
        {
            let c = counts_at(&p, lower);
            assert!(
                wilson_lower_bound(c.tp, c.fired()) < 0.95,
                "threshold {lower} also qualifies, so {} was not the lowest",
                op.threshold
            );
        }
    }

    #[test]
    fn operating_point_refuses_a_threshold_backed_by_too_few_predictions() {
        // Two perfect predictions then a mess. The bound alone rejects this:
        // 2 of 2 is not evidence of 95%, and the first run read cymbal that way.
        let p = pts(&[
            (0.99, true),
            (0.98, true),
            (0.5, false),
            (0.5, false),
            (0.5, false),
            (0.4, true),
        ]);
        assert!(
            operating_point(&p, 0.95, 1).is_none(),
            "two perfect predictions are not 95% confidence of 95% precision"
        );
    }

    #[test]
    fn the_bound_demands_more_evidence_from_a_smaller_sample() {
        // The property the whole calibration rests on: the same observed ratio
        // qualifies at scale and does not qualify on a handful.
        assert!(wilson_lower_bound(19, 20) < 0.85);
        assert!(wilson_lower_bound(190, 200) > 0.9);
        // Monotone in sample size at a fixed ratio.
        assert!(wilson_lower_bound(9, 10) < wilson_lower_bound(90, 100));
        // Degenerate inputs stay in range rather than producing a NaN that
        // would silently compare false against every target.
        assert!(wilson_lower_bound(0, 0).abs() < f64::EPSILON);
        assert!(wilson_lower_bound(0, 5) >= 0.0);
        assert!(wilson_lower_bound(5, 5) <= 1.0);
    }

    #[test]
    fn operating_point_never_returns_a_zero_threshold() {
        // Every sample scores zero: the index had nothing to say. A zero
        // threshold would "apply" the tag to all of them.
        let spec: Vec<(f64, bool)> = (0..40).map(|i| (0.0, i % 2 == 0)).collect();
        assert!(operating_point(&pts(&spec), 0.4, 1).is_none());
    }

    #[test]
    fn operating_point_is_none_when_precision_is_unreachable() {
        let p = pts(&[(0.9, false), (0.8, false), (0.7, true)]);
        assert!(operating_point(&p, 0.95, 1).is_none());
    }

    #[test]
    fn tied_scores_are_taken_together() {
        // A threshold cannot split two equal scores, so the walk must not report
        // a precision that only holds if it does. A clean run of positives, then
        // a tied group of one positive and five negatives: taking the group
        // whole fails the bar, and splitting it would pass.
        let mut spec: Vec<(f64, bool)> = (0..100)
            .map(|i| (0.9 - f64::from(i) * 0.001, true))
            .collect();
        spec.push((0.5, true));
        spec.extend((0..5).map(|_| (0.5, false)));

        let op = operating_point(&pts(&spec), 0.95, 1).unwrap();
        assert!(
            op.threshold > 0.5,
            "the tie was split; got {}",
            op.threshold
        );
    }

    #[test]
    fn out_of_fold_counts_an_uncalibratable_fold_as_misses() {
        // Nothing separates these, so no fold finds a threshold and every
        // positive must be a miss rather than silently vanishing.
        let p = pts(&[
            (0.5, true),
            (0.5, false),
            (0.5, true),
            (0.5, false),
            (0.5, true),
            (0.5, false),
            (0.5, true),
            (0.5, false),
            (0.5, true),
            (0.5, false),
        ]);
        let (counts, thresholds) = out_of_fold(&p, 5, 0.95, 1);
        assert!(thresholds.is_empty());
        assert_eq!(counts.tp, 0);
        assert_eq!(counts.actual(), 5, "all five positives are accounted for");
    }

    #[test]
    fn out_of_fold_is_not_more_optimistic_than_the_data_supports() {
        // Separable data: out-of-fold should recover it, since the threshold
        // learned on four folds transfers to the fifth.
        let mut spec: Vec<(f64, bool)> = (0..100)
            .map(|i| (0.9 - f64::from(i) * 0.001, true))
            .collect();
        spec.extend((0..100).map(|i| (0.3 - f64::from(i) * 0.001, false)));
        let p = pts(&spec);
        let (counts, thresholds) = out_of_fold(&p, 5, 0.95, 4);
        assert_eq!(thresholds.len(), 5, "every fold calibrates");
        // Precision holds on folds the threshold never saw, which is the whole
        // contract. Not 100%: the target is 0.95, so the most permissive
        // qualifying threshold deliberately admits a few negatives, and demanding
        // perfection here would be asserting a stricter bar than was asked for.
        assert!(
            counts.precision().unwrap() >= 0.95,
            "{:?}",
            counts.precision()
        );
        assert!(counts.recall().unwrap() > 0.9, "{:?}", counts.recall());
    }

    #[test]
    fn threshold_spread_reports_min_and_max() {
        let (mean, min, max) = threshold_spread(&[0.4, 0.6, 0.5]).unwrap();
        assert!((mean - 0.5).abs() < 1e-9);
        assert!((min - 0.4).abs() < 1e-9);
        assert!((max - 0.6).abs() < 1e-9);
        assert!(threshold_spread(&[]).is_none());
    }
}