foxhunt/tests/e2e/tests/ml_model_integration_tests.rs

//! ML Model Integration E2E Tests
//!
//! Comprehensive end-to-end testing of ML model integration with the Foxhunt trading system.
//! Tests cover model inference, feature extraction, ensemble predictions, and trading integration.

use anyhow::Result;
use foxhunt_e2e::*;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tracing::{debug, info};

/// ML Model Integration Test Suite
pub struct MLModelIntegrationTests {
    framework: Arc<E2ETestFramework>,
}

impl MLModelIntegrationTests {
    pub fn new(framework: Arc<E2ETestFramework>) -> Self {
        Self { framework }
    }

    /// Test 1: Basic ML Model Health Check
    ///
    /// Validates that all ML models are available and healthy
    pub async fn test_ml_model_health(&self) -> Result<WorkflowTestResult> {
        let start_time = Instant::now();
        let workflow_name = "ml_model_health_check".to_string();

        info!("🔍 Starting ML Model Health Check Test");

        let mut metrics = HashMap::new();
        let mut steps_completed = 0;

        // Create ML pipeline instance for health check
        let ml_pipeline = ml_pipeline::MLPipelineTestHarness::new().await?;
        let model_status = ml_pipeline.check_models_health().await?;

        metrics.insert(
            "mamba_available".to_string(),
            if model_status.mamba_available {
                1.0
            } else {
                0.0
            },
        );
        metrics.insert(
            "dqn_available".to_string(),
            if model_status.dqn_available { 1.0 } else { 0.0 },
        );
        metrics.insert(
            "tft_available".to_string(),
            if model_status.tft_available { 1.0 } else { 0.0 },
        );
        metrics.insert(
            "tlob_available".to_string(),
            if model_status.tlob_available {
                1.0
            } else {
                0.0
            },
        );
        metrics.insert(
            "models_available".to_string(),
            model_status.available_count() as f64,
        );

        assert!(
            model_status.any_available(),
            "At least one ML model should be available"
        );

        info!(
            "✅ ML Model Health Check: {} models available",
            model_status.available_count()
        );
        steps_completed += 1;

        let duration = start_time.elapsed();
        let mut result = WorkflowTestResult::success(workflow_name, duration, steps_completed);
        result.metrics = metrics;

        Ok(result)
    }

    /// Test 2: Feature Extraction Pipeline
    ///
    /// Tests extraction of features from market data for ML model input
    pub async fn test_feature_extraction(&self) -> Result<WorkflowTestResult> {
        let start_time = Instant::now();
        let workflow_name = "feature_extraction_pipeline".to_string();

        info!("🔧 Starting Feature Extraction Pipeline Test");

        let mut metrics = HashMap::new();
        let mut steps_completed = 0;

        // Generate test market data
        let test_data = test_utils::generate_market_data("AAPL", 100);
        metrics.insert("market_data_count".to_string(), test_data.len() as f64);
        steps_completed += 1;

        // Extract features using a new ML pipeline instance
        let feature_start = Instant::now();
        let mut ml_pipeline = ml_pipeline::MLPipelineTestHarness::new().await?;
        let features = ml_pipeline.extract_features(&test_data).await?;
        let feature_duration = feature_start.elapsed();

        metrics.insert("features_extracted".to_string(), features.len() as f64);
        metrics.insert(
            "feature_extraction_ms".to_string(),
            feature_duration.as_millis() as f64,
        );

        assert!(
            !features.is_empty(),
            "Should extract at least one feature vector"
        );
        assert!(
            feature_duration < Duration::from_secs(1),
            "Feature extraction should complete in < 1s"
        );

        if let Some(feature_vec) = features.first() {
            metrics.insert(
                "feature_dimension".to_string(),
                feature_vec.features.len() as f64,
            );
            info!(
                "✅ Extracted {} feature vectors with {} dimensions",
                features.len(),
                feature_vec.features.len()
            );
        }

        steps_completed += 1;

        let duration = start_time.elapsed();
        let mut result = WorkflowTestResult::success(workflow_name, duration, steps_completed);
        result.metrics = metrics;

        Ok(result)
    }

