foxhunt/crates/ml/src/benchmark/stability_validator.rs

/// Metrics describing training stability
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, Default)]
pub struct StabilityMetrics {
    pub is_stable: bool,
    pub has_nan_inf: bool,
    pub gradient_health: GradientHealth,
    pub loss_trend: LossTrend,
    pub warnings: Vec<String>,
}

/// Health classification for gradient norms
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, Default)]
pub enum GradientHealth {
    #[default]
    /// Gradient norm in healthy range (1e-6 < norm < 10.0)
    Healthy,
    /// Gradient norm too large (norm >= 10.0), likely learning rate too high
    Exploding,
    /// Gradient norm too small (norm <= 1e-6), likely learning rate too low
    Vanishing,
}

/// Trend classification for loss values over time
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize, Default)]
pub enum LossTrend {
    #[default]
    /// Loss consistently decreasing
    Converging,
    /// Loss consistently increasing (training diverging)
    Diverging,
    /// Loss flat for >5 epochs (no progress)
    Stagnant,
}

/// Validator for training stability across epochs
#[derive(Debug)]
pub struct StabilityValidator {
    loss_history: Vec<f64>,
    gradient_norms: Vec<f64>,
    stagnant_threshold: usize,
}

impl StabilityValidator {
    /// Create a new stability validator
    pub fn new() -> Self {
        Self {
            loss_history: Vec::new(),
            gradient_norms: Vec::new(),
            stagnant_threshold: 5,
        }
    }

    /// Create a validator with custom stagnant detection threshold
    pub fn with_stagnant_threshold(threshold: usize) -> Self {
        Self {
            loss_history: Vec::new(),
            gradient_norms: Vec::new(),
            stagnant_threshold: threshold,
        }
    }

    /// Record a loss value from an epoch
    pub fn record_loss(&mut self, loss: f64) {
        self.loss_history.push(loss);
    }

    /// Record a gradient norm from an epoch
    pub fn record_gradient_norm(&mut self, norm: f64) {
        self.gradient_norms.push(norm);
    }

    /// Calculate gradient norm from a flat f32 slice (L2 norm).
    pub fn calculate_gradient_norm_from_slice(&self, data: &[f32]) -> f64 {
        let sum_sq: f64 = data.iter().map(|&x| (x as f64) * (x as f64)).sum();
        sum_sq.sqrt()
    }

    /// Validate training stability and return metrics
    pub fn validate(&self) -> StabilityMetrics {
        let mut warnings = Vec::new();

        // Check for NaN/Inf in loss history
        let has_nan_inf_loss = self
            .loss_history
            .iter()
            .any(|&loss| loss.is_nan() || loss.is_infinite());

        // Check for NaN/Inf in gradient norms
        let has_nan_inf_grads = self
            .gradient_norms
            .iter()
            .any(|&norm| norm.is_nan() || norm.is_infinite());

        let has_nan_inf = has_nan_inf_loss || has_nan_inf_grads;

        if has_nan_inf_loss {
            warnings.push("NaN or Inf detected in loss values".to_owned());
        }
        if has_nan_inf_grads {
            warnings.push("NaN or Inf detected in gradient norms".to_owned());
        }

        // Analyze gradient health (use most recent gradient norm)
        let gradient_health = if let Some(&latest_norm) = self.gradient_norms.last() {
            if latest_norm.is_nan() || latest_norm.is_infinite() {
                warnings.push("Invalid gradient norm value".to_owned());
                GradientHealth::Exploding
            } else if latest_norm >= 10.0 {
                warnings.push(format!(
                    "Exploding gradients detected (norm: {:.2e})",
                    latest_norm
                ));
                GradientHealth::Exploding
            } else if latest_norm <= 1e-6 {
                warnings.push(format!(
                    "Vanishing gradients detected (norm: {:.2e})",
                    latest_norm
                ));
                GradientHealth::Vanishing
            } else {
                GradientHealth::Healthy
            }
        } else {
            warnings.push("No gradient norms recorded".to_owned());
            GradientHealth::Healthy
        };

        // Analyze loss trend
        let loss_trend = self.analyze_loss_trend(&mut warnings);

        // Overall stability: no NaN/Inf, healthy gradients, and converging/stagnant loss
        let is_stable = !has_nan_inf
            && gradient_health == GradientHealth::Healthy
            && loss_trend != LossTrend::Diverging;

        StabilityMetrics {
            is_stable,
            has_nan_inf,
            gradient_health,
            loss_trend,
            warnings,
        }
    }

