helena 0.1.0

use std::num::NonZeroUsize;

use crate::{Error, Result};

/// A bootstrap confidence interval for a scalar metric.
///
/// The point estimate (the statistic on the observed data) plus lower and upper
/// percentile bounds at a stated confidence level. Produced by
/// [`ClusterBootstrap::measure`] and read wherever a metric would otherwise be
/// reported as a bare number, so a difference between two systems can be judged
/// against its own sampling noise rather than taken at face value. This is the read
/// the small-sample evaluation brief calls for: an objective bakeoff that "picks a
/// winner" between two checkpoints whose curves differ by less than the interval is
/// reading noise, not signal.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct ConfidenceInterval {
    point: f64,
    lo: f64,
    hi: f64,
    confidence: f64,
}

impl ConfidenceInterval {
    /// The statistic on the observed data: the mean of per-cluster means.
    pub fn point(&self) -> f64 {
        self.point
    }

    /// The lower percentile bound.
    pub fn lo(&self) -> f64 {
        self.lo
    }

    /// The upper percentile bound.
    pub fn hi(&self) -> f64 {
        self.hi
    }

    /// The confidence level the bounds were computed at, in `(0, 1)`.
    pub fn confidence(&self) -> f64 {
        self.confidence
    }

    /// The interval width, `hi - lo`.
    pub fn width(&self) -> f64 {
        self.hi - self.lo
    }
}

/// A clustered (nested) percentile bootstrap over per-conditioning-item observations.
///
/// The uncertainty layer for the latent and audio diagnostics: wrap any per-item
/// scalar metric (per-item dispersion, spectral distance, preservation correlation,
/// a per-SNR bin's losses) in a confidence interval so a metric *delta* is read
/// against its own sampling noise. The small-sample evaluation regime makes this
/// non-optional — a point estimate over a few hundred clustered observations can
/// reorder two systems on resampling alone.
///
/// **Clustering and the statistic.** Each `cluster` holds the observations for one
/// conditioning item; both estimators resample *items* with replacement, then
/// *observations within each drawn item* with replacement — the brief's "resample
/// items first, then seeds within items," the coherent unit because outputs from one
/// condition are not independent. They differ only in how a resample is summarized:
/// [`measure`](Self::measure) takes the mean of per-cluster means, so an item with more
/// observations does not outweigh one with fewer (the estimand when the conditioning
/// item is itself the unit of analysis, e.g. per-item dispersion or spectral distance);
/// [`measure_pooled`](Self::measure_pooled) takes the observation-weighted pooled mean
/// (the estimand when a per-observation headline is being bracketed, e.g. a per-SNR
/// bin's pooled loss). The two coincide when clusters are equal-sized. A singleton
/// cluster makes the inner resample a no-op, so the flat per-item case degrades to a
/// plain item-level bootstrap.
///
/// **Bounds.** Percentile method: the `[alpha/2, 1 - alpha/2]` quantiles of the
/// resampled statistic (`alpha = 1 - confidence`), by linear interpolation between
/// order statistics. Both bounds lie within the observed range, so the interval is
/// scale-faithful to the data.
///
/// **Determinism and dependencies.** Resampling is driven by an inlined
/// `SplitMix64` seeded from `seed`, so a `(data, resamples, confidence, seed)`
/// tuple maps to exactly one interval and pulls no `rand` dependency into the core —
/// the same boundary the rest of `helena::eval` holds. Total, host-side, no Burn.
///
/// Non-mean statistics use [`measure_cluster_statistic`](Self::measure_cluster_statistic):
/// resample conditioning-item clusters, then let the caller recompute the whole
/// statistic over the resampled cluster set.
#[derive(Clone, Copy, Debug)]
pub struct ClusterBootstrap {
    resamples: NonZeroUsize,
    confidence: f64,
    seed: u64,
}

impl ClusterBootstrap {
    /// Configure a bootstrap of `resamples` replicates at a `confidence` level,
    /// seeded by `seed`.
    ///
    /// Returns [`Error::Validation`] unless `confidence` is finite and strictly
    /// inside `(0, 1)`.
    pub fn new(resamples: NonZeroUsize, confidence: f64, seed: u64) -> Result<Self> {
        if !(confidence.is_finite() && confidence > 0.0 && confidence < 1.0) {
            return Err(Error::validation(format!(
                "bootstrap confidence must be in (0, 1); got {confidence}"
            )));
        }
        Ok(Self {
            resamples,
            confidence,
            seed,
        })
    }

