foxhunt/services/ml_training_service/tests/ensemble_training_tests.rs

//! TDD Tests for Ensemble Training Coordination
//!
//! These tests define the behavior we expect from the ensemble training system
//! BEFORE implementing the actual functionality.
//!
//! Test Coverage:
//! 1. Ensemble training configuration
//! 2. Multi-model coordination during training
//! 3. Ensemble weight optimization
//! 4. Checkpoint synchronization (all 4 models)
//! 5. Integration with ML Training Service

use std::collections::HashMap;
use std::sync::Arc;

use chrono::Utc;
use ml::training_pipeline::{ProductionTrainingConfig, ModelArchitectureConfig, TrainingHyperparameters, FinancialValidationConfig, PerformanceConfig};
use ml::safety::{MLSafetyConfig, GradientSafetyConfig};
use ml_training_service::ensemble_training_coordinator::{
    EnsembleTrainingCoordinator, EnsembleTrainingConfig, ModelTrainingStatus
};
use uuid::Uuid;

/// Test 1: Ensemble training configuration validation
#[tokio::test]
async fn test_ensemble_training_config_validation() {
    // Test 1.1: Valid configuration should be accepted
    let config = create_valid_ensemble_config();
    assert!(config.is_valid(), "Valid ensemble config should pass validation");

    // Test 1.2: Must have all 4 models (DQN, PPO, MAMBA-2, TFT)
    let mut models = config.model_configs.keys().cloned().collect::<Vec<_>>();
    models.sort();
    assert_eq!(models, vec!["DQN", "MAMBA2", "PPO", "TFT"], "Must configure all 4 models");

    // Test 1.3: Weights must sum to 1.0
    let weight_sum: f64 = config.model_weights.values().sum();
    assert!((weight_sum - 1.0).abs() < 1e-6, "Model weights must sum to 1.0, got {}", weight_sum);

    // Test 1.4: Each model must have matching training and weight configuration
    for model_name in &["DQN", "PPO", "MAMBA2", "TFT"] {
        assert!(config.model_configs.contains_key(model_name), "Missing config for {}", model_name);
        assert!(config.model_weights.contains_key(model_name), "Missing weight for {}", model_name);
    }
}

/// Test 2: Multi-model coordination during training
#[tokio::test]
async fn test_multi_model_training_coordination() {
    let config = create_valid_ensemble_config();
    let coordinator = create_ensemble_coordinator(config).await;

    // Test 2.1: All models should start in Pending state
    for model_name in &["DQN", "PPO", "MAMBA2", "TFT"] {
        let status = coordinator.get_model_status(model_name).await.unwrap();
        assert_eq!(status, ModelTrainingStatus::Pending, "{} should start Pending", model_name);
    }

    // Test 2.2: Can start training for all models
    let job_id = coordinator.start_ensemble_training().await.unwrap();
    assert_ne!(job_id, Uuid::nil(), "Should return valid job ID");

    // Test 2.3: At least one model should be training after start
    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
    let any_training = ["DQN", "PPO", "MAMBA2", "TFT"].iter().any(|model| {
        matches!(
            coordinator.get_model_status(model).await,
            Ok(ModelTrainingStatus::Training)
        )
    });
    assert!(any_training, "At least one model should be training after start");
}

