foxhunt/ml/src/benchmark/statistical_sampler.rs

//! Statistical Sampler for GPU Training Benchmarks
//!
//! Provides rigorous statistical analysis of training epoch times including:
//! - 95% confidence intervals using t-distribution
//! - P50 (median), P95, P99 percentiles
//! - Coefficient of variation for stability assessment
//! - Outlier detection and removal
//!
//! # Examples
//!
//! ```no_run
//! use ml::benchmark::statistical_sampler::StatisticalSampler;
//!
//! let mut sampler = StatisticalSampler::new(2); // 2 warmup epochs
//!
//! // Add epoch timing samples
//! for epoch_time in &[1.5, 1.6, 1.55, 1.58, 1.52] {
//!     sampler.add_sample(*epoch_time);
//! }
//!
//! // Compute statistics
//! let stats = sampler.compute_statistics().expect("Failed to compute statistics");
//! println!("Mean: {:.3}s, 95% CI: ({:.3}, {:.3})",
//!          stats.mean_seconds,
//!          stats.confidence_interval_95.0,
//!          stats.confidence_interval_95.1);
//! ```

use crate::MLError;
use statrs::distribution::{ContinuousCDF, StudentsT};

/// Complete benchmark statistics with confidence intervals and percentiles
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
pub struct BenchmarkStatistics {
    /// Mean epoch time in seconds
    pub mean_seconds: f64,
    /// Standard deviation of epoch times
    pub std_dev: f64,
    /// 95% confidence interval (lower bound, upper bound) in seconds
    pub confidence_interval_95: (f64, f64),
    /// P50 (median) epoch time in seconds
    pub p50_median: f64,
    /// P95 percentile epoch time in seconds
    pub p95: f64,
    /// P99 percentile epoch time in seconds
    pub p99: f64,
    /// Coefficient of variation (std_dev / mean)
    pub coefficient_of_variation: f64,
    /// Number of samples used for statistics (after outlier removal)
    pub num_samples: usize,
    /// Number of outliers removed
    pub outliers_removed: usize,
}

impl BenchmarkStatistics {
    /// Check if training is stable (CV < 0.15)
    pub fn is_stable(&self) -> bool {
        self.coefficient_of_variation < 0.15
    }

    /// Get percentage of samples that were outliers
    pub fn outlier_percentage(&self) -> f64 {
        if self.num_samples + self.outliers_removed == 0 {
            0.0
        } else {
            (self.outliers_removed as f64 / (self.num_samples + self.outliers_removed) as f64)
                * 100.0
        }
    }

    /// Get margin of error for 95% confidence interval
    pub fn margin_of_error_95(&self) -> f64 {
        self.confidence_interval_95.1 - self.mean_seconds
    }

    /// Check if sufficient samples exist for reliable statistics
    pub fn has_sufficient_samples(&self) -> bool {
        self.num_samples >= 10
    }
}

/// Statistical sampler for epoch timing measurements
#[derive(Debug, Clone)]
pub struct StatisticalSampler {
    /// Raw epoch time samples in seconds
    samples: Vec<f64>,
    /// Number of warmup epochs to skip
    warmup_epochs: usize,
}

impl StatisticalSampler {
    /// Create a new statistical sampler
    ///
    /// # Arguments
    ///
    /// * `warmup_epochs` - Number of initial epochs to skip (typically 2-3)
    ///
    /// # Examples
    ///
    /// ```
    /// use ml::benchmark::statistical_sampler::StatisticalSampler;
    ///
    /// let sampler = StatisticalSampler::new(2);
    /// ```
    pub fn new(warmup_epochs: usize) -> Self {
        Self {
            samples: Vec::new(),
            warmup_epochs,
        }
    }

    /// Add a new epoch time sample
    ///
    /// # Arguments
    ///
    /// * `epoch_time_seconds` - Duration of epoch in seconds
    ///
    /// # Examples
    ///
    /// ```
    /// use ml::benchmark::statistical_sampler::StatisticalSampler;
    ///
    /// let mut sampler = StatisticalSampler::new(2);
    /// sampler.add_sample(1.523);
    /// ```
    pub fn add_sample(&mut self, epoch_time_seconds: f64) {
        self.samples.push(epoch_time_seconds);
    }