    /// Test 3: MAMBA Model Inference
    ///
    /// Tests MAMBA-2 state space model predictions
    pub async fn test_mamba_inference(&self) -> Result<WorkflowTestResult> {
        let start_time = Instant::now();
        let workflow_name = "mamba_model_inference".to_string();

        info!("🧬 Starting MAMBA Model Inference Test");

        let mut metrics = HashMap::new();
        let mut steps_completed = 0;

        // Generate test data and extract features
        let test_data = test_utils::generate_market_data("AAPL", 50);
        let mut ml_pipeline = ml_pipeline::MLPipelineTestHarness::new().await?;
        let features = ml_pipeline.extract_features(&test_data).await?;
        steps_completed += 1;

        // Test MAMBA inference
        let inference_start = Instant::now();
        let prediction = ml_pipeline.predict_with_mamba(&features).await?;
        let inference_duration = inference_start.elapsed();

        metrics.insert("mamba_signal".to_string(), prediction.signal);
        metrics.insert("mamba_confidence".to_string(), prediction.confidence);
        metrics.insert(
            "mamba_inference_ms".to_string(),
            inference_duration.as_millis() as f64,
        );

        assert!(
            prediction.confidence >= 0.0 && prediction.confidence <= 1.0,
            "Confidence should be in [0, 1]"
        );
        assert!(
            prediction.signal >= -1.0 && prediction.signal <= 1.0,
            "Signal should be in [-1, 1]"
        );
        assert!(
            inference_duration < Duration::from_millis(500),
            "MAMBA inference should complete in < 500ms"
        );

        info!(
            "✅ MAMBA prediction: signal={:.4}, confidence={:.4}, latency={:?}",
            prediction.signal, prediction.confidence, inference_duration
        );
        steps_completed += 1;

        let duration = start_time.elapsed();
        let mut result = WorkflowTestResult::success(workflow_name, duration, steps_completed);
        result.metrics = metrics;

        Ok(result)
    }

    /// Test 4: DQN Reinforcement Learning
    ///
    /// Tests Deep Q-Network model for trading decisions
    pub async fn test_dqn_inference(&self) -> Result<WorkflowTestResult> {
        let start_time = Instant::now();
        let workflow_name = "dqn_reinforcement_learning".to_string();

        info!("🎮 Starting DQN Reinforcement Learning Test");

        let mut metrics = HashMap::new();
        let mut steps_completed = 0;

        // Generate test data and extract features
        let test_data = test_utils::generate_market_data("MSFT", 50);
        let mut ml_pipeline = ml_pipeline::MLPipelineTestHarness::new().await?;
        let features = ml_pipeline.extract_features(&test_data).await?;
        steps_completed += 1;

        // Test DQN inference
        let inference_start = Instant::now();
        let prediction = ml_pipeline.predict_with_dqn(&features).await?;
        let inference_duration = inference_start.elapsed();

        metrics.insert("dqn_signal".to_string(), prediction.signal);
        metrics.insert("dqn_confidence".to_string(), prediction.confidence);
        metrics.insert(
            "dqn_inference_ms".to_string(),
            inference_duration.as_millis() as f64,
        );

        assert!(
            prediction.confidence >= 0.0 && prediction.confidence <= 1.0,
            "Confidence should be in [0, 1]"
        );
        assert!(
            inference_duration < Duration::from_millis(500),
            "DQN inference should complete in < 500ms"
        );

        info!(
            "✅ DQN prediction: signal={:.4}, confidence={:.4}, latency={:?}",
            prediction.signal, prediction.confidence, inference_duration
        );
        steps_completed += 1;

        let duration = start_time.elapsed();
        let mut result = WorkflowTestResult::success(workflow_name, duration, steps_completed);
        result.metrics = metrics;

        Ok(result)
    }

