foxhunt/services/ml_training_service/tests/batch_tuning_tests.rs

//! Batch Tuning Tests - Complete TDD Implementation
//!
//! These tests validate the BatchTuningManager with mock TuningManager
//! to avoid spawning actual Optuna subprocesses.

use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use uuid::Uuid;
use chrono::Utc;
use async_trait::async_trait;
use anyhow::Result;

use ml_training_service::batch_tuning_manager::{
    BatchTuningManager, BatchJobStatus, ModelTuningResult, BatchTuningJob,
};
use ml_training_service::tuning_manager::{
    TuningManagerTrait, TuningJob, TuningJobStatus, TrialResult, TrialState,
};

// ============================================================================
// MOCK TUNING MANAGER FOR TESTING
// ============================================================================

/// Mock TuningManager for unit testing without subprocess overhead
struct MockTuningManager {
    jobs: Arc<RwLock<HashMap<Uuid, TuningJob>>>,
    auto_complete: bool,
    failure_models: Vec<String>,
}

impl MockTuningManager {
    fn new() -> Self {
        Self {
            jobs: Arc::new(RwLock::new(HashMap::new())),
            auto_complete: true,
            failure_models: Vec::new(),
        }
    }

    fn with_failures(failure_models: Vec<String>) -> Self {
        Self {
            jobs: Arc::new(RwLock::new(HashMap::new())),
            auto_complete: true,
            failure_models,
        }
    }
}

#[async_trait]
impl TuningManagerTrait for MockTuningManager {
    async fn start_tuning_job(
        &self,
        model_type: String,
        num_trials: u32,
        _config_path: String,
        description: String,
        tags: HashMap<String, String>,
    ) -> Result<Uuid> {
        let mut job = TuningJob::new(model_type.clone(), num_trials, description, tags);
        let job_id = job.id;

        // Simulate failure for specific models
        if self.failure_models.contains(&model_type) {
            job.status = TuningJobStatus::Failed;
            job.error_message = Some(format!("Mock failure for {}", model_type));
        } else if self.auto_complete {
            // Auto-complete job with mock results
            job.status = TuningJobStatus::Completed;
            job.current_trial = num_trials;
            
            // Generate mock best params
            let mut best_params = HashMap::new();
            best_params.insert("learning_rate".to_string(), 0.001);
            best_params.insert("batch_size".to_string(), 128.0);
            job.best_params = best_params;

            // Generate mock metrics
            let mut best_metrics = HashMap::new();
            best_metrics.insert("sharpe_ratio".to_string(), 1.5 + (model_type.len() as f32) * 0.1);
            best_metrics.insert("training_loss".to_string(), 0.05);
            job.best_metrics = best_metrics;

            // Add mock trial history
            for i in 1..=num_trials {
                let mut trial_params = HashMap::new();
                trial_params.insert("learning_rate".to_string(), 0.001 * (i as f32));
                trial_params.insert("batch_size".to_string(), 64.0 + (i as f32) * 2.0);

                let mut trial_metrics = HashMap::new();
                trial_metrics.insert("sharpe_ratio".to_string(), 1.0 + (i as f32) * 0.05);

                job.trial_history.push(TrialResult {
                    trial_number: i,
                    params: trial_params,
                    objective_value: 1.0 + (i as f32) * 0.05,
                    metrics: trial_metrics,
                    state: TrialState::Complete,
                    started_at: Utc::now(),
                    completed_at: Some(Utc::now()),
                });
            }
        } else {
            job.status = TuningJobStatus::Running;
        }

        let mut jobs = self.jobs.write().await;
        jobs.insert(job_id, job);

        Ok(job_id)
    }

    async fn get_tuning_job_status(&self, job_id: Uuid) -> Result<TuningJob> {
        let jobs = self.jobs.read().await;
        jobs.get(&job_id)
            .cloned()
            .ok_or_else(|| anyhow::anyhow!("Job {} not found", job_id))
    }

    async fn stop_tuning_job(&self, job_id: Uuid, _reason: String) -> Result<()> {
        let mut jobs = self.jobs.write().await;
        if let Some(job) = jobs.get_mut(&job_id) {
            job.status = TuningJobStatus::Stopped;
            job.completed_at = Some(Utc::now());
        }
        Ok(())
    }
}