/// Test 3: Ensemble weight optimization
#[tokio::test]
async fn test_ensemble_weight_optimization() {
    let mut config = create_valid_ensemble_config();
    config.enable_weight_optimization = true;
    config.weight_optimization_interval_epochs = 5;

    let coordinator = create_ensemble_coordinator(config).await;

    // Test 3.1: Initial weights should match configuration
    let initial_weights = coordinator.get_current_weights().await.unwrap();
    assert_eq!(initial_weights.len(), 4, "Should have 4 model weights");

    // Test 3.2: Simulate training progress and weight updates
    coordinator.simulate_training_epochs(10).await.unwrap();

    // Test 3.3: Weights should be updated after optimization interval
    let updated_weights = coordinator.get_current_weights().await.unwrap();
    assert_ne!(initial_weights, updated_weights, "Weights should be updated after optimization");

    // Test 3.4: Updated weights should still sum to 1.0
    let weight_sum: f64 = updated_weights.values().sum();
    assert!((weight_sum - 1.0).abs() < 1e-6, "Optimized weights must sum to 1.0, got {}", weight_sum);

    // Test 3.5: Better-performing models should get higher weights
    // (This test assumes DQN performs better in simulation)
    let dqn_initial = initial_weights.get("DQN").unwrap();
    let dqn_updated = updated_weights.get("DQN").unwrap();
    // Weight adjustment logic will determine if this increases or decreases
    assert_ne!(dqn_initial, dqn_updated, "DQN weight should be adjusted based on performance");
}

/// Test 4: Checkpoint synchronization for all models
#[tokio::test]
async fn test_checkpoint_synchronization() {
    let config = create_valid_ensemble_config();
    let coordinator = create_ensemble_coordinator(config).await;

    // Test 4.1: Start training and wait for first checkpoint
    let job_id = coordinator.start_ensemble_training().await.unwrap();
    coordinator.simulate_training_epochs(1).await.unwrap();

    // Test 4.2: All models should have checkpoint paths after first epoch
    for model_name in &["DQN", "PPO", "MAMBA2", "TFT"] {
        let checkpoint = coordinator.get_latest_checkpoint(model_name).await.unwrap();
        assert!(checkpoint.is_some(), "{} should have checkpoint after epoch 1", model_name);

        let checkpoint_path = checkpoint.unwrap();
        assert!(checkpoint_path.contains(model_name), "Checkpoint path should contain model name");
        assert!(checkpoint_path.contains("epoch_1"), "Checkpoint should be from epoch 1");
    }

    // Test 4.3: Checkpoints should be synchronized (all from same epoch)
    let checkpoints = coordinator.get_all_checkpoints().await.unwrap();
    let epochs: Vec<_> = checkpoints.iter().map(|(_, cp)| {
        cp.split("epoch_").last().unwrap().split('_').next().unwrap().parse::<u32>().unwrap()
    }).collect();

    let first_epoch = epochs[0];
    assert!(epochs.iter().all(|&e| e == first_epoch), "All checkpoints should be from same epoch");

    // Test 4.4: Can load synchronized ensemble from checkpoints
    let ensemble_restored = coordinator.load_synchronized_ensemble(first_epoch).await;
    assert!(ensemble_restored.is_ok(), "Should be able to load synchronized ensemble");
}

/// Test 5: Performance-based weight adjustment
#[tokio::test]
async fn test_performance_based_weight_adjustment() {
    let mut config = create_valid_ensemble_config();
    config.enable_weight_optimization = true;

    let coordinator = create_ensemble_coordinator(config).await;

    // Test 5.1: Set different performance metrics for each model
    coordinator.set_model_performance("DQN", 0.85, 0.15).await.unwrap(); // High accuracy, low loss
    coordinator.set_model_performance("PPO", 0.75, 0.25).await.unwrap(); // Medium
    coordinator.set_model_performance("MAMBA2", 0.65, 0.35).await.unwrap(); // Lower
    coordinator.set_model_performance("TFT", 0.90, 0.10).await.unwrap(); // Highest

    // Test 5.2: Trigger weight optimization
    coordinator.optimize_weights().await.unwrap();

    // Test 5.3: TFT should have highest weight (best performance)
    let weights = coordinator.get_current_weights().await.unwrap();
    let tft_weight = weights.get("TFT").unwrap();

    for (model, weight) in weights.iter() {
        if model != "TFT" {
            assert!(tft_weight >= weight, "TFT (best performer) should have highest or equal weight");
        }
    }

    // Test 5.4: MAMBA2 should have lowest weight (worst performance)
    let mamba2_weight = weights.get("MAMBA2").unwrap();
    for (model, weight) in weights.iter() {
        if model != "MAMBA2" {
            assert!(mamba2_weight <= weight, "MAMBA2 (worst performer) should have lowest or equal weight");
        }
    }
}

