foxhunt/crates/ml-security/src/prediction_validator.rs

//! Prediction validation and model poisoning detection
//!
//! Provides statistical bounds checking and outlier detection to identify
//! poisoned or adversarial models in the ML inference pipeline.

use std::collections::VecDeque;
use std::time::{Duration, Instant};

use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{debug, warn};

use ml_core::MLError;

/// Prediction validator with statistical bounds checking
///
/// Detects model poisoning through:
/// 1. Range validation (predictions must be in [-1.0, 1.0])
/// 2. Z-score outlier detection (>3sigma from mean)
/// 3. Confidence sanity checks
/// 4. Extreme prediction rate limiting
///
/// # Performance
/// - Validation: ~5us per prediction
/// - Statistics update: ~2us
pub struct PredictionValidator {
    /// Historical prediction statistics (rolling window)
    stats: RwLock<PredictionStatistics>,

    /// Configuration
    config: ValidationConfig,

    /// Extreme prediction tracking (for rate limiting)
    extreme_tracker: RwLock<ExtremeTracker>,
}

impl std::fmt::Debug for PredictionValidator {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("PredictionValidator")
            .field("config", &self.config)
            .field("stats", &format_args!("<RwLock<PredictionStatistics>>"))
            .field("extreme_tracker", &format_args!("<RwLock<ExtremeTracker>>"))
            .finish()
    }
}

/// Configuration for prediction validation
#[derive(Debug, Clone)]
pub struct ValidationConfig {
    /// Z-score threshold for outlier detection (default: 3.0)
    pub z_score_threshold: f64,

    /// Minimum confidence threshold (default: 0.5)
    pub confidence_threshold: f64,

    /// Extreme prediction threshold (default: 0.9)
    pub extreme_prediction_threshold: f64,

    /// Maximum allowed extreme prediction rate (default: 0.1 = 10%)
    pub max_extreme_rate: f64,

    /// Window duration for rate limiting (default: 60 seconds)
    pub window_duration: Duration,

    /// Bootstrap sample count (default: 1000)
    pub bootstrap_samples: u64,

    /// Exponential moving average weight (default: 0.05)
    pub ema_alpha: f64,
}

impl Default for ValidationConfig {
    fn default() -> Self {
        Self {
            z_score_threshold: 3.0,
            confidence_threshold: 0.5,
            extreme_prediction_threshold: 0.9,
            max_extreme_rate: 0.1,
            window_duration: Duration::from_secs(60),
            bootstrap_samples: 1000,
            ema_alpha: 0.05,
        }
    }
}

/// Historical prediction statistics
#[derive(Debug, Clone)]
struct PredictionStatistics {
    /// Exponential moving average of predictions
    mean: f64,

    /// Standard deviation (Welford's online algorithm)
    std_dev: f64,

    /// Variance (for std_dev calculation)
    variance: f64,

    /// Total number of samples processed
    sample_count: u64,

    /// Last update timestamp
    last_updated: Instant,

    /// Sum of squared differences (for variance)
    m2: f64,
}

impl PredictionStatistics {
    fn new() -> Self {
        Self {
            mean: 0.0,
            std_dev: 0.3, // Conservative bootstrap value
            variance: 0.09,
            sample_count: 0,
            last_updated: Instant::now(),
            m2: 0.0,
        }
    }

    /// Update statistics with new prediction (Welford's algorithm)
    fn update(&mut self, prediction: f64, ema_alpha: f64, bootstrap_samples: u64) {
        self.sample_count += 1;

        if self.sample_count <= bootstrap_samples {
            // Bootstrap phase: Use Welford's algorithm for stable variance
            let delta = prediction - self.mean;
            self.mean += delta / self.sample_count as f64;
            let delta2 = prediction - self.mean;
            self.m2 += delta * delta2;

            if self.sample_count > 1 {
                self.variance = self.m2 / (self.sample_count - 1) as f64;
                self.std_dev = self.variance.sqrt().max(0.01); // Minimum 0.01 to avoid division by zero
            }
        } else {
            // Production phase: Use exponential moving average
            self.mean = (1.0 - ema_alpha) * self.mean + ema_alpha * prediction;

            // Update variance with EMA
            let delta = prediction - self.mean;
            self.variance = (1.0 - ema_alpha) * self.variance + ema_alpha * delta * delta;
            self.std_dev = self.variance.sqrt().max(0.01);
        }

        self.last_updated = Instant::now();
    }
}