    /// Bootstrap a [`ConfidenceInterval`] for the *equal-weight* mean of per-cluster
    /// means of `clusters`.
    ///
    /// Each item counts once regardless of its observation count, so this is the
    /// estimand when the conditioning item is the unit of analysis (per-item dispersion,
    /// spectral distance). For an observation-weighted headline use
    /// [`measure_pooled`](Self::measure_pooled); the two coincide when clusters are
    /// equal-sized.
    ///
    /// `clusters` holds one inner vector of observations per conditioning item.
    /// Returns [`Error::Validation`] if there are no clusters, if any cluster is
    /// empty, or if any observation is non-finite (a diverged upstream metric must
    /// surface as an error, not a `NaN` interval).
    pub fn measure(&self, clusters: &[Vec<f64>]) -> Result<ConfidenceInterval> {
        validate_clusters(clusters)?;
        let point = mean_of_cluster_means(clusters);
        let k = clusters.len();
        let mut stats = Vec::with_capacity(self.resamples.get());
        let mut rng = SplitMix64::new(self.seed);
        for _ in 0..self.resamples.get() {
            let mut sum_means = 0.0f64;
            for _ in 0..k {
                let cluster = &clusters[rng.index(k)];
                let m = cluster.len();
                let mut sum = 0.0f64;
                for _ in 0..m {
                    sum += cluster[rng.index(m)];
                }
                sum_means += sum / m as f64;
            }
            stats.push(sum_means / k as f64);
        }
        Ok(self.percentile_interval(point, stats))
    }

    /// Bootstrap a [`ConfidenceInterval`] for the *pooled* (observation-weighted) mean
    /// of `clusters`.
    ///
    /// The point estimate is the flat mean of every observation, so a cluster with more
    /// observations contributes proportionally more — the estimand that brackets a
    /// per-observation headline such as a per-SNR bin's pooled loss. The resampling is
    /// identical to [`measure`](Self::measure) (items with replacement, then
    /// observations within each drawn item), so only the *estimand* differs, not the
    /// clustering: the interval still reads the headline against item-to-item, not
    /// observation-to-observation, sampling noise. Same validation and determinism as
    /// [`measure`](Self::measure).
    pub fn measure_pooled(&self, clusters: &[Vec<f64>]) -> Result<ConfidenceInterval> {
        validate_clusters(clusters)?;
        let point = pooled_mean(clusters);
        let k = clusters.len();
        let mut stats = Vec::with_capacity(self.resamples.get());
        let mut rng = SplitMix64::new(self.seed);
        for _ in 0..self.resamples.get() {
            let mut total = 0.0f64;
            let mut count = 0usize;
            for _ in 0..k {
                let cluster = &clusters[rng.index(k)];
                let m = cluster.len();
                for _ in 0..m {
                    total += cluster[rng.index(m)];
                    count += 1;
                }
            }
            // Each of the k drawn clusters contributes at least one observation.
            stats.push(total / count as f64);
        }
        Ok(self.percentile_interval(point, stats))
    }

    /// Bootstrap a [`ConfidenceInterval`] for a caller-defined statistic over
    /// conditioning-item clusters.
    ///
    /// Unlike [`measure`](Self::measure) and [`measure_pooled`](Self::measure_pooled),
    /// this does not average observations inside a cluster. Each resample draws the
    /// same number of clusters with replacement, then calls `statistic` on that cluster
    /// set. This is the primitive for correlations and system-level ranking accuracy:
    /// the statistic must see item grouping directly instead of receiving a collapsed
    /// mean.
    ///
    /// `statistic` is responsible for domain validation (for example, enough rank
    /// variability after resampling). If it returns an error for the observed data or
    /// any resample, the interval is rejected rather than silently conditioning on the
    /// successful resamples.
    pub fn measure_cluster_statistic<T, F>(
        &self,
        clusters: &[T],
        statistic: F,
    ) -> Result<ConfidenceInterval>
    where
        T: Clone,
        F: Fn(&[T]) -> Result<f64>,
    {
        if clusters.is_empty() {
            return Err(Error::validation("no clusters to bootstrap"));
        }
        let point = finite_statistic(statistic(clusters)?)?;
        let k = clusters.len();
        let mut stats = Vec::with_capacity(self.resamples.get());
        let mut rng = SplitMix64::new(self.seed);
        for _ in 0..self.resamples.get() {
            let mut sample = Vec::with_capacity(k);
            for _ in 0..k {
                sample.push(clusters[rng.index(k)].clone());
            }
            stats.push(finite_statistic(statistic(&sample)?)?);
        }
        Ok(self.percentile_interval(point, stats))
    }