    /// Test 5: TFT Time Series Forecasting
    ///
    /// Tests Temporal Fusion Transformer for price prediction
    pub async fn test_tft_inference(&self) -> Result<WorkflowTestResult> {
        let start_time = Instant::now();
        let workflow_name = "tft_time_series_forecasting".to_string();

        info!("📈 Starting TFT Time Series Forecasting Test");

        let mut metrics = HashMap::new();
        let mut steps_completed = 0;

        // Generate test data and extract features
        let test_data = test_utils::generate_market_data("GOOGL", 100);
        let mut ml_pipeline = ml_pipeline::MLPipelineTestHarness::new().await?;
        let features = ml_pipeline.extract_features(&test_data).await?;
        steps_completed += 1;

        // Test TFT inference
        let inference_start = Instant::now();
        let prediction = ml_pipeline.predict_with_tft(&features).await?;
        let inference_duration = inference_start.elapsed();

        metrics.insert("tft_signal".to_string(), prediction.signal);
        metrics.insert("tft_confidence".to_string(), prediction.confidence);
        metrics.insert(
            "tft_inference_ms".to_string(),
            inference_duration.as_millis() as f64,
        );

        assert!(
            prediction.confidence >= 0.0 && prediction.confidence <= 1.0,
            "Confidence should be in [0, 1]"
        );
        assert!(
            inference_duration < Duration::from_secs(1),
            "TFT inference should complete in < 1s"
        );

        info!(
            "✅ TFT prediction: signal={:.4}, confidence={:.4}, latency={:?}",
            prediction.signal, prediction.confidence, inference_duration
        );
        steps_completed += 1;

        let duration = start_time.elapsed();
        let mut result = WorkflowTestResult::success(workflow_name, duration, steps_completed);
        result.metrics = metrics;

        Ok(result)
    }

    /// Test 6: TLOB Order Book Analysis
    ///
    /// Tests TLOB transformer for order book microstructure
    pub async fn test_tlob_inference(&self) -> Result<WorkflowTestResult> {
        let start_time = Instant::now();
        let workflow_name = "tlob_order_book_analysis".to_string();

        info!("📊 Starting TLOB Order Book Analysis Test");

        let mut metrics = HashMap::new();
        let mut steps_completed = 0;

        // Generate test data and extract features
        let test_data = test_utils::generate_market_data("TSLA", 50);
        let mut ml_pipeline = ml_pipeline::MLPipelineTestHarness::new().await?;
        let features = ml_pipeline.extract_features(&test_data).await?;
        steps_completed += 1;

        // Test TLOB inference
        let inference_start = Instant::now();
        let prediction = ml_pipeline.predict_with_tlob(&features).await?;
        let inference_duration = inference_start.elapsed();

        metrics.insert("tlob_signal".to_string(), prediction.signal);
        metrics.insert("tlob_confidence".to_string(), prediction.confidence);
        metrics.insert(
            "tlob_inference_ms".to_string(),
            inference_duration.as_millis() as f64,
        );

        assert!(
            prediction.confidence >= 0.0 && prediction.confidence <= 1.0,
            "Confidence should be in [0, 1]"
        );
        assert!(
            inference_duration < Duration::from_millis(200),
            "TLOB inference should complete in < 200ms (fast for order book analysis)"
        );

        info!(
            "✅ TLOB prediction: signal={:.4}, confidence={:.4}, latency={:?}",
            prediction.signal, prediction.confidence, inference_duration
        );
        steps_completed += 1;

        let duration = start_time.elapsed();
        let mut result = WorkflowTestResult::success(workflow_name, duration, steps_completed);
        result.metrics = metrics;

        Ok(result)
    }