/// Test 6: Training failure recovery
#[tokio::test]
async fn test_training_failure_recovery() {
    let config = create_valid_ensemble_config();
    let coordinator = create_ensemble_coordinator(config).await;

    // Test 6.1: Start training
    let job_id = coordinator.start_ensemble_training().await.unwrap();

    // Test 6.2: Simulate one model failing
    coordinator.simulate_model_failure("PPO").await.unwrap();

    // Test 6.3: PPO should be in Failed state
    let ppo_status = coordinator.get_model_status("PPO").await.unwrap();
    assert_eq!(ppo_status, ModelTrainingStatus::Failed, "PPO should be in Failed state");

    // Test 6.4: Other models should continue training
    for model_name in &["DQN", "MAMBA2", "TFT"] {
        let status = coordinator.get_model_status(model_name).await.unwrap();
        assert_ne!(status, ModelTrainingStatus::Failed, "{} should not fail due to PPO failure", model_name);
    }

    // Test 6.5: Can retry failed model
    let retry_result = coordinator.retry_failed_model("PPO").await;
    assert!(retry_result.is_ok(), "Should be able to retry failed model");

    // Test 6.6: PPO should return to training after retry
    tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
    let ppo_status_after_retry = coordinator.get_model_status("PPO").await.unwrap();
    assert_ne!(ppo_status_after_retry, ModelTrainingStatus::Failed, "PPO should not be Failed after retry");
}

/// Test 7: Ensemble validation metrics
#[tokio::test]
async fn test_ensemble_validation_metrics() {
    let config = create_valid_ensemble_config();
    let coordinator = create_ensemble_coordinator(config).await;

    // Test 7.1: Start training
    coordinator.start_ensemble_training().await.unwrap();
    coordinator.simulate_training_epochs(5).await.unwrap();

    // Test 7.2: Should have ensemble-level metrics
    let metrics = coordinator.get_ensemble_metrics().await.unwrap();
    assert!(metrics.contains_key("ensemble_train_loss"), "Should have ensemble train loss");
    assert!(metrics.contains_key("ensemble_val_loss"), "Should have ensemble val loss");
    assert!(metrics.contains_key("ensemble_accuracy"), "Should have ensemble accuracy");

    // Test 7.3: Ensemble metrics should be aggregated from all models
    let ensemble_loss = metrics.get("ensemble_train_loss").unwrap();
    assert!(ensemble_loss > &0.0, "Ensemble loss should be positive");

    // Test 7.4: Should track diversity metrics
    assert!(metrics.contains_key("prediction_diversity"), "Should track prediction diversity");
    let diversity = metrics.get("prediction_diversity").unwrap();
    assert!(diversity >= &0.0 && diversity <= &1.0, "Diversity should be in [0, 1]");
}

/// Test 8: Integration with ML Training Service
#[tokio::test]
async fn test_integration_with_ml_training_service() {
    // This test verifies the coordinator integrates with existing ML training infrastructure

    let config = create_valid_ensemble_config();
    let coordinator = create_ensemble_coordinator(config).await;

    // Test 8.1: Should use existing ProductionTrainingConfig
    for model_name in &["DQN", "PPO", "MAMBA2", "TFT"] {
        let model_config = coordinator.get_model_training_config(model_name).await.unwrap();
        assert!(model_config.model_config.input_dim > 0, "Should have valid input dimension");
        assert!(!model_config.model_config.hidden_dims.is_empty(), "Should have hidden layers");
    }

    // Test 8.2: Should respect existing safety configurations
    let dqn_config = coordinator.get_model_training_config("DQN").await.unwrap();
    assert!(dqn_config.safety_config.max_loss_value > 0.0, "Should have safety limits");
    assert!(dqn_config.gradient_config.max_gradient_norm > 0.0, "Should have gradient clipping");

    // Test 8.3: Should integrate with checkpoint manager
    coordinator.start_ensemble_training().await.unwrap();
    coordinator.simulate_training_epochs(1).await.unwrap();

    let checkpoints = coordinator.get_all_checkpoints().await.unwrap();
    assert_eq!(checkpoints.len(), 4, "Should have checkpoints for all 4 models");
}