/// Extreme prediction tracker for rate limiting
#[derive(Debug)]
struct ExtremeTracker {
    /// Recent extreme prediction timestamps
    timestamps: VecDeque<Instant>,
}

impl ExtremeTracker {
    fn new() -> Self {
        Self {
            timestamps: VecDeque::new(),
        }
    }

    /// Record an extreme prediction
    fn record_extreme(&mut self, now: Instant, window_duration: Duration) {
        // Remove old timestamps outside window
        while let Some(oldest) = self.timestamps.front() {
            if now.duration_since(*oldest) > window_duration {
                self.timestamps.pop_front();
            } else {
                break;
            }
        }

        self.timestamps.push_back(now);
    }

    /// Calculate current extreme prediction rate
    fn calculate_rate(&self, window_duration: Duration) -> f64 {
        let now = Instant::now();
        let count = self
            .timestamps
            .iter()
            .filter(|&&ts| now.duration_since(ts) <= window_duration)
            .count();

        // Estimate total predictions based on typical inference rate (~1000/sec)
        let window_secs = window_duration.as_secs_f64();
        let estimated_total = (window_secs * 1000.0).max(1.0);

        count as f64 / estimated_total
    }
}

/// Validated prediction result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ValidatedPrediction {
    /// Original prediction value
    pub value: f64,

    /// Whether prediction is an outlier
    pub is_outlier: bool,

    /// Z-score (standard deviations from mean)
    pub z_score: f64,

    /// Whether to override with ensemble fallback
    pub should_override: bool,

    /// Validation flags
    pub validation_flags: Vec<ValidationFlag>,
}

/// Validation flags for different failure modes
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ValidationFlag {
    /// Prediction outside valid range [-1.0, 1.0]
    OutOfBounds { value: f64 },

    /// Prediction is statistical outlier
    Outlier { z_score: f64 },

    /// Confidence below threshold
    LowConfidence { confidence: f64 },

    /// Extreme prediction rate exceeded
    ExtremeRateLimit { rate: f64 },
}

impl PredictionValidator {
    /// Create a new prediction validator with default configuration
    pub fn new() -> Self {
        Self::with_config(ValidationConfig::default())
    }

    /// Create a new prediction validator with custom configuration
    pub fn with_config(config: ValidationConfig) -> Self {
        Self {
            stats: RwLock::new(PredictionStatistics::new()),
            config,
            extreme_tracker: RwLock::new(ExtremeTracker::new()),
        }
    }