    /// Compute comprehensive benchmark statistics
    ///
    /// # Returns
    ///
    /// Returns `Ok(BenchmarkStatistics)` on success, or an error if:
    /// - Insufficient samples after warmup and outlier removal
    /// - Standard deviation is 0 or NaN
    /// - Coefficient of variation exceeds 0.15 (unstable training)
    ///
    /// # Errors
    ///
    /// This function will return an error if:
    /// - Fewer than 3 samples remain after warmup
    /// - All samples are identical (std_dev = 0)
    /// - Coefficient of variation > 0.15 (warning)
    /// - Less than 10 samples after outlier removal (warning)
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use ml::benchmark::statistical_sampler::StatisticalSampler;
    ///
    /// let mut sampler = StatisticalSampler::new(2);
    /// for time in &[1.5, 1.6, 1.55, 1.58, 1.52] {
    ///     sampler.add_sample(*time);
    /// }
    ///
    /// let stats = sampler.compute_statistics()?;
    /// println!("Mean: {:.3}s ± {:.3}s", stats.mean_seconds, stats.std_dev);
    /// # Ok::<(), ml::MLError>(())
    /// ```
    pub fn compute_statistics(&self) -> Result<BenchmarkStatistics, MLError> {
        // Step 1: Remove warmup epochs
        if self.samples.len() <= self.warmup_epochs {
            return Err(MLError::InsufficientData(format!(
                "Only {} samples, but {} warmup epochs specified",
                self.samples.len(),
                self.warmup_epochs
            )));
        }

        let samples_after_warmup: Vec<f64> = self
            .samples
            .iter()
            .skip(self.warmup_epochs)
            .copied()
            .collect();

        if samples_after_warmup.len() < 3 {
            return Err(MLError::InsufficientData(format!(
                "Only {} samples after warmup (minimum 3 required)",
                samples_after_warmup.len()
            )));
        }

        // Step 2: Calculate initial statistics for outlier detection
        let initial_mean = Self::calculate_mean(&samples_after_warmup);
        let initial_std_dev = Self::calculate_std_dev(&samples_after_warmup, initial_mean);

        if initial_std_dev.is_nan() || initial_std_dev.is_infinite() {
            return Err(MLError::ValidationError {
                message: "Standard deviation is NaN or infinite".to_string(),
            });
        }

        // Step 3: Remove outliers (>3 standard deviations from mean)
        let (samples_clean, outliers_removed) = if initial_std_dev == 0.0 {
            // All samples identical - no outliers to remove
            (samples_after_warmup, 0)
        } else {
            Self::remove_outliers(&samples_after_warmup, initial_mean, initial_std_dev)
        };

        if samples_clean.is_empty() {
            return Err(MLError::InsufficientData(
                "All samples were identified as outliers".to_string(),
            ));
        }

        // Step 4: Calculate final statistics on clean data
        let mean = Self::calculate_mean(&samples_clean);
        let std_dev = Self::calculate_std_dev(&samples_clean, mean);

        if std_dev.is_nan() || std_dev.is_infinite() {
            return Err(MLError::ValidationError {
                message: "Final standard deviation is NaN or infinite".to_string(),
            });
        }

        if std_dev == 0.0 {
            return Err(MLError::ValidationError {
                message: "Standard deviation is zero (all samples identical)".to_string(),
            });
        }

        // Step 5: Calculate coefficient of variation
        let coefficient_of_variation = std_dev / mean;

        // Validation: Warn if CV > 0.15 (unstable training)
        if coefficient_of_variation > 0.15 {
            tracing::warn!(
                "High coefficient of variation ({:.3}): training may be unstable",
                coefficient_of_variation
            );
        }

        // Validation: Warn if < 10 samples
        if samples_clean.len() < 10 {
            tracing::warn!(
                "Low sample count ({}): statistics may not be reliable",
                samples_clean.len()
            );
        }

        // Step 6: Calculate 95% confidence interval using t-distribution
        let confidence_interval_95 =
            Self::calculate_confidence_interval(&samples_clean, mean, std_dev)?;

        // Step 7: Calculate percentiles
        let mut sorted_samples = samples_clean.clone();
        sorted_samples.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));

        let p50_median = Self::calculate_percentile(&sorted_samples, 50.0);
        let p95 = Self::calculate_percentile(&sorted_samples, 95.0);
        let p99 = Self::calculate_percentile(&sorted_samples, 99.0);

        Ok(BenchmarkStatistics {
            mean_seconds: mean,
            std_dev,
            confidence_interval_95,
            p50_median,
            p95,
            p99,
            coefficient_of_variation,
            num_samples: samples_clean.len(),
            outliers_removed,
        })
    }

    /// Calculate mean of samples
    fn calculate_mean(samples: &[f64]) -> f64 {
        let sum: f64 = samples.iter().sum();
        sum / samples.len() as f64
    }

    /// Calculate standard deviation of samples
    fn calculate_std_dev(samples: &[f64], mean: f64) -> f64 {
        if samples.len() <= 1 {
            return 0.0;
        }

        let variance: f64 =
            samples.iter().map(|x| (x - mean).powi(2)).sum::<f64>() / (samples.len() - 1) as f64; // Sample variance (n-1)

        variance.sqrt()
    }

    /// Remove outliers using iterative 3-sigma rule
    ///
    /// Uses an iterative approach to handle cases where outliers skew the initial
    /// mean and standard deviation. Continues removing outliers until no more are found.
    ///
    /// Returns (clean_samples, number_of_outliers_removed)
    fn remove_outliers(samples: &[f64], _mean: f64, _std_dev: f64) -> (Vec<f64>, usize) {
        let mut current_samples = samples.to_vec();
        let mut total_outliers_removed = 0;
        let max_iterations = 10; // Prevent infinite loops

        for iteration in 0..max_iterations {
            // Calculate current mean and std_dev
            let current_mean = Self::calculate_mean(&current_samples);
            let current_std_dev = Self::calculate_std_dev(&current_samples, current_mean);

            if current_std_dev == 0.0 {
                // All remaining samples are identical - no more outliers
                break;
            }

            let threshold = 3.0 * current_std_dev;
            let mut clean_samples = Vec::new();
            let mut outliers_in_iteration = 0;

            for &sample in &current_samples {
                if (sample - current_mean).abs() <= threshold {
                    clean_samples.push(sample);
                } else {
                    outliers_in_iteration += 1;
                    total_outliers_removed += 1;
                    tracing::debug!(
                        "Removed outlier (iteration {}): {:.3}s (mean: {:.3}s, std_dev: {:.3}s, threshold: {:.3}s)",
                        iteration + 1,
                        sample,
                        current_mean,
                        current_std_dev,
                        threshold
                    );
                }
            }

            // If no outliers found in this iteration, we're done
            if outliers_in_iteration == 0 {
                break;
            }

            current_samples = clean_samples;
        }

        (current_samples, total_outliers_removed)
    }

    /// Calculate 95% confidence interval using t-distribution
    ///
    /// Formula: CI = mean ± t_critical * (std_dev / sqrt(n))
    fn calculate_confidence_interval(
        samples: &[f64],
        mean: f64,
        std_dev: f64,
    ) -> Result<(f64, f64), MLError> {
        let n = samples.len() as f64;
        let degrees_of_freedom = n - 1.0;

        // Create t-distribution
        let t_dist =
            StudentsT::new(0.0, 1.0, degrees_of_freedom).map_err(|e| MLError::ValidationError {
                message: format!("Failed to create t-distribution: {}", e),
            })?;

        // Calculate t-critical value for 95% confidence (two-tailed)
        // For 95% CI, we need the 97.5th percentile (0.975)
        let t_critical = t_dist.inverse_cdf(0.975);

        // Calculate standard error
        let standard_error = std_dev / n.sqrt();

        // Calculate margin of error
        let margin_of_error = t_critical * standard_error;

        // Calculate confidence interval bounds
        let lower_bound = mean - margin_of_error;
        let upper_bound = mean + margin_of_error;

        Ok((lower_bound, upper_bound))
    }

    /// Calculate percentile using linear interpolation
    ///
    /// # Arguments
    ///
    /// * `sorted_samples` - Samples sorted in ascending order
    /// * `percentile` - Percentile to calculate (0.0 to 100.0)
    fn calculate_percentile(sorted_samples: &[f64], percentile: f64) -> f64 {
        if sorted_samples.is_empty() {
            return 0.0;
        }

        if sorted_samples.len() == 1 {
            return sorted_samples[0];
        }

        // Calculate position using linear interpolation between ranks
        let n = sorted_samples.len() as f64;
        let position = (percentile / 100.0) * (n - 1.0);

        // Get integer and fractional parts
        let lower_index = position.floor() as usize;
        let upper_index = position.ceil() as usize;
        let fraction = position - position.floor();

        // Handle edge cases
        if upper_index >= sorted_samples.len() {
            return sorted_samples[sorted_samples.len() - 1];
        }

        // Linear interpolation
        let lower_value = sorted_samples[lower_index];
        let upper_value = sorted_samples[upper_index];

        lower_value + fraction * (upper_value - lower_value)
    }

    /// Get number of samples collected
    pub fn sample_count(&self) -> usize {
        self.samples.len()
    }

    /// Get number of samples after warmup
    pub fn samples_after_warmup_count(&self) -> usize {
        self.samples.len().saturating_sub(self.warmup_epochs)
    }

    /// Clear all samples (useful for restarting collection)
    pub fn clear(&mut self) {
        self.samples.clear();
    }
}

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

    #[test]
    fn test_basic_statistics() {
        let mut sampler = StatisticalSampler::new(2);

        // Add warmup samples
        sampler.add_sample(10.0); // Warmup
        sampler.add_sample(9.5); // Warmup

        // Add actual samples
        let samples = vec![1.5, 1.6, 1.55, 1.58, 1.52, 1.54, 1.56, 1.57, 1.53, 1.59];
        for &sample in &samples {
            sampler.add_sample(sample);
        }

        let stats = sampler.compute_statistics().unwrap();

        assert_eq!(stats.num_samples, 10);
        assert_eq!(stats.outliers_removed, 0);
        assert_relative_eq!(stats.mean_seconds, 1.554, epsilon = 0.001);
        assert!(stats.std_dev > 0.0);
        assert!(stats.coefficient_of_variation < 0.15);
        assert!(stats.is_stable());
    }

    #[test]
    fn test_outlier_detection() {
        let mut sampler = StatisticalSampler::new(0);

        // Add normal samples
        for i in 0..10 {
            sampler.add_sample(1.5 + (i as f64 * 0.01));
        }

        // Add outliers
        sampler.add_sample(10.0); // Far outlier
        sampler.add_sample(0.1); // Far outlier

        let stats = sampler.compute_statistics().unwrap();

        assert_eq!(stats.outliers_removed, 2);
        assert_eq!(stats.num_samples, 10);
        assert_relative_eq!(stats.mean_seconds, 1.545, epsilon = 0.01);
    }

    #[test]
    fn test_percentile_calculation() {
        let mut sampler = StatisticalSampler::new(0);

        // Add samples from 1.0 to 2.0
        for i in 0..=100 {
            sampler.add_sample(1.0 + (i as f64 * 0.01));
        }

        let stats = sampler.compute_statistics().unwrap();

        assert_relative_eq!(stats.p50_median, 1.5, epsilon = 0.05);
        assert_relative_eq!(stats.p95, 1.95, epsilon = 0.05);
        assert_relative_eq!(stats.p99, 1.99, epsilon = 0.05);
    }

    #[test]
    fn test_confidence_interval() {
        let mut sampler = StatisticalSampler::new(0);

        // Add samples with known distribution
        let samples = vec![
            1.5, 1.6, 1.55, 1.58, 1.52, 1.54, 1.56, 1.57, 1.53, 1.59, 1.51, 1.62, 1.48, 1.61, 1.49,
        ];
        for &sample in &samples {
            sampler.add_sample(sample);
        }

        let stats = sampler.compute_statistics().unwrap();

        // CI should contain the mean
        assert!(stats.confidence_interval_95.0 < stats.mean_seconds);
        assert!(stats.confidence_interval_95.1 > stats.mean_seconds);

        // CI width should be reasonable (not too wide)
        let ci_width = stats.confidence_interval_95.1 - stats.confidence_interval_95.0;
        assert!(ci_width < stats.mean_seconds * 0.5); // Less than 50% of mean
    }

    #[test]
    fn test_insufficient_samples() {
        let mut sampler = StatisticalSampler::new(2);

        // Add only warmup samples
        sampler.add_sample(1.5);
        sampler.add_sample(1.6);

        let result = sampler.compute_statistics();
        assert!(result.is_err());
        assert!(matches!(result, Err(MLError::InsufficientData(_))));
    }

    #[test]
    fn test_zero_variance_error() {
        let mut sampler = StatisticalSampler::new(0);

        // Add identical samples
        for _ in 0..10 {
            sampler.add_sample(1.5);
        }

        let result = sampler.compute_statistics();
        assert!(result.is_err());
        assert!(matches!(result, Err(MLError::ValidationError { .. })));
    }

    #[test]
    fn test_high_coefficient_of_variation() {
        let mut sampler = StatisticalSampler::new(0);

        // Add samples with high variance
        sampler.add_sample(1.0);
        sampler.add_sample(2.0);
        sampler.add_sample(1.5);
        sampler.add_sample(2.5);
        sampler.add_sample(1.2);
        sampler.add_sample(2.3);
        sampler.add_sample(1.8);
        sampler.add_sample(2.1);
        sampler.add_sample(1.4);
        sampler.add_sample(2.4);

        let stats = sampler.compute_statistics();
        // Should still compute (just warns), but CV should be high
        assert!(stats.is_ok());
        if let Ok(s) = stats {
            assert!(s.coefficient_of_variation > 0.15);
            assert!(!s.is_stable());
        }
    }

    #[test]
    fn test_all_outliers_error() {
        let mut sampler = StatisticalSampler::new(0);

        // Add samples where all will be detected as outliers
        // (this is a pathological case)
        sampler.add_sample(1.0);
        sampler.add_sample(100.0);
        sampler.add_sample(1.1);
        sampler.add_sample(200.0);

        let result = sampler.compute_statistics();
        // Should handle gracefully - may succeed with remaining samples or fail
        // depending on which samples survive outlier detection
        assert!(result.is_ok() || matches!(result, Err(MLError::InsufficientData(_))));
    }

    #[test]
    fn test_margin_of_error() {
        let mut sampler = StatisticalSampler::new(0);

        let samples = vec![1.5, 1.6, 1.55, 1.58, 1.52, 1.54, 1.56, 1.57, 1.53, 1.59];
        for &sample in &samples {
            sampler.add_sample(sample);
        }

        let stats = sampler.compute_statistics().unwrap();
        let moe = stats.margin_of_error_95();

        // Margin of error should be positive and reasonable
        assert!(moe > 0.0);
        assert!(moe < stats.mean_seconds * 0.2); // Less than 20% of mean
    }

    #[test]
    fn test_outlier_percentage() {
        let mut sampler = StatisticalSampler::new(0);

        // 10 normal samples
        for i in 0..10 {
            sampler.add_sample(1.5 + (i as f64 * 0.01));
        }

        // 2 outliers
        sampler.add_sample(10.0);
        sampler.add_sample(0.1);

        let stats = sampler.compute_statistics().unwrap();
        let outlier_pct = stats.outlier_percentage();

        // Should be approximately 16.67% (2 out of 12)
        assert_relative_eq!(outlier_pct, 16.67, epsilon = 0.1);
    }

    #[test]
    fn test_clear_samples() {
        let mut sampler = StatisticalSampler::new(0);

        sampler.add_sample(1.5);
        sampler.add_sample(1.6);
        assert_eq!(sampler.sample_count(), 2);

        sampler.clear();
        assert_eq!(sampler.sample_count(), 0);
    }

    #[test]
    fn test_samples_after_warmup_count() {
        let mut sampler = StatisticalSampler::new(3);

        for i in 0..10 {
            sampler.add_sample(1.5 + (i as f64 * 0.01));
        }

        assert_eq!(sampler.sample_count(), 10);
        assert_eq!(sampler.samples_after_warmup_count(), 7);
    }

    #[test]
    fn test_known_distribution_statistics() {
        // Test with a known normal-like distribution
        let mut sampler = StatisticalSampler::new(0);

        // Create samples centered at 1.5 with small variance
        let mean_target = 1.5;
        let samples = vec![
            1.48, 1.49, 1.50, 1.51, 1.52, 1.48, 1.49, 1.50, 1.51, 1.52, 1.48, 1.49, 1.50, 1.51,
            1.52,
        ];

        for &sample in &samples {
            sampler.add_sample(sample);
        }

        let stats = sampler.compute_statistics().unwrap();

        // Mean should be close to target
        assert_relative_eq!(stats.mean_seconds, mean_target, epsilon = 0.01);

        // Median should be close to mean for symmetric distribution
        assert_relative_eq!(stats.p50_median, stats.mean_seconds, epsilon = 0.05);

        // CV should be low for tight distribution
        assert!(stats.coefficient_of_variation < 0.02);
        assert!(stats.is_stable());
    }
}