foxhunt/crates/ml-stress-testing/src/lib.rs

#![deny(clippy::unwrap_used, clippy::expect_used)]
#![cfg_attr(test, allow(clippy::unwrap_used, clippy::expect_used))]
#![allow(clippy::module_name_repetitions)]
#![allow(clippy::integer_division)]
#![allow(dead_code)]
#![allow(missing_docs)]
#![allow(missing_debug_implementations)]
#![allow(clippy::float_arithmetic)]
#![allow(clippy::doc_markdown)]
#![allow(clippy::indexing_slicing)]
#![allow(clippy::missing_const_for_fn)]
#![allow(clippy::shadow_reuse, clippy::shadow_same, clippy::shadow_unrelated)] // Tensor ops: let x = x.relu() is idiomatic
#![allow(clippy::non_ascii_literal)] // Math symbols in ML documentation and error messages
#![allow(clippy::partial_pub_fields)] // ML config structs: some fields are pub API, some internal
#![allow(clippy::same_name_method)] // Intentional: inherent methods shadow trait defaults for ML-specific behavior

//! Stress testing framework for ML models under high-volume market data conditions
//!
//! This crate provides comprehensive stress testing capabilities to validate ML model
//! performance under realistic HFT market conditions with high-frequency data feeds.

// Re-export key types from ml-core that are used throughout this crate
pub use common::model_types::ModelType;
pub use ml_core::{Features, MLModel, ModelPrediction};

pub mod load_generator;
pub mod market_simulator;
pub mod performance_analyzer;

// Re-export key types that are needed by external users
pub use load_generator::{LoadGenerator, LoadProfile, TrafficPattern};
pub use market_simulator::{MarketCondition, MarketDataSimulator, SimulatorConfig};
pub use performance_analyzer::{LatencyStats, PerformanceAnalyzer, StressTestReport};

use common::types::Price;
use rust_decimal::Decimal;

use anyhow::Result;
use rust_decimal::prelude::ToPrimitive;
use serde::{Deserialize, Serialize};
use std::time::{Duration, Instant};
use tokio::sync::mpsc;