    /// Validate prediction with multi-layer checks
    ///
    /// # Arguments
    /// * `prediction` - Prediction value to validate
    /// * `confidence` - Model confidence (0.0-1.0)
    /// * `model_id` - Model identifier for logging
    ///
    /// # Returns
    /// ValidatedPrediction with flags and z-score
    ///
    /// # Errors
    /// Returns error if prediction is invalid (out of bounds, extreme rate exceeded)
    ///
    /// # Performance
    /// ~5us per call (includes statistics update)
    pub async fn validate(
        &self,
        prediction: f64,
        confidence: f64,
        model_id: &str,
    ) -> Result<ValidatedPrediction, MLError> {
        let mut flags = Vec::new();

        // Layer 1: Range check
        if !(-1.0..=1.0).contains(&prediction) {
            warn!(
                "Prediction out of bounds for model {}: {} (valid range: [-1.0, 1.0])",
                model_id, prediction
            );

            flags.push(ValidationFlag::OutOfBounds { value: prediction });

            return Err(MLError::ValidationError {
                message: format!(
                    "Prediction {} out of bounds [-1.0, 1.0] for model {}",
                    prediction, model_id
                ),
            });
        }

        // Layer 2: Z-score outlier detection
        let stats = self.stats.read().await;
        let z_score = if stats.std_dev > 0.001 {
            (prediction - stats.mean) / stats.std_dev
        } else {
            0.0
        };
        drop(stats); // Release read lock

        if z_score.abs() > self.config.z_score_threshold {
            warn!(
                "Outlier prediction detected for model {}: {} (z-score: {:.2}, mean: {:.3}, std: {:.3})",
                model_id,
                prediction,
                z_score,
                self.stats.read().await.mean,
                self.stats.read().await.std_dev
            );

            flags.push(ValidationFlag::Outlier { z_score });
        }

        // Layer 3: Confidence sanity check
        if confidence < self.config.confidence_threshold {
            warn!(
                "Low confidence prediction for model {}: {:.3}",
                model_id, confidence
            );

            flags.push(ValidationFlag::LowConfidence { confidence });
        }

        // Layer 4: Rate limit extreme predictions
        let is_extreme = prediction.abs() > self.config.extreme_prediction_threshold;
        if is_extreme {
            let now = Instant::now();
            let mut tracker = self.extreme_tracker.write().await;
            tracker.record_extreme(now, self.config.window_duration);

            let extreme_rate = tracker.calculate_rate(self.config.window_duration);

            if extreme_rate > self.config.max_extreme_rate {
                warn!(
                    "Extreme prediction rate exceeded for model {}: {:.3} (threshold: {:.3})",
                    model_id, extreme_rate, self.config.max_extreme_rate
                );

                flags.push(ValidationFlag::ExtremeRateLimit { rate: extreme_rate });

                return Err(MLError::ValidationError {
                    message: format!(
                        "Too many extreme predictions from model {} - possible poisoning (rate: {:.3}, threshold: {:.3})",
                        model_id, extreme_rate, self.config.max_extreme_rate
                    ),
                });
            }
        }

        // Update statistics (after validation to avoid corrupting stats with invalid data)
        self.update_statistics(prediction).await;

        let is_outlier = !flags.is_empty();
        let should_override = z_score.abs() > self.config.z_score_threshold;

        debug!(
            "Validated prediction for {}: value={:.3}, z={:.2}, outlier={}, override={}",
            model_id, prediction, z_score, is_outlier, should_override
        );

        Ok(ValidatedPrediction {
            value: prediction,
            is_outlier,
            z_score,
            should_override,
            validation_flags: flags,
        })
    }

    /// Update prediction statistics with new value
    pub async fn update_statistics(&self, prediction: f64) {
        let mut stats = self.stats.write().await;
        stats.update(
            prediction,
            self.config.ema_alpha,
            self.config.bootstrap_samples,
        );
    }

    /// Get current prediction statistics
    pub async fn get_statistics(&self) -> PredictionStats {
        let stats = self.stats.read().await;
        PredictionStats {
            mean: stats.mean,
            std_dev: stats.std_dev,
            sample_count: stats.sample_count,
        }
    }

    /// Reset statistics (useful for testing or model retraining)
    pub async fn reset_statistics(&self) {
        let mut stats = self.stats.write().await;
        *stats = PredictionStatistics::new();

        let mut tracker = self.extreme_tracker.write().await;
        tracker.timestamps.clear();
    }
}

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

/// Public statistics snapshot
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PredictionStats {
    pub mean: f64,
    pub std_dev: f64,
    pub sample_count: u64,
}