// ============================================================================
// TEST HELPERS
// ============================================================================

/// Create a mock tuning manager with auto-complete enabled
fn create_mock_manager() -> Arc<dyn TuningManagerTrait> {
    Arc::new(MockTuningManager::new())
}

/// Create a mock tuning manager with specific failure models
fn create_mock_manager_with_failures(failure_models: Vec<String>) -> Arc<dyn TuningManagerTrait> {
    Arc::new(MockTuningManager::with_failures(failure_models))
}

/// Wait for batch job to complete (with timeout)
async fn wait_for_completion(
    manager: &BatchTuningManager,
    batch_id: Uuid,
    timeout_secs: u64,
) -> Result<BatchTuningJob> {
    let start = std::time::Instant::now();
    let timeout = std::time::Duration::from_secs(timeout_secs);

    loop {
        let status = manager.get_batch_status(batch_id).await?;
        
        match status.status {
            BatchJobStatus::Completed 
            | BatchJobStatus::Failed 
            | BatchJobStatus::PartiallyCompleted 
            | BatchJobStatus::Stopped => {
                return Ok(status);
            }
            _ => {
                if start.elapsed() > timeout {
                    return Err(anyhow::anyhow!("Batch job timed out after {}s", timeout_secs));
                }
                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
            }
        }
    }
}

// ============================================================================
// TEST 1: Batch Job Creation
// ============================================================================

#[tokio::test]
async fn test_batch_job_creation() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());
    let models = vec!["DQN".to_string(), "PPO".to_string()];

    let result = manager
        .start_batch_tuning(
            models.clone(),
            10,
            "tuning_config.yaml".to_string(),
            None,
            false,
            None,
        )
        .await;

    assert!(result.is_ok(), "Failed to create batch job: {:?}", result.err());
    let batch_id = result.unwrap();
    assert_ne!(batch_id, Uuid::nil());

    // Verify job can be retrieved
    let status = manager.get_batch_status(batch_id).await;
    assert!(status.is_ok());
    let job = status.unwrap();
    assert_eq!(job.batch_id, batch_id);
    assert_eq!(job.models, models);
}

// ============================================================================
// TEST 2: Model Dependency Resolution
// ============================================================================

#[tokio::test]
async fn test_model_dependency_resolution() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    // Test case 1: TFT depends on MAMBA_2 (should order MAMBA_2 first)
    let models = vec!["TFT".to_string(), "MAMBA_2".to_string()];
    let resolved = manager.resolve_model_dependencies(&models);

    assert_eq!(resolved.len(), 2);
    let mamba_idx = resolved.iter().position(|m| m == "MAMBA_2").unwrap();
    let tft_idx = resolved.iter().position(|m| m == "TFT").unwrap();
    assert!(mamba_idx < tft_idx, "MAMBA_2 must come before TFT");
}

#[tokio::test]
async fn test_independent_models_no_ordering() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    // DQN and PPO are independent - can run in any order
    let models = vec!["DQN".to_string(), "PPO".to_string()];
    let resolved = manager.resolve_model_dependencies(&models);

    assert_eq!(resolved.len(), 2);
    assert!(resolved.contains(&"DQN".to_string()));
    assert!(resolved.contains(&"PPO".to_string()));
}

#[tokio::test]
async fn test_complex_dependency_chain() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    // Complex case: DQN, PPO (independent), MAMBA_2, TFT (depends on MAMBA_2)
    let models = vec![
        "TFT".to_string(),
        "DQN".to_string(),
        "MAMBA_2".to_string(),
        "PPO".to_string(),
    ];
    let resolved = manager.resolve_model_dependencies(&models);

    assert_eq!(resolved.len(), 4);
    
    // MAMBA_2 must come before TFT
    let mamba_idx = resolved.iter().position(|m| m == "MAMBA_2").unwrap();
    let tft_idx = resolved.iter().position(|m| m == "TFT").unwrap();
    assert!(mamba_idx < tft_idx, "MAMBA_2 must come before TFT");
}