    /// Test 7: Ensemble Model Prediction
    ///
    /// Tests aggregation of predictions from multiple models
    pub async fn test_ensemble_prediction(&self) -> Result<WorkflowTestResult> {
        let start_time = Instant::now();
        let workflow_name = "ensemble_model_prediction".to_string();

        info!("🤝 Starting Ensemble Model Prediction Test");

        let mut metrics = HashMap::new();
        let mut steps_completed = 0;

        // Generate test data and extract features
        let test_data = test_utils::generate_market_data("AAPL", 50);
        let mut ml_pipeline = ml_pipeline::MLPipelineTestHarness::new().await?;
        let features = ml_pipeline.extract_features(&test_data).await?;
        steps_completed += 1;

        // Test ensemble prediction
        let ensemble_start = Instant::now();
        let ensemble_prediction = ml_pipeline.predict_ensemble(&features).await?;
        let ensemble_duration = ensemble_start.elapsed();

        metrics.insert("ensemble_signal".to_string(), ensemble_prediction.signal);
        metrics.insert(
            "ensemble_confidence".to_string(),
            ensemble_prediction.confidence,
        );
        metrics.insert(
            "ensemble_models_count".to_string(),
            ensemble_prediction.individual_predictions.len() as f64,
        );
        metrics.insert(
            "ensemble_inference_ms".to_string(),
            ensemble_duration.as_millis() as f64,
        );
        metrics.insert(
            "ensemble_signal_strength".to_string(),
            ensemble_prediction.signal_strength,
        );

        assert!(
            !ensemble_prediction.individual_predictions.is_empty(),
            "Ensemble should aggregate at least one model"
        );
        assert!(
            ensemble_prediction.confidence >= 0.0 && ensemble_prediction.confidence <= 1.0,
            "Ensemble confidence should be in [0, 1]"
        );

        info!(
            "✅ Ensemble prediction: signal={:.4}, confidence={:.4}, models={}, prediction={}, latency={:?}",
            ensemble_prediction.signal,
            ensemble_prediction.confidence,
            ensemble_prediction.individual_predictions.len(),
            ensemble_prediction.prediction,
            ensemble_duration
        );

        // Log individual model contributions
        for (i, pred) in ensemble_prediction
            .individual_predictions
            .into_iter()
            .enumerate()
        {
            debug!(
                "   Model {}: {} signal={:.4}, confidence={:.4}",
                i + 1,
                pred.model_name,
                pred.signal,
                pred.confidence
            );
        }

        steps_completed += 1;

        let duration = start_time.elapsed();
        let mut result = WorkflowTestResult::success(workflow_name, duration, steps_completed);
        result.metrics = metrics;

        Ok(result)
    }