    /// Analyze loss trend over recent epochs
    fn analyze_loss_trend(&self, warnings: &mut Vec<String>) -> LossTrend {
        if self.loss_history.len() < 2 {
            return LossTrend::Converging; // Not enough data
        }

        // Use last 5 epochs or all available if less than 5
        let window_size = self.stagnant_threshold.min(self.loss_history.len());
        let recent_losses: Vec<f64> = self
            .loss_history
            .iter()
            .rev()
            .take(window_size)
            .copied()
            .collect();

        // Check if all recent losses are very similar (stagnant)
        if self.loss_history.len() >= self.stagnant_threshold {
            let mean = recent_losses.iter().sum::<f64>() / recent_losses.len() as f64;
            let variance = recent_losses
                .iter()
                .map(|&x| (x - mean).powi(2))
                .sum::<f64>()
                / recent_losses.len() as f64;
            let std_dev = variance.sqrt();

            // If std dev is very small relative to mean, loss is stagnant
            let relative_std = if mean.abs() > 1e-10 {
                std_dev / mean.abs()
            } else {
                std_dev
            };

            if relative_std < 0.001 {
                warnings.push(format!(
                    "Loss stagnant for {} epochs (relative std: {:.2e})",
                    window_size, relative_std
                ));
                return LossTrend::Stagnant;
            }
        }

        // Calculate moving average trend (are we going up or down?)
        let first_half_mean = recent_losses
            .iter()
            .skip(window_size / 2)
            .copied()
            .sum::<f64>()
            / (window_size - window_size / 2) as f64;

        let second_half_mean = recent_losses
            .iter()
            .take(window_size / 2)
            .copied()
            .sum::<f64>()
            / (window_size / 2) as f64;

        // Note: recent_losses is reversed, so second_half is more recent
        if second_half_mean < first_half_mean * 0.999 {
            // Loss is decreasing (second half lower than first half)
            LossTrend::Converging
        } else if second_half_mean > first_half_mean * 1.001 {
            // Loss is increasing (diverging!)
            warnings.push(format!(
                "Loss diverging: increased from {:.6} to {:.6}",
                first_half_mean, second_half_mean
            ));
            LossTrend::Diverging
        } else {
            // Loss relatively flat but not stagnant yet
            LossTrend::Converging
        }
    }

    /// Clear all recorded metrics
    pub fn clear(&mut self) {
        self.loss_history.clear();
        self.gradient_norms.clear();
    }

    /// Get the number of recorded epochs
    pub fn epoch_count(&self) -> usize {
        self.loss_history.len()
    }

    /// Get the loss history
    pub fn loss_history(&self) -> &[f64] {
        &self.loss_history
    }