// ============================================================================
// TEST 3: Batch Job Status Tracking
// ============================================================================

#[tokio::test]
async fn test_batch_status_retrieval() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    let batch_id = manager
        .start_batch_tuning(
            vec!["DQN".to_string()],
            5,
            "tuning_config.yaml".to_string(),
            None,
            false,
            None,
        )
        .await
        .expect("Failed to start batch");

    let result = manager.get_batch_status(batch_id).await;
    assert!(result.is_ok());
    
    let status = result.unwrap();
    assert_eq!(status.batch_id, batch_id);
}

#[tokio::test]
async fn test_batch_status_progress_tracking() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    let models = vec!["DQN".to_string(), "PPO".to_string()];
    let batch_id = manager
        .start_batch_tuning(models, 10, "tuning_config.yaml".to_string(), None, false, None)
        .await
        .expect("Failed to start batch");

    // Wait for completion
    let final_status = wait_for_completion(&manager, batch_id, 30).await;
    assert!(final_status.is_ok(), "Batch did not complete: {:?}", final_status.err());

    let status = final_status.unwrap();
    assert_eq!(status.status, BatchJobStatus::Completed);
    assert_eq!(status.results.len(), 2);
}

// ============================================================================
// TEST 4: Automatic YAML Export
// ============================================================================

#[tokio::test]
async fn test_automatic_yaml_export() {
    use std::fs;

    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());
    let output_path = "/tmp/test_best_hyperparameters.yaml";

    let batch_id = manager
        .start_batch_tuning(
            vec!["DQN".to_string()],
            5,
            "tuning_config.yaml".to_string(),
            None,
            true, // auto_export_yaml
            Some(output_path.to_string()),
        )
        .await
        .expect("Failed to start batch");

    // Wait for completion
    let _ = wait_for_completion(&manager, batch_id, 30).await.expect("Batch did not complete");

    // Verify YAML was exported
    assert!(std::path::Path::new(output_path).exists(), "YAML file was not created");
    
    let yaml_content = fs::read_to_string(output_path).expect("Failed to read YAML");
    assert!(yaml_content.contains("DQN"), "YAML does not contain DQN");
    assert!(yaml_content.contains("learning_rate"), "YAML does not contain learning_rate");

    // Cleanup
    let _ = fs::remove_file(output_path);
}

#[tokio::test]
async fn test_yaml_export_format() {
    use std::fs;

    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());
    let output_path = "/tmp/test_yaml_format.yaml";

    let batch_id = manager
        .start_batch_tuning(
            vec!["DQN".to_string(), "PPO".to_string()],
            5,
            "tuning_config.yaml".to_string(),
            None,
            false,
            Some(output_path.to_string()),
        )
        .await
        .expect("Failed to start batch");

    // Wait for completion
    let _ = wait_for_completion(&manager, batch_id, 30).await.expect("Batch did not complete");

    // Manual export
    let export_result = manager.export_best_hyperparameters(batch_id, output_path).await;
    assert!(export_result.is_ok(), "Failed to export YAML: {:?}", export_result.err());

    let yaml_content = fs::read_to_string(output_path).expect("Failed to read YAML");
    assert!(yaml_content.contains("models:"));
    assert!(yaml_content.contains("hyperparameters:"));
    assert!(yaml_content.contains("metrics:"));

    // Cleanup
    let _ = fs::remove_file(output_path);
}

// ============================================================================
// TEST 5: Consolidated Reporting
// ============================================================================

#[tokio::test]
async fn test_consolidated_report_generation() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    let batch_id = manager
        .start_batch_tuning(
            vec!["DQN".to_string(), "PPO".to_string()],
            5,
            "tuning_config.yaml".to_string(),
            None,
            false,
            None,
        )
        .await
        .expect("Failed to start batch");

    // Wait for completion
    let _ = wait_for_completion(&manager, batch_id, 30).await.expect("Batch did not complete");

    let result = manager.generate_consolidated_report(batch_id).await;
    assert!(result.is_ok(), "Failed to generate report: {:?}", result.err());

    let report = result.unwrap();
    assert!(report.contains("BATCH TUNING CONSOLIDATED REPORT"));
    assert!(report.contains("DQN"));
    assert!(report.contains("PPO"));
    assert!(report.contains("Best Sharpe Ratio"));
}