/// Comprehensive stress test configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StressTestConfig {
    /// Test duration in seconds
    pub duration_seconds: u64,
    /// Target requests per second
    pub target_rps: u32,
    /// Number of concurrent connections
    pub concurrent_connections: u32,
    /// Market data feed rate (updates per second)
    pub market_data_rate: u32,
    /// Test phases with different load patterns
    pub test_phases: Vec<TestPhase>,
    /// Models to test
    pub models_to_test: Vec<String>,
    /// Market conditions to simulate
    pub market_conditions: Vec<MarketCondition>,
    /// Performance requirements
    pub requirements: PerformanceRequirements,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TestPhase {
    pub name: String,
    pub duration_seconds: u64,
    pub load_multiplier: f64,
    pub market_volatility: f64,
    pub error_injection_rate: f64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PerformanceRequirements {
    /// Maximum acceptable latency in microseconds
    pub max_latency_us: u64,
    /// 95th percentile latency requirement
    pub p95_latency_us: u64,
    /// 99th percentile latency requirement
    pub p99_latency_us: u64,
    /// Maximum acceptable error rate
    pub max_error_rate: f64,
    /// Minimum throughput (predictions per second)
    pub min_throughput: u32,
    /// Maximum memory usage in MB
    pub max_memory_mb: u64,
    /// Maximum CPU utilization percentage
    pub max_cpu_percent: f64,
}

impl Default for PerformanceRequirements {
    fn default() -> Self {
        Self {
            max_latency_us: 100,   // 100us max latency
            p95_latency_us: 50,    // 50us 95th percentile
            p99_latency_us: 80,    // 80us 99th percentile
            max_error_rate: 0.01,  // 1% max error rate
            min_throughput: 10000, // 10k predictions/sec
            max_memory_mb: 1024,   // 1GB max memory
            max_cpu_percent: 80.0, // 80% max CPU
        }
    }
}

/// Main stress testing orchestrator
#[derive(Debug)]
pub struct StressTestOrchestrator {
    config: StressTestConfig,
    market_simulator: MarketDataSimulator,
    load_generator: LoadGenerator,
    performance_analyzer: PerformanceAnalyzer,
}

impl StressTestOrchestrator {
    /// Create new stress test orchestrator with configuration-driven market simulation
    pub fn new(config: StressTestConfig) -> Result<Self> {
        // Use configuration-driven approach - load from config crate
        let simulation_config = config::SimulationConfig::default();

        // Use test fixtures for symbol configuration
        let test_symbols = vec!["TEST_LARGE_1", "TEST_LARGE_2", "TEST_MID_1", "TEST_SMALL_1"];
        let simulator_config = SimulatorConfig {
            symbols: test_symbols.into_iter().map(|s| s.to_owned()).collect(),
            update_rate_hz: config.market_data_rate,
            volatility: 0.02,
            trend: 0.0,
            simulation_config: Some(simulation_config),
        };

        let market_simulator = MarketDataSimulator::new(simulator_config)?;
        let load_generator = LoadGenerator::new(config.target_rps, config.concurrent_connections)?;
        let performance_analyzer = PerformanceAnalyzer::new();

        Ok(Self {
            config,
            market_simulator,
            load_generator,
            performance_analyzer,
        })
    }

    /// Create stress test orchestrator with custom simulation configuration
    pub fn new_with_simulation_config(
        config: StressTestConfig,
        simulation_config: config::SimulationConfig,
    ) -> Result<Self> {
        let symbols: Vec<String> = simulation_config
            .initial_market_state
            .symbols
            .keys()
            .cloned()
            .collect();

        let simulator_config = SimulatorConfig {
            symbols,
            update_rate_hz: config.market_data_rate,
            volatility: simulation_config.parameters.base_volatility,
            trend: simulation_config.parameters.trend,
            simulation_config: Some(simulation_config),
        };

        let market_simulator = MarketDataSimulator::new(simulator_config)?;
        let load_generator = LoadGenerator::new(config.target_rps, config.concurrent_connections)?;
        let performance_analyzer = PerformanceAnalyzer::new();

        Ok(Self {
            config,
            market_simulator,
            load_generator,
            performance_analyzer,
        })
    }

    /// Create stress test orchestrator for testing with generated test symbols
    pub fn new_for_testing(config: StressTestConfig) -> Result<Self> {
        let simulation_config = config::SimulationConfig::default();
        let test_symbols = MarketDataSimulator::generate_test_symbols(&simulation_config);

        let simulator_config = SimulatorConfig {
            symbols: test_symbols,
            update_rate_hz: config.market_data_rate,
            volatility: simulation_config.parameters.base_volatility,
            trend: simulation_config.parameters.trend,
            simulation_config: Some(simulation_config),
        };

        let market_simulator = MarketDataSimulator::new(simulator_config)?;
        let load_generator = LoadGenerator::new(config.target_rps, config.concurrent_connections)?;
        let performance_analyzer = PerformanceAnalyzer::new();

        Ok(Self {
            config,
            market_simulator,
            load_generator,
            performance_analyzer,
        })
    }

    /// Run comprehensive stress test
    pub async fn run_stress_test(
        &mut self,
        models: Vec<std::sync::Arc<dyn MLModel>>,
    ) -> Result<StressTestReport> {
        tracing::info!("Starting stress test with {} models", models.len());

        // Create channels for communication
        let (market_tx, mut market_rx) = mpsc::channel(10000);
        let (prediction_tx, _prediction_rx) = mpsc::channel(10000);

        // Start market data simulation
        let simulator_handle = {
            let mut simulator = self.market_simulator.clone();
            tokio::spawn(async move { simulator.start_simulation(market_tx).await })
        };

        // Start performance monitoring
        let analyzer = self.performance_analyzer.clone();
        let monitor_handle = tokio::spawn(async move { analyzer.start_monitoring().await });

        // Execute test phases
        let mut phase_results = Vec::new();
        let test_start = Instant::now();

        // Clone the test phases to avoid borrowing conflicts
        let test_phases = self.config.test_phases.clone();
        for (phase_idx, phase) in test_phases.into_iter().enumerate() {
            tracing::info!("Starting test phase {}: {}", phase_idx + 1, phase.name);

            let phase_result = self
                .run_test_phase(&phase, &models, &mut market_rx, &prediction_tx)
                .await?;

            phase_results.push(phase_result);
        }

        // Stop simulation and monitoring
        simulator_handle.abort();
        monitor_handle.abort();

        // Generate comprehensive report
        let total_duration = test_start.elapsed();
        let report = self
            .generate_stress_test_report(phase_results, total_duration)
            .await?;

        tracing::info!(
            "Stress test completed in {:.2}s",
            total_duration.as_secs_f64()
        );
        Ok(report)
    }

    /// Run individual test phase
    async fn run_test_phase(
        &mut self,
        phase: &TestPhase,
        models: &[std::sync::Arc<dyn MLModel>],
        market_rx: &mut mpsc::Receiver<MarketDataUpdate>,
        prediction_tx: &mpsc::Sender<PredictionResult>,
    ) -> Result<PhaseResult> {
        let phase_start = Instant::now();
        let phase_duration = Duration::from_secs(phase.duration_seconds);

        let mut phase_stats = PhaseStats::new();

        // Adjust load generator for this phase
        self.load_generator
            .set_load_multiplier(phase.load_multiplier);

        while phase_start.elapsed() < phase_duration {
            // Process market data updates
            while let Ok(market_update) = market_rx.try_recv() {
                // Convert market data to features
                let features = self.convert_market_data_to_features(&market_update)?;

                // Run predictions on all models
                for model in models {
                    let model_start = Instant::now();

                    match model.predict(&features).await {
                        Ok(prediction) => {
                            let latency_us = model_start.elapsed().as_micros() as u64;

                            phase_stats.record_successful_prediction(latency_us);

                            let result = PredictionResult {
                                model_name: model.name().to_owned(),
                                model_type: model.model_type(),
                                prediction,
                                latency_us,
                                timestamp: std::time::SystemTime::now(),
                                success: true,
                                error_message: None,
                            };

                            drop(prediction_tx.send(result).await);
                        },
                        Err(e) => {
                            let latency_us = model_start.elapsed().as_micros() as u64;
                            phase_stats.record_failed_prediction(latency_us);

                            let result = PredictionResult {
                                model_name: model.name().to_owned(),
                                model_type: model.model_type(),
                                prediction: ModelPrediction::new(
                                    model.name().to_owned(),
                                    0.0,
                                    0.0,
                                ),
                                latency_us,
                                timestamp: std::time::SystemTime::now(),
                                success: false,
                                error_message: Some(e.to_string()),
                            };

                            drop(prediction_tx.send(result).await);
                        },
                    }
                }
            }

            // Small delay to prevent busy waiting
            tokio::time::sleep(Duration::from_micros(100)).await;
        }

        Ok(PhaseResult {
            phase_name: phase.name.clone(),
            duration: phase_start.elapsed(),
            stats: phase_stats,
        })
    }

    /// Convert market data to ML features
    #[allow(clippy::unnecessary_wraps)]
    fn convert_market_data_to_features(&self, market_data: &MarketDataUpdate) -> Result<Features> {
        let values = vec![
            market_data.price.to_f64(),
            market_data.volume.to_f64().unwrap_or(0.0),
            market_data.bid.to_f64(),
            market_data.ask.to_f64(),
            market_data.spread().to_f64(),
            market_data.mid_price().to_f64(),
        ];

        let names = vec![
            "price".to_owned(),
            "volume".to_owned(),
            "bid".to_owned(),
            "ask".to_owned(),
            "spread".to_owned(),
            "mid_price".to_owned(),
        ];

        Ok(Features::new(values, names).with_symbol(market_data.symbol.clone()))
    }

    /// Generate comprehensive stress test report
    async fn generate_stress_test_report(
        &self,
        phase_results: Vec<PhaseResult>,
        total_duration: Duration,
    ) -> Result<StressTestReport> {
        let mut total_predictions = 0;
        let mut total_errors = 0;
        let mut all_latencies = Vec::new();

        for phase in &phase_results {
            total_predictions +=
                phase.stats.successful_predictions + phase.stats.failed_predictions;
            total_errors += phase.stats.failed_predictions;
            all_latencies.extend(&phase.stats.latencies);
        }

        let error_rate = if total_predictions > 0 {
            total_errors as f64 / total_predictions as f64
        } else {
            0.0
        };

        let latency_stats = self.calculate_latency_statistics(&all_latencies);

        // Check if requirements are met
        let requirements_met = self.check_requirements(&latency_stats, error_rate);
        let recommendations = self.generate_recommendations(&latency_stats, error_rate);

        Ok(StressTestReport {
            config: self.config.clone(),
            total_duration,
            phase_results,
            total_predictions,
            total_errors,
            error_rate,
            latency_stats,
            requirements_met,
            throughput_achieved: total_predictions as f64 / total_duration.as_secs_f64(),
            recommendations,
        })
    }

    fn calculate_latency_statistics(&self, latencies: &[u64]) -> LatencyStats {
        if latencies.is_empty() {
            return LatencyStats::default();
        }

        let mut sorted_latencies = latencies.to_vec();
        sorted_latencies.sort_unstable();

        let len = sorted_latencies.len();
        let mean = sorted_latencies.iter().sum::<u64>() as f64 / len as f64;
        let min = sorted_latencies[0];
        let max = sorted_latencies[len - 1];
        let p50 = sorted_latencies[len * 50 / 100];
        let p95 = sorted_latencies[len * 95 / 100];
        let p99 = sorted_latencies[len * 99 / 100];

        LatencyStats {
            mean,
            min,
            max,
            p50,
            p95,
            p99,
            count: len as u64,
        }
    }

    fn check_requirements(
        &self,
        latency_stats: &LatencyStats,
        error_rate: f64,
    ) -> RequirementsCheck {
        RequirementsCheck {
            latency_ok: latency_stats.max <= self.config.requirements.max_latency_us,
            p95_latency_ok: latency_stats.p95 <= self.config.requirements.p95_latency_us,
            p99_latency_ok: latency_stats.p99 <= self.config.requirements.p99_latency_us,
            error_rate_ok: error_rate <= self.config.requirements.max_error_rate,
            overall_pass: latency_stats.max <= self.config.requirements.max_latency_us
                && latency_stats.p95 <= self.config.requirements.p95_latency_us
                && latency_stats.p99 <= self.config.requirements.p99_latency_us
                && error_rate <= self.config.requirements.max_error_rate,
        }
    }

    fn generate_recommendations(
        &self,
        latency_stats: &LatencyStats,
        error_rate: f64,
    ) -> Vec<String> {
        let mut recommendations = Vec::new();

        if latency_stats.p99 > self.config.requirements.p99_latency_us {
            recommendations.push(format!(
                "P99 latency ({}us) exceeds requirement ({}us). Consider model optimization or hardware upgrades.",
                latency_stats.p99, self.config.requirements.p99_latency_us
            ));
        }

        if error_rate > self.config.requirements.max_error_rate {
            recommendations.push(format!(
                "Error rate ({:.2}%) exceeds requirement ({:.2}%). Investigate model reliability.",
                error_rate * 100.0,
                self.config.requirements.max_error_rate * 100.0
            ));
        }

        if latency_stats.mean > 50.0 {
            recommendations
                .push("Consider enabling GPU acceleration for better performance.".to_owned());
        }

        if recommendations.is_empty() {
            recommendations
                .push("All performance requirements met. System ready for production.".to_owned());
        }

        recommendations
    }
}

/// Market data update structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MarketDataUpdate {
    pub symbol: String,
    pub price: Price,
    pub volume: Decimal,
    pub bid: Price,
    pub ask: Price,
    pub timestamp: std::time::SystemTime,
}