    /// The `[alpha/2, 1 - alpha/2]` percentile interval of the resampled `stats` about
    /// `point`, by linear interpolation between order statistics. Sorts `stats` in
    /// place; shared by [`measure`](Self::measure) and
    /// [`measure_pooled`](Self::measure_pooled).
    fn percentile_interval(&self, point: f64, mut stats: Vec<f64>) -> ConfidenceInterval {
        stats.sort_by(f64::total_cmp);
        let alpha = 1.0 - self.confidence;
        ConfidenceInterval {
            point,
            lo: quantile_sorted(&stats, alpha / 2.0),
            hi: quantile_sorted(&stats, 1.0 - alpha / 2.0),
            confidence: self.confidence,
        }
    }
}

/// The shared precondition of [`ClusterBootstrap::measure`] and
/// [`ClusterBootstrap::measure_pooled`]: at least one cluster, every cluster non-empty,
/// every observation finite (a diverged upstream metric surfaces as an error, not a
/// `NaN` interval).
fn validate_clusters(clusters: &[Vec<f64>]) -> Result<()> {
    if clusters.is_empty() {
        return Err(Error::validation("no clusters to bootstrap"));
    }
    for cluster in clusters {
        if cluster.is_empty() {
            return Err(Error::validation(
                "every cluster must have at least one observation",
            ));
        }
        if let Some(bad) = cluster.iter().find(|v| !v.is_finite()) {
            return Err(Error::validation(format!(
                "bootstrap observations must be finite; found {bad}"
            )));
        }
    }
    Ok(())
}

fn finite_statistic(value: f64) -> Result<f64> {
    if value.is_finite() {
        Ok(value)
    } else {
        Err(Error::validation(format!(
            "bootstrap statistic must be finite; found {value}"
        )))
    }
}

/// The observation-weighted mean over every observation in `clusters` (the pooled
/// mean), as opposed to [`mean_of_cluster_means`]'s equal weight per cluster.
///
/// Summed over the flattened observations in one pass, so the result depends only on
/// their order, not on how they are grouped into clusters: a headline computed by
/// flattening the same observations (a per-SNR bin's pooled `mean_loss`, a step gap's
/// pooled cell mean) is reproduced bit-for-bit, so the interval is centered exactly.
fn pooled_mean(clusters: &[Vec<f64>]) -> f64 {
    let total: f64 = clusters.iter().flatten().copied().sum();
    let count: usize = clusters.iter().map(Vec::len).sum();
    total / count as f64
}

/// The mean of each cluster's mean (equal weight per cluster).
fn mean_of_cluster_means(clusters: &[Vec<f64>]) -> f64 {
    let k = clusters.len() as f64;
    clusters
        .iter()
        .map(|c| c.iter().sum::<f64>() / c.len() as f64)
        .sum::<f64>()
        / k
}

/// The `q`-quantile of an ascending-sorted, non-empty slice by linear interpolation
/// between order statistics (the type-7 / R default). `q` is clamped to `[0, 1]`.
fn quantile_sorted(xs: &[f64], q: f64) -> f64 {
    if xs.len() == 1 {
        return xs[0];
    }
    let h = q.clamp(0.0, 1.0) * (xs.len() as f64 - 1.0);
    let lo = h.floor() as usize;
    let hi = h.ceil() as usize;
    xs[lo] + (xs[hi] - xs[lo]) * (h - lo as f64)
}

/// SplitMix64 (Steele, Lea & Flood, 2014): a tiny, well-distributed generator that
/// is exactly reproducible from a `u64` seed. Inlined so the clustered bootstrap
/// stays deterministic without pulling `rand` into the core — the same "no new
/// dependency" boundary the rest of `helena::eval` holds. Not cryptographic; a
/// bootstrap only needs uniform index draws.
pub(crate) struct SplitMix64 {
    state: u64,
}

impl SplitMix64 {
    pub(crate) fn new(seed: u64) -> Self {
        Self { state: seed }
    }

    pub(crate) fn next_u64(&mut self) -> u64 {
        self.state = self.state.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = self.state;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    }