    /// Test 8: Model Performance Benchmarking
    ///
    /// Benchmarks inference latency and throughput for all models
    pub async fn test_model_performance(&self) -> Result<WorkflowTestResult> {
        let start_time = Instant::now();
        let workflow_name = "model_performance_benchmark".to_string();

        info!("⚡ Starting Model Performance Benchmark");

        let mut metrics = HashMap::new();
        let mut steps_completed = 0;

        // Generate test data once
        let test_data = test_utils::generate_market_data("AAPL", 50);
        let mut ml_pipeline = ml_pipeline::MLPipelineTestHarness::new().await?;
        let features = ml_pipeline.extract_features(&test_data).await?;
        steps_completed += 1;

        // Benchmark each model with multiple runs
        let benchmark_runs = 10;

        // MAMBA benchmark
        let mut mamba_latencies = Vec::new();
        for _ in 0..benchmark_runs {
            let start = Instant::now();
            let _ = ml_pipeline.predict_with_mamba(&features).await?;
            mamba_latencies.push(start.elapsed().as_micros() as f64);
        }
        let mamba_avg = mamba_latencies.iter().sum::<f64>() / mamba_latencies.len() as f64;
        metrics.insert("mamba_avg_latency_us".to_string(), mamba_avg);
        steps_completed += 1;

        // DQN benchmark
        let mut dqn_latencies = Vec::new();
        for _ in 0..benchmark_runs {
            let start = Instant::now();
            let _ = ml_pipeline.predict_with_dqn(&features).await?;
            dqn_latencies.push(start.elapsed().as_micros() as f64);
        }
        let dqn_avg = dqn_latencies.iter().sum::<f64>() / dqn_latencies.len() as f64;
        metrics.insert("dqn_avg_latency_us".to_string(), dqn_avg);
        steps_completed += 1;

        // TFT benchmark
        let mut tft_latencies = Vec::new();
        for _ in 0..benchmark_runs {
            let start = Instant::now();
            let _ = ml_pipeline.predict_with_tft(&features).await?;
            tft_latencies.push(start.elapsed().as_micros() as f64);
        }
        let tft_avg = tft_latencies.iter().sum::<f64>() / tft_latencies.len() as f64;
        metrics.insert("tft_avg_latency_us".to_string(), tft_avg);
        steps_completed += 1;

        // TLOB benchmark
        let mut tlob_latencies = Vec::new();
        for _ in 0..benchmark_runs {
            let start = Instant::now();
            let _ = ml_pipeline.predict_with_tlob(&features).await?;
            tlob_latencies.push(start.elapsed().as_micros() as f64);
        }
        let tlob_avg = tlob_latencies.iter().sum::<f64>() / tlob_latencies.len() as f64;
        metrics.insert("tlob_avg_latency_us".to_string(), tlob_avg);
        steps_completed += 1;

        info!(
            "✅ Performance Benchmark Results (avg over {} runs):",
            benchmark_runs
        );
        info!("   MAMBA: {:.2} μs", mamba_avg);
        info!("   DQN:   {:.2} μs", dqn_avg);
        info!("   TFT:   {:.2} μs", tft_avg);
        info!("   TLOB:  {:.2} μs", tlob_avg);

        let duration = start_time.elapsed();
        let mut result = WorkflowTestResult::success(workflow_name, duration, steps_completed);
        result.metrics = metrics;

        Ok(result)
    }