impl MarketDataUpdate {
    pub fn spread(&self) -> Price {
        Price::from_f64(self.ask.to_f64() - self.bid.to_f64()).unwrap_or_default()
    }

    pub fn mid_price(&self) -> Price {
        Price::from_f64((self.bid.to_f64() + self.ask.to_f64()) / 2.0).unwrap_or_default()
    }
}

/// Prediction result with timing information
#[derive(Debug, Clone)]
pub struct PredictionResult {
    pub model_name: String,
    pub model_type: ModelType,
    pub prediction: ModelPrediction,
    pub latency_us: u64,
    pub timestamp: std::time::SystemTime,
    pub success: bool,
    pub error_message: Option<String>,
}

/// Phase execution statistics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PhaseStats {
    pub successful_predictions: u64,
    pub failed_predictions: u64,
    pub latencies: Vec<u64>,
}

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

impl PhaseStats {
    pub fn new() -> Self {
        Self {
            successful_predictions: 0,
            failed_predictions: 0,
            latencies: Vec::new(),
        }
    }

    pub fn record_successful_prediction(&mut self, latency_us: u64) {
        self.successful_predictions += 1;
        self.latencies.push(latency_us);
    }

    pub fn record_failed_prediction(&mut self, latency_us: u64) {
        self.failed_predictions += 1;
        self.latencies.push(latency_us);
    }
}