    /// A nearly-unbiased index in `[0, n)` via Lemire's multiply-high reduction.
    /// `n` must be non-zero, which holds: clusters and their members are validated
    /// non-empty before any draw.
    pub(crate) fn index(&mut self, n: usize) -> usize {
        ((u128::from(self.next_u64()) * n as u128) >> 64) as usize
    }
}

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

    fn boot(resamples: usize, confidence: f64, seed: u64) -> ClusterBootstrap {
        ClusterBootstrap::new(NonZeroUsize::new(resamples).unwrap(), confidence, seed).unwrap()
    }

    #[test]
    fn point_is_mean_of_cluster_means() {
        // Cluster means are 2.0 and 5.0, so the equal-weight statistic is 3.5 —
        // not the pooled mean (3.0), which would over-weight the larger cluster.
        let clusters = vec![vec![1.0, 3.0], vec![5.0]];
        let ci = boot(500, 0.95, 1).measure(&clusters).unwrap();
        assert!((ci.point() - 3.5).abs() < 1e-12, "{}", ci.point());
    }

    #[test]
    fn pooled_point_weights_by_observation() {
        // The same clusters: the pooled estimand weights by observation count, so
        // (1 + 3 + 5) / 3 = 3.0, the very thing `measure` deliberately avoids (3.5).
        let clusters = vec![vec![1.0, 3.0], vec![5.0]];
        let ci = boot(500, 0.95, 1).measure_pooled(&clusters).unwrap();
        assert!((ci.point() - 3.0).abs() < 1e-12, "{}", ci.point());
    }

    #[test]
    fn pooled_coincides_with_mean_of_means_for_equal_clusters() {
        // Equal-sized clusters make the two estimands identical, and the resampling
        // draw order is the same, so the whole interval is byte-for-byte equal.
        let clusters = vec![vec![1.0, 4.0], vec![2.0, 5.0], vec![9.0, -3.0]];
        let equal = boot(500, 0.95, 1234).measure(&clusters).unwrap();
        let pooled = boot(500, 0.95, 1234).measure_pooled(&clusters).unwrap();
        assert_eq!(equal, pooled);
    }

    #[test]
    fn degenerate_data_has_zero_width() {
        let clusters = vec![vec![2.0], vec![2.0, 2.0], vec![2.0]];
        let ci = boot(256, 0.95, 7).measure(&clusters).unwrap();
        assert_eq!(ci.point(), 2.0);
        assert_eq!(ci.lo(), 2.0);
        assert_eq!(ci.hi(), 2.0);
        assert_eq!(ci.width(), 0.0);
    }

    #[test]
    fn single_singleton_cluster_collapses() {
        let ci = boot(64, 0.9, 3).measure(&[vec![4.0]]).unwrap();
        assert_eq!(ci.point(), 4.0);
        assert_eq!(ci.lo(), 4.0);
        assert_eq!(ci.hi(), 4.0);
    }

    #[test]
    fn bounds_lie_within_observation_range() {
        let clusters = vec![vec![1.0, 2.0], vec![3.0], vec![4.0, 5.0, 6.0]];
        let ci = boot(1000, 0.95, 42).measure(&clusters).unwrap();
        assert!(ci.lo() >= 1.0, "lo {} below min", ci.lo());
        assert!(ci.hi() <= 6.0, "hi {} above max", ci.hi());
        assert!(ci.lo() <= ci.point() + 1e-9 && ci.point() <= ci.hi() + 1e-9);
        assert!(ci.width() > 0.0);
    }

    #[test]
    fn higher_confidence_is_wider() {
        let clusters = vec![vec![0.0], vec![1.0], vec![2.0], vec![3.0], vec![10.0]];
        let narrow = boot(1000, 0.80, 99).measure(&clusters).unwrap();
        let wide = boot(1000, 0.99, 99).measure(&clusters).unwrap();
        assert!(
            wide.width() >= narrow.width() - 1e-12,
            "0.99 width {} < 0.80 width {}",
            wide.width(),
            narrow.width()
        );
    }

    #[test]
    fn is_deterministic() {
        let clusters = vec![vec![1.0, 4.0], vec![2.0], vec![9.0, -3.0]];
        let a = boot(500, 0.95, 1234).measure(&clusters).unwrap();
        let b = boot(500, 0.95, 1234).measure(&clusters).unwrap();
        assert_eq!(a, b);
    }

    #[test]
    fn rejects_invalid_inputs() {
        let b = boot(100, 0.95, 0);
        assert!(b.measure(&[]).is_err());
        assert!(b.measure(&[vec![1.0], vec![]]).is_err());
        assert!(b.measure(&[vec![1.0, f64::NAN]]).is_err());
        assert!(b.measure(&[vec![f64::INFINITY]]).is_err());
        // The shared validation rejects the same inputs through the pooled estimand.
        assert!(b.measure_pooled(&[]).is_err());
        assert!(b.measure_pooled(&[vec![1.0], vec![]]).is_err());
        assert!(b.measure_pooled(&[vec![1.0, f64::NAN]]).is_err());

        let one = NonZeroUsize::new(100).unwrap();
        for bad in [0.0, 1.0, -0.1, 1.5, f64::NAN, f64::INFINITY] {
            assert!(
                ClusterBootstrap::new(one, bad, 0).is_err(),
                "confidence {bad} should be rejected"
            );
        }
    }

    #[test]
    fn cluster_statistic_bootstrap_resamples_whole_clusters() {
        let clusters = vec![vec![1.0, 2.0], vec![10.0], vec![100.0]];
        let ci = boot(128, 0.9, 4)
            .measure_cluster_statistic(&clusters, |sample| {
                Ok(sample.iter().map(|cluster| cluster[0]).sum::<f64>())
            })
            .unwrap();
        assert_eq!(ci.point(), 111.0);
        assert!(ci.width() > 0.0);
    }

    #[test]
    fn cluster_statistic_rejects_empty_and_nonfinite() {
        let b = boot(16, 0.95, 0);
        let empty: Vec<Vec<f64>> = Vec::new();
        assert!(
            b.measure_cluster_statistic(&empty, |_| unreachable!())
                .is_err()
        );
        assert!(
            b.measure_cluster_statistic(&[vec![1.0]], |_| Ok(f64::NAN))
                .is_err()
        );
    }

    use proptest::prelude::*;

    fn arb_clusters() -> impl Strategy<Value = Vec<Vec<f64>>> {
        proptest::collection::vec(proptest::collection::vec(-50.0f64..50.0, 1..6), 1..10)
    }

    proptest! {
        #[test]
        fn bounds_are_ordered_and_in_range(
            clusters in arb_clusters(),
            confidence in 0.50f64..0.999,
            seed in any::<u64>(),
        ) {
            let ci = boot(200, confidence, seed).measure(&clusters).unwrap();
            let flat: Vec<f64> = clusters.iter().flatten().copied().collect();
            let min = flat.iter().copied().fold(f64::INFINITY, f64::min);
            let max = flat.iter().copied().fold(f64::NEG_INFINITY, f64::max);
            prop_assert!(ci.lo() <= ci.hi());
            prop_assert!(ci.lo() >= min - 1e-9 && ci.hi() <= max + 1e-9);
            prop_assert!(ci.point() >= min - 1e-9 && ci.point() <= max + 1e-9);
        }

        #[test]
        fn point_equals_independent_mean_of_means(clusters in arb_clusters()) {
            let ci = boot(64, 0.95, 5).measure(&clusters).unwrap();
            let want = mean_of_cluster_means(&clusters);
            prop_assert!((ci.point() - want).abs() <= 1e-9 * want.abs().max(1.0));
        }

        #[test]
        fn pooled_point_equals_flat_mean(clusters in arb_clusters()) {
            let ci = boot(64, 0.95, 5).measure_pooled(&clusters).unwrap();
            let flat: Vec<f64> = clusters.iter().flatten().copied().collect();
            let want = flat.iter().sum::<f64>() / flat.len() as f64;
            prop_assert!((ci.point() - want).abs() <= 1e-9 * want.abs().max(1.0));
        }

        #[test]
        fn confidence_is_monotone_in_width(
            clusters in arb_clusters(),
            lower in 0.50f64..0.90,
            delta in 0.01f64..0.09,
            seed in any::<u64>(),
        ) {
            let wide = boot(200, lower + delta, seed).measure(&clusters).unwrap();
            let narrow = boot(200, lower, seed).measure(&clusters).unwrap();
            prop_assert!(wide.width() >= narrow.width() - 1e-9);
        }

        #[test]
        fn is_reproducible(
            clusters in arb_clusters(),
            confidence in 0.50f64..0.999,
            seed in any::<u64>(),
        ) {
            let a = boot(128, confidence, seed).measure(&clusters).unwrap();
            let b = boot(128, confidence, seed).measure(&clusters).unwrap();
            prop_assert_eq!(a, b);
        }
    }
}