#[cfg(test)]
#[allow(
    clippy::let_underscore_must_use,
    clippy::assertions_on_result_states,
    clippy::field_reassign_with_default
)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_validate_normal_prediction() {
        let validator = PredictionValidator::new();

        // Add some bootstrap samples
        for i in 0..100 {
            let _ = validator
                .validate(0.0 + (i as f64 / 1000.0), 0.8, "test_model")
                .await;
        }

        // Validate normal prediction
        let result = validator.validate(0.05, 0.8, "test_model").await;
        assert!(result.is_ok());

        let validated = result.unwrap();
        assert!(!validated.should_override);
        assert!(validated.validation_flags.is_empty());
    }

    #[tokio::test]
    async fn test_validate_out_of_bounds() {
        let validator = PredictionValidator::new();

        // Test upper bound
        let result = validator.validate(1.5, 0.8, "test_model").await;
        assert!(result.is_err());

        // Test lower bound
        let result = validator.validate(-1.5, 0.8, "test_model").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_validate_outlier() {
        let validator = PredictionValidator::new();

        // Build statistics with normal predictions around 0.0
        for _ in 0..1000 {
            let _ = validator.validate(0.0, 0.8, "test_model").await;
        }

        // Inject extreme outlier
        let result = validator.validate(0.95, 0.8, "test_model").await.unwrap();

        assert!(result.is_outlier);
        assert!(result.z_score.abs() > 3.0);
        assert!(result.should_override);
    }

    #[tokio::test]
    async fn test_low_confidence_flag() {
        let validator = PredictionValidator::new();

        let result = validator.validate(0.5, 0.3, "test_model").await.unwrap();

        // Should flag low confidence but not reject
        let has_low_conf = result
            .validation_flags
            .iter()
            .any(|f| matches!(f, ValidationFlag::LowConfidence { .. }));

        assert!(has_low_conf);
    }

    #[tokio::test]
    async fn test_extreme_rate_limiting() {
        let mut config = ValidationConfig::default();
        config.max_extreme_rate = 0.05; // 5%
        config.window_duration = Duration::from_secs(1);

        let validator = PredictionValidator::with_config(config);

        // Inject many extreme predictions quickly
        for _ in 0..100 {
            let result = validator.validate(0.95, 0.8, "test_model").await;

            // Should eventually hit rate limit
            if result.is_err() {
                // Verify it's a rate limit error
                let err = result.unwrap_err();
                assert!(err.to_string().contains("extreme predictions"));
                return;
            }

            tokio::time::sleep(Duration::from_millis(1)).await;
        }

        panic!("Rate limit should have been triggered");
    }

    #[tokio::test]
    async fn test_statistics_update() {
        let validator = PredictionValidator::new();

        // Add predictions
        for i in 0..100 {
            let value = (i as f64 / 100.0) - 0.5; // Range: -0.5 to 0.5
            validator.update_statistics(value).await;
        }

        let stats = validator.get_statistics().await;

        assert!(stats.sample_count == 100);
        assert!(stats.mean.abs() < 0.1); // Should be close to 0
        assert!(stats.std_dev > 0.0);
    }

    #[tokio::test]
    async fn test_reset_statistics() {
        let validator = PredictionValidator::new();

        // Add some data
        for _ in 0..100 {
            validator.update_statistics(0.5).await;
        }

        let stats_before = validator.get_statistics().await;
        assert!(stats_before.sample_count == 100);

        // Reset
        validator.reset_statistics().await;

        let stats_after = validator.get_statistics().await;
        assert_eq!(stats_after.sample_count, 0);
        assert_eq!(stats_after.mean, 0.0);
    }

    #[tokio::test]
    async fn test_bootstrap_phase() {
        let mut config = ValidationConfig::default();
        config.bootstrap_samples = 10;

        let validator = PredictionValidator::with_config(config);

        // Add bootstrap samples
        for i in 0..10 {
            validator.update_statistics(i as f64 / 10.0).await;
        }

        let stats = validator.get_statistics().await;
        assert!(stats.std_dev > 0.0);
    }
}