    /// Get the gradient norm history
    pub fn gradient_norms(&self) -> &[f64] {
        &self.gradient_norms
    }
}

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

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

    fn _cuda_device() -> ml_core::native_types::NativeDevice {
        ml_core::native_types::NativeDevice::Cuda(0)
    }


    #[test]
    fn test_converging_loss() {
        let mut validator = StabilityValidator::new();

        // Simulate converging training
        for i in 0..10 {
            let loss = 1.0 / (i as f64 + 1.0); // Decreasing loss
            let grad_norm = 0.1; // Healthy gradient
            validator.record_loss(loss);
            validator.record_gradient_norm(grad_norm);
        }

        let metrics = validator.validate();
        assert!(metrics.is_stable);
        assert!(!metrics.has_nan_inf);
        assert_eq!(metrics.gradient_health, GradientHealth::Healthy);
        assert_eq!(metrics.loss_trend, LossTrend::Converging);
    }


    #[test]
    fn test_diverging_loss() {
        let mut validator = StabilityValidator::new();

        // Simulate diverging training
        for i in 0..10 {
            let loss = (i as f64 + 1.0) * 0.5; // Increasing loss
            let grad_norm = 0.1;
            validator.record_loss(loss);
            validator.record_gradient_norm(grad_norm);
        }

        let metrics = validator.validate();
        assert!(!metrics.is_stable); // Diverging is unstable
        assert_eq!(metrics.loss_trend, LossTrend::Diverging);
    }


    #[test]
    fn test_stagnant_loss() {
        let mut validator = StabilityValidator::new();

        // Simulate stagnant training
        for _ in 0..10 {
            validator.record_loss(1.0); // Same loss every epoch
            validator.record_gradient_norm(0.1);
        }

        let metrics = validator.validate();
        assert!(metrics.is_stable); // Stagnant is stable, just not improving
        assert_eq!(metrics.loss_trend, LossTrend::Stagnant);
    }


    #[test]
    fn test_nan_detection() {
        let mut validator = StabilityValidator::new();

        validator.record_loss(1.0);
        validator.record_loss(f64::NAN);
        validator.record_gradient_norm(0.1);

        let metrics = validator.validate();
        assert!(!metrics.is_stable);
        assert!(metrics.has_nan_inf);
        assert!(metrics.warnings.iter().any(|w| w.contains("NaN")));
    }


    #[test]
    fn test_inf_detection() {
        let mut validator = StabilityValidator::new();

        validator.record_loss(1.0);
        validator.record_loss(f64::INFINITY);
        validator.record_gradient_norm(0.1);

        let metrics = validator.validate();
        assert!(!metrics.is_stable);
        assert!(metrics.has_nan_inf);
    }


    #[test]
    fn test_exploding_gradients() {
        let mut validator = StabilityValidator::new();

        validator.record_loss(1.0);
        validator.record_loss(0.9);
        validator.record_gradient_norm(15.0); // Exploding!

        let metrics = validator.validate();
        assert!(!metrics.is_stable);
        assert_eq!(metrics.gradient_health, GradientHealth::Exploding);
        assert!(metrics.warnings.iter().any(|w| w.contains("Exploding")));
    }


    #[test]
    fn test_vanishing_gradients() {
        let mut validator = StabilityValidator::new();

        validator.record_loss(1.0);
        validator.record_loss(0.9);
        validator.record_gradient_norm(1e-7); // Vanishing!

        let metrics = validator.validate();
        assert!(!metrics.is_stable);
        assert_eq!(metrics.gradient_health, GradientHealth::Vanishing);
        assert!(metrics.warnings.iter().any(|w| w.contains("Vanishing")));
    }


    #[test]
    fn test_healthy_gradients() {
        let mut validator = StabilityValidator::new();

        validator.record_loss(1.0);
        validator.record_gradient_norm(1.0); // Healthy range

        let metrics = validator.validate();
        assert_eq!(metrics.gradient_health, GradientHealth::Healthy);
    }


    #[test]
    fn test_gradient_norm_calculation() {
        let validator = StabilityValidator::new();

        // Known L2 norm of [3, 4] is 5
        let data = [3.0_f32, 4.0_f32];
        let norm = validator.calculate_gradient_norm_from_slice(&data);

        assert!((norm - 5.0).abs() < 1e-6);
    }


    #[test]
    fn test_nan_in_gradients() {
        let mut validator = StabilityValidator::new();

        validator.record_loss(1.0);
        validator.record_gradient_norm(f64::NAN);

        let metrics = validator.validate();
        assert!(!metrics.is_stable);
        assert!(metrics.has_nan_inf);
    }


    #[test]
    fn test_early_divergence_detection() {
        let mut validator = StabilityValidator::new();

        // Simulate divergence within 5 epochs
        for i in 0..5 {
            let loss = 1.0 + (i as f64) * 0.2; // Clearly increasing
            validator.record_loss(loss);
            validator.record_gradient_norm(0.1);
        }

        let metrics = validator.validate();
        assert_eq!(metrics.loss_trend, LossTrend::Diverging);
    }


    #[test]
    fn test_clear_metrics() {
        let mut validator = StabilityValidator::new();

        validator.record_loss(1.0);
        validator.record_gradient_norm(0.1);
        assert_eq!(validator.epoch_count(), 1);

        validator.clear();
        assert_eq!(validator.epoch_count(), 0);
        assert_eq!(validator.loss_history().len(), 0);
        assert_eq!(validator.gradient_norms().len(), 0);
    }


    #[test]
    fn test_custom_stagnant_threshold() {
        let mut validator = StabilityValidator::with_stagnant_threshold(3);

        // Only 3 epochs needed for stagnant detection
        for _ in 0..3 {
            validator.record_loss(1.0);
            validator.record_gradient_norm(0.1);
        }

        let metrics = validator.validate();
        assert_eq!(metrics.loss_trend, LossTrend::Stagnant);
    }


    #[test]
    fn test_insufficient_data() {
        let mut validator = StabilityValidator::new();

        // Only one epoch
        validator.record_loss(1.0);
        validator.record_gradient_norm(0.1);

        let metrics = validator.validate();
        assert!(metrics.is_stable);
        assert_eq!(metrics.loss_trend, LossTrend::Converging); // Default with insufficient data
    }


    #[test]
    fn test_all_nan_scenario() {
        let mut validator = StabilityValidator::new();

        // Everything is NaN
        for _ in 0..5 {
            validator.record_loss(f64::NAN);
            validator.record_gradient_norm(f64::NAN);
        }

        let metrics = validator.validate();
        assert!(!metrics.is_stable);
        assert!(metrics.has_nan_inf);
        assert!(metrics.warnings.len() >= 2); // Should have warnings for both loss and grads
    }


    #[test]
    fn test_mixed_gradient_health() {
        let mut validator = StabilityValidator::new();

        // Record several gradient norms, only latest matters
        validator.record_gradient_norm(15.0); // Exploding
        validator.record_gradient_norm(1e-7); // Vanishing
        validator.record_gradient_norm(1.0); // Healthy (this is what counts)

        validator.record_loss(1.0);

        let metrics = validator.validate();
        assert_eq!(metrics.gradient_health, GradientHealth::Healthy);
    }


    #[test]
    fn test_no_gradients_recorded() {
        let mut validator = StabilityValidator::new();

        validator.record_loss(1.0);
        validator.record_loss(0.9);
        // No gradient norms recorded

        let metrics = validator.validate();
        assert_eq!(metrics.gradient_health, GradientHealth::Healthy); // Default when no data
        assert!(metrics
            .warnings
            .iter()
            .any(|w| w.contains("No gradient norms")));
    }
}