/// Phase execution result
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct PhaseResult {
    pub phase_name: String,
    pub duration: Duration,
    pub stats: PhaseStats,
}

/// Requirements check result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RequirementsCheck {
    pub latency_ok: bool,
    pub p95_latency_ok: bool,
    pub p99_latency_ok: bool,
    pub error_rate_ok: bool,
    pub overall_pass: bool,
}

/// Create default HFT stress test configuration with production-ready settings
pub fn create_hft_stress_test_config() -> StressTestConfig {
    StressTestConfig {
        duration_seconds: 300, // 5 minutes
        target_rps: 50000,     // 50k requests per second
        concurrent_connections: 100,
        market_data_rate: 10000, // 10k market updates per second
        test_phases: vec![
            TestPhase {
                name: "warmup".to_owned(),
                duration_seconds: 60,
                load_multiplier: 0.5,
                market_volatility: 0.01,
                error_injection_rate: 0.0,
            },
            TestPhase {
                name: "normal_load".to_owned(),
                duration_seconds: 120,
                load_multiplier: 1.0,
                market_volatility: 0.02,
                error_injection_rate: 0.001,
            },
            TestPhase {
                name: "peak_load".to_owned(),
                duration_seconds: 60,
                load_multiplier: 2.0,
                market_volatility: 0.05,
                error_injection_rate: 0.005,
            },
            TestPhase {
                name: "stress_load".to_owned(),
                duration_seconds: 60,
                load_multiplier: 5.0,
                market_volatility: 0.1,
                error_injection_rate: 0.01,
            },
        ],
        models_to_test: vec![
            "TLOB_Transformer".to_owned(),
            "MAMBA_SSM".to_owned(),
            "DQN_Agent".to_owned(),
        ],
        market_conditions: vec![
            MarketCondition::Normal,
            MarketCondition::HighVolatility,
            MarketCondition::Flash,
        ],
        requirements: PerformanceRequirements::default(),
    }
}