    /// Test 9: Model Failover and Redundancy
    ///
    /// Tests ensemble prediction when individual models fail
    pub async fn test_model_failover(&self) -> Result<WorkflowTestResult> {
        let start_time = Instant::now();
        let workflow_name = "model_failover_redundancy".to_string();

        info!("🔄 Starting Model Failover Test");

        let mut metrics = HashMap::new();
        let mut steps_completed = 0;

        // Generate test data
        let test_data = test_utils::generate_market_data("AAPL", 50);
        let mut ml_pipeline = ml_pipeline::MLPipelineTestHarness::new().await?;
        let features = ml_pipeline.extract_features(&test_data).await?;
        steps_completed += 1;

        // Get baseline ensemble prediction with all models
        let baseline_prediction = ml_pipeline.predict_ensemble(&features).await?;
        let baseline_models = baseline_prediction.individual_predictions.len();
        metrics.insert("baseline_models".to_string(), baseline_models as f64);
        steps_completed += 1;

        // Disable MAMBA and test failover
        ml_pipeline.disable_model("mamba").await?;
        let failover_prediction = ml_pipeline.predict_ensemble(&features).await?;
        let failover_models = failover_prediction.individual_predictions.len();
        metrics.insert("failover_models".to_string(), failover_models as f64);

        assert!(
            failover_models == baseline_models - 1,
            "Should have one fewer model after disabling MAMBA"
        );
        assert!(
            !failover_prediction.individual_predictions.is_empty(),
            "Ensemble should still work with remaining models"
        );

        info!(
            "✅ Failover test: {} models → {} models after MAMBA disabled",
            baseline_models, failover_models
        );
        steps_completed += 1;

        // Re-enable MAMBA
        ml_pipeline.enable_model("mamba").await?;
        let recovered_prediction = ml_pipeline.predict_ensemble(&features).await?;
        let recovered_models = recovered_prediction.individual_predictions.len();
        metrics.insert("recovered_models".to_string(), recovered_models as f64);

        assert!(
            recovered_models == baseline_models,
            "Should restore to baseline after re-enabling MAMBA"
        );

        info!(
            "✅ Recovery test: {} models restored after MAMBA re-enabled",
            recovered_models
        );
        steps_completed += 1;

        let duration = start_time.elapsed();
        let mut result = WorkflowTestResult::success(workflow_name, duration, steps_completed);
        result.metrics = metrics;

        Ok(result)
    }
}

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

    e2e_test!(
        test_ml_model_health,
        |framework: Arc<E2ETestFramework>| async move {
            let ml_tests = MLModelIntegrationTests::new(framework);
            let result = ml_tests.test_ml_model_health().await?;
            assert!(
                result.success,
                "ML model health check failed: {:?}",
                result.error_message
            );
            Ok(())
        }
    );

    e2e_test!(test_feature_extraction, |framework: Arc<
        E2ETestFramework,
    >| async move {
        let ml_tests = MLModelIntegrationTests::new(framework);
        let result = ml_tests.test_feature_extraction().await?;
        assert!(
            result.success,
            "Feature extraction failed: {:?}",
            result.error_message
        );
        Ok(())
    });

    e2e_test!(
        test_mamba_inference,
        |framework: Arc<E2ETestFramework>| async move {
            let ml_tests = MLModelIntegrationTests::new(framework);
            let result = ml_tests.test_mamba_inference().await?;
            assert!(
                result.success,
                "MAMBA inference failed: {:?}",
                result.error_message
            );
            Ok(())
        }
    );

    e2e_test!(
        test_dqn_inference,
        |framework: Arc<E2ETestFramework>| async move {
            let ml_tests = MLModelIntegrationTests::new(framework);
            let result = ml_tests.test_dqn_inference().await?;
            assert!(
                result.success,
                "DQN inference failed: {:?}",
                result.error_message
            );
            Ok(())
        }
    );

    e2e_test!(
        test_tft_inference,
        |framework: Arc<E2ETestFramework>| async move {
            let ml_tests = MLModelIntegrationTests::new(framework);
            let result = ml_tests.test_tft_inference().await?;
            assert!(
                result.success,
                "TFT inference failed: {:?}",
                result.error_message
            );
            Ok(())
        }
    );

    e2e_test!(
        test_tlob_inference,
        |framework: Arc<E2ETestFramework>| async move {
            let ml_tests = MLModelIntegrationTests::new(framework);
            let result = ml_tests.test_tlob_inference().await?;
            assert!(
                result.success,
                "TLOB inference failed: {:?}",
                result.error_message
            );
            Ok(())
        }
    );

    e2e_test!(test_ensemble_prediction, |framework: Arc<
        E2ETestFramework,
    >| async move {
        let ml_tests = MLModelIntegrationTests::new(framework);
        let result = ml_tests.test_ensemble_prediction().await?;
        assert!(
            result.success,
            "Ensemble prediction failed: {:?}",
            result.error_message
        );
        Ok(())
    });

    e2e_test!(test_model_performance, |framework: Arc<
        E2ETestFramework,
    >| async move {
        let ml_tests = MLModelIntegrationTests::new(framework);
        let result = ml_tests.test_model_performance().await?;
        assert!(
            result.success,
            "Model performance test failed: {:?}",
            result.error_message
        );
        Ok(())
    });

    e2e_test!(
        test_model_failover,
        |framework: Arc<E2ETestFramework>| async move {
            let ml_tests = MLModelIntegrationTests::new(framework);
            let result = ml_tests.test_model_failover().await?;
            assert!(
                result.success,
                "Model failover test failed: {:?}",
                result.error_message
            );
            Ok(())
        }
    );
}