// Helper functions for tests

/// Create valid ensemble configuration
fn create_valid_ensemble_config() -> EnsembleTrainingConfig {
    let mut model_configs = HashMap::new();
    let mut model_weights = HashMap::new();

    // DQN configuration (33% weight)
    model_configs.insert("DQN".to_string(), create_model_config("DQN", 64, vec![256, 128], 32));
    model_weights.insert("DQN".to_string(), 0.33);

    // PPO configuration (33% weight)
    model_configs.insert("PPO".to_string(), create_model_config("PPO", 64, vec![256, 128], 32));
    model_weights.insert("PPO".to_string(), 0.33);

    // MAMBA-2 configuration (17% weight)
    model_configs.insert("MAMBA2".to_string(), create_model_config("MAMBA2", 64, vec![512, 256], 32));
    model_weights.insert("MAMBA2".to_string(), 0.17);

    // TFT configuration (17% weight)
    model_configs.insert("TFT".to_string(), create_model_config("TFT", 64, vec![512, 256, 128], 32));
    model_weights.insert("TFT".to_string(), 0.17);

    EnsembleTrainingConfig {
        job_id: Uuid::new_v4(),
        model_configs,
        model_weights,
        enable_weight_optimization: true,
        weight_optimization_interval_epochs: 10,
        checkpoint_interval_epochs: 1,
        max_epochs: 100,
        parallel_training: true,
        created_at: Utc::now(),
    }
}

/// Create model-specific training configuration
fn create_model_config(model_type: &str, input_dim: usize, hidden_dims: Vec<usize>, output_dim: usize) -> ProductionTrainingConfig {
    ProductionTrainingConfig {
        model_config: ModelArchitectureConfig {
            input_dim,
            hidden_dims,
            output_dim,
            dropout_rate: 0.1,
            activation: "relu".to_string(),
            batch_norm: true,
            residual_connections: false,
        },
        training_params: TrainingHyperparameters {
            learning_rate: 0.001,
            batch_size: 64,
            max_epochs: 100,
            patience: 10,
            validation_split: 0.2,
            l2_regularization: 0.0001,
            lr_decay_factor: 0.5,
            lr_decay_patience: 5,
        },
        safety_config: MLSafetyConfig {
            max_loss_value: 1000.0,
            max_prediction_value: 100.0,
            nan_check_interval: 10,
            enable_loss_scaling: true,
            convergence_window: 20,
        },
        gradient_config: GradientSafetyConfig {
            max_gradient_norm: 1.0,
            min_gradient_norm: 1e-8,
            gradient_clip_threshold: 5.0,
            enable_gradient_monitoring: true,
            gradient_check_interval: 1,
        },
        financial_config: FinancialValidationConfig {
            max_prediction_multiple: 2.0,
            min_prediction_confidence: 0.6,
            validate_position_sizing: true,
            max_position_fraction: 0.2,
            min_sharpe_threshold: 0.5,
        },
        performance_config: PerformanceConfig {
            device_preference: "cpu".to_string(),
            max_memory_bytes: 4_000_000_000,
            mixed_precision: false,
            num_workers: 2,
            gradient_accumulation_steps: 1,
        },
    }
}

/// Create ensemble coordinator instance
async fn create_ensemble_coordinator(config: EnsembleTrainingConfig) -> EnsembleTrainingCoordinator {
    EnsembleTrainingCoordinator::new(config).await.expect("Failed to create coordinator")
}