/// Create stress test configuration optimized for testing with generic symbols
pub fn create_test_stress_test_config() -> StressTestConfig {
    let mut config = create_hft_stress_test_config();

    // Reduce intensity for testing
    config.duration_seconds = 60; // 1 minute for testing
    config.target_rps = 1000; // 1k requests per second for testing
    config.market_data_rate = 100; // 100 updates per second for testing

    // Simplified test phases
    config.test_phases = vec![TestPhase {
        name: "test_phase".to_owned(),
        duration_seconds: 60,
        load_multiplier: 1.0,
        market_volatility: 0.02,
        error_injection_rate: 0.0,
    }];

    config
}

/// Create stress test configuration with custom simulation parameters
pub fn create_custom_stress_test_config(
    _symbols: Vec<String>,
    simulation_config: config::SimulationConfig,
) -> StressTestConfig {
    let mut config = create_hft_stress_test_config();

    // Use simulation config parameters
    config.market_data_rate = simulation_config.parameters.update_rate_hz;

    // Update volatility in test phases based on simulation config
    for phase in &mut config.test_phases {
        phase.market_volatility = simulation_config.parameters.base_volatility;
    }

    config
}

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

    #[test]
    fn test_stress_test_config_creation() {
        let config = create_hft_stress_test_config();
        assert_eq!(config.test_phases.len(), 4);
        assert_eq!(config.target_rps, 50000);
    }

    #[test]
    fn test_market_data_calculations() {
        let update = MarketDataUpdate {
            symbol: "TEST_LARGE_1".to_owned(),
            price: Price::from_f64(150.0).unwrap(),
            volume: Decimal::try_from(1000.0).unwrap(),
            bid: Price::from_f64(149.95).unwrap(),
            ask: Price::from_f64(150.05).unwrap(),
            timestamp: std::time::SystemTime::now(),
        };

        assert!((update.spread().to_f64() - 0.10).abs() < 0.001);
        assert!((update.mid_price().to_f64() - 150.0).abs() < 0.001);
    }

    #[test]
    fn test_configuration_driven_simulator() {
        let simulation_config = config::SimulationConfig::default();
        let test_symbols = MarketDataSimulator::generate_test_symbols(&simulation_config);

        assert_eq!(test_symbols.len(), simulation_config.test_symbols.count);
        assert!(test_symbols[0].starts_with(&simulation_config.test_symbols.symbol_prefix));
    }

    #[test]
    fn test_custom_stress_test_config() {
        let symbols = vec!["TEST001".to_owned(), "TEST002".to_owned()];
        let simulation_config = config::SimulationConfig::default();
        let stress_config = create_custom_stress_test_config(symbols, simulation_config.clone());

        assert_eq!(
            stress_config.market_data_rate,
            simulation_config.parameters.update_rate_hz
        );
    }

    #[test]
    fn test_phase_stats() {
        let mut stats = PhaseStats::new();
        stats.record_successful_prediction(25);
        stats.record_successful_prediction(30);
        stats.record_failed_prediction(100);

        assert_eq!(stats.successful_predictions, 2);
        assert_eq!(stats.failed_predictions, 1);
        assert_eq!(stats.latencies.len(), 3);
    }
}