#[tokio::test]
async fn test_consolidated_report_content() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    let batch_id = manager
        .start_batch_tuning(
            vec!["DQN".to_string(), "PPO".to_string()],
            10,
            "tuning_config.yaml".to_string(),
            None,
            false,
            None,
        )
        .await
        .expect("Failed to start batch");

    // Wait for completion
    let _ = wait_for_completion(&manager, batch_id, 30).await.expect("Batch did not complete");

    let report = manager.generate_consolidated_report(batch_id).await.unwrap();

    // Verify report contains key sections
    assert!(report.contains("Batch ID:"));
    assert!(report.contains("PER-MODEL RESULTS"));
    assert!(report.contains("MODEL COMPARISON"));
    assert!(report.contains("RECOMMENDATION"));
    assert!(report.contains("EXPORT INFORMATION"));
}

// ============================================================================
// TEST 6: Sequential Execution with Dependencies
// ============================================================================

#[tokio::test]
async fn test_sequential_execution_order() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    let models = vec![
        "TFT".to_string(),
        "MAMBA_2".to_string(),
        "DQN".to_string(),
    ];

    let batch_id = manager
        .start_batch_tuning(models, 5, "tuning_config.yaml".to_string(), None, false, None)
        .await
        .expect("Failed to start batch");

    // Wait for completion
    let final_status = wait_for_completion(&manager, batch_id, 30).await.unwrap();

    // Check that MAMBA_2 completed before TFT
    let mamba_result = final_status.results.iter()
        .find(|r| r.model_type == "MAMBA_2")
        .expect("MAMBA_2 result not found");
    
    let tft_result = final_status.results.iter()
        .find(|r| r.model_type == "TFT")
        .expect("TFT result not found");

    assert!(mamba_result.completed_at < tft_result.completed_at,
            "MAMBA_2 should complete before TFT");
}

// ============================================================================
// TEST 7: Error Handling - Model Failure
// ============================================================================

#[tokio::test]
async fn test_model_failure_continues_batch() {
    let mock_tuning = create_mock_manager_with_failures(vec!["INVALID_MODEL".to_string()]);
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    // Invalid model should be rejected at validation
    let result = manager
        .start_batch_tuning(
            vec!["DQN".to_string(), "INVALID_MODEL".to_string(), "PPO".to_string()],
            5,
            "tuning_config.yaml".to_string(),
            None,
            false,
            None,
        )
        .await;

    // Should fail validation
    assert!(result.is_err(), "Expected validation error for INVALID_MODEL");
}

#[tokio::test]
async fn test_model_failure_partial_completion() {
    let mock_tuning = create_mock_manager_with_failures(vec!["PPO".to_string()]);
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    let batch_id = manager
        .start_batch_tuning(
            vec!["DQN".to_string(), "PPO".to_string()],
            5,
            "tuning_config.yaml".to_string(),
            None,
            false,
            None,
        )
        .await
        .expect("Failed to start batch");

    // Wait for completion
    let final_status = wait_for_completion(&manager, batch_id, 30).await.unwrap();

    // Status should be PartiallyCompleted
    assert_eq!(final_status.status, BatchJobStatus::PartiallyCompleted);
    assert_eq!(final_status.results.len(), 2);

    // DQN should succeed, PPO should fail
    let dqn_result = final_status.results.iter()
        .find(|r| r.model_type == "DQN")
        .unwrap();
    assert_eq!(dqn_result.status, TuningJobStatus::Completed);

    let ppo_result = final_status.results.iter()
        .find(|r| r.model_type == "PPO")
        .unwrap();
    assert_eq!(ppo_result.status, TuningJobStatus::Failed);
    assert!(ppo_result.error_message.is_some());
}

// ============================================================================
// TEST 8: Batch Job Cancellation
// ============================================================================

#[tokio::test]
async fn test_batch_job_cancellation() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    let batch_id = manager
        .start_batch_tuning(
            vec!["DQN".to_string(), "PPO".to_string(), "MAMBA_2".to_string()],
            50,
            "tuning_config.yaml".to_string(),
            None,
            false,
            None,
        )
        .await
        .expect("Failed to start batch");

    // Wait briefly for job to start
    tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;

    // Cancel the batch
    let cancel_result = manager.stop_batch_job(batch_id, "User cancellation".to_string()).await;
    assert!(cancel_result.is_ok(), "Failed to cancel batch: {:?}", cancel_result.err());

    let status = manager.get_batch_status(batch_id).await.unwrap();
    assert_eq!(status.status, BatchJobStatus::Stopped);
}

// ============================================================================
// TEST 9: Results Comparison
// ============================================================================

#[tokio::test]
async fn test_results_comparison() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    let batch_id = manager
        .start_batch_tuning(
            vec!["DQN".to_string(), "PPO".to_string(), "MAMBA_2".to_string()],
            10,
            "tuning_config.yaml".to_string(),
            None,
            false,
            None,
        )
        .await
        .expect("Failed to start batch");

    // Wait for completion
    let _ = wait_for_completion(&manager, batch_id, 30).await.expect("Batch did not complete");

    let report = manager.generate_consolidated_report(batch_id).await.unwrap();
    
    // Report should contain comparison and recommendation
    assert!(report.contains("Best Overall Model:"));
    assert!(report.contains("Sharpe Ratio"));
    assert!(report.contains("RECOMMENDATION"));
}

// ============================================================================
// TEST 10: YAML Export Path Validation
// ============================================================================

#[tokio::test]
async fn test_yaml_export_path_validation() {
    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch".to_string());

    let batch_id = manager
        .start_batch_tuning(
            vec!["DQN".to_string()],
            5,
            "tuning_config.yaml".to_string(),
            None,
            false,
            None,
        )
        .await
        .expect("Failed to start batch");

    // Wait for completion
    let _ = wait_for_completion(&manager, batch_id, 30).await.expect("Batch did not complete");

    // Test with valid path (should create directories)
    let valid_path = "/tmp/test_batch_export/best_params.yaml";
    let result = manager.export_best_hyperparameters(batch_id, valid_path).await;
    assert!(result.is_ok(), "Failed to export to valid path: {:?}", result.err());

    // Cleanup
    let _ = std::fs::remove_file(valid_path);
    let _ = std::fs::remove_dir("/tmp/test_batch_export");
}

// ============================================================================
// INTEGRATION TEST: Full Batch Tuning Flow
// ============================================================================

#[tokio::test]
#[ignore] // Only run with --ignored flag (integration test)
async fn test_full_batch_tuning_flow_e2e() {
    use std::fs;

    let mock_tuning = create_mock_manager();
    let manager = BatchTuningManager::new(mock_tuning, "/tmp/test_batch_e2e".to_string());

    // Full E2E test with 2 models, 10 trials each
    let models = vec!["DQN".to_string(), "PPO".to_string()];
    let batch_id = manager
        .start_batch_tuning(models, 10, "tuning_config.yaml".to_string(), None, true, None)
        .await
        .expect("Failed to start batch job");

    println!("Batch job started: {}", batch_id);

    // Wait for completion (timeout 5 minutes for safety)
    let final_status = wait_for_completion(&manager, batch_id, 300)
        .await
        .expect("Batch job did not complete");

    println!("Batch job completed with status: {:?}", final_status.status);

    // Verify all models ran
    assert_eq!(final_status.results.len(), 2);
    assert!(matches!(
        final_status.status,
        BatchJobStatus::Completed | BatchJobStatus::PartiallyCompleted
    ));

    // Generate report
    let report = manager.generate_consolidated_report(batch_id).await.unwrap();
    println!("=== CONSOLIDATED REPORT ===\n{}", report);

    // Verify report contains expected sections
    assert!(report.contains("BATCH TUNING CONSOLIDATED REPORT"));
    assert!(report.contains("PER-MODEL RESULTS"));
    assert!(report.contains("MODEL COMPARISON"));
}