foxhunt/tests/e2e/tests/e2e_ml_training_test.rs

//! End-to-End ML Training Pipeline Test
//!
//! Mission: Validate complete training pipeline from DBN data → checkpoint → registry
//! Methodology: TDD (RED → GREEN → REFACTOR)
//!
//! Tests cover:
//! 1. DBN data loading → checkpoint creation → registry storage
//! 2. All 4 models training end-to-end (DQN, PPO, MAMBA2, TFT)
//! 3. Multi-symbol training
//! 4. Checkpoint loading and inference validation
//! 5. Training metrics validation
//! 6. GPU memory optimization during training

use anyhow::{Context, Result};
use candle_core::Device;
use sqlx::PgPool;
use std::collections::HashMap;
use std::path::PathBuf;
use tokio::fs;
use tracing::{info, warn};
use uuid::Uuid;

// Import ML training infrastructure
use ml::data_loaders::dbn_sequence_loader::DbnSequenceLoader;
use ml::training::unified_trainer::{TrainingConfig, UnifiedTrainer};
use ml_training_service::checkpoint_manager::{CheckpointManager, RetentionPolicy};
use ml::checkpoint::CheckpointMetadata;
use ml::ModelType;

// ============================================================================
// Helper Functions (Test Infrastructure)
// ============================================================================

/// Get test database pool
async fn get_test_db_pool() -> PgPool {
    let database_url = std::env::var("DATABASE_URL").unwrap_or_else(|_| {
        "postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt".to_string()
    });

    PgPool::connect(&database_url)
        .await
        .expect("Failed to connect to test database")
}

/// Get test data directory path
fn get_test_data_path() -> PathBuf {
    let manifest_dir = env!("CARGO_MANIFEST_DIR");
    let workspace_root = PathBuf::from(manifest_dir)
        .parent()
        .unwrap()
        .parent()
        .unwrap()
        .to_path_buf();
    workspace_root.join("test_data")
}

/// Check if test data exists
fn test_data_available(symbol: &str) -> bool {
    let data_path = get_test_data_path();
    let dbn_file = data_path.join(format!("{}.20240102.dbn", symbol));
    dbn_file.exists()
}

/// Load DBN data for testing
async fn load_dbn_bars(symbol: &str, num_bars: usize) -> Result<Vec<(f64, f64, f64, f64, f64)>> {
    let data_path = get_test_data_path();
    let dbn_file = data_path.join(format!("{}.20240102.dbn", symbol));

    if !dbn_file.exists() {
        return Err(anyhow::anyhow!(
            "DBN file not found: {}",
            dbn_file.display()
        ));
    }

    let loader = DbnSequenceLoader::new(
        dbn_file.to_str().unwrap(),
        32, // batch_size
        1,  // sequence_length
        Some(symbol.to_string()),
    ).await?;

    // Extract OHLCV bars (simplified for test)
    let bars = vec![(4500.0, 4510.0, 4495.0, 4505.0, 1000.0); num_bars];
    Ok(bars)
}

/// Create test output directory
async fn create_test_output_dir(test_name: &str) -> Result<PathBuf> {
    let output_dir = PathBuf::from(format!("/tmp/foxhunt_e2e_test_{}", test_name));

    if output_dir.exists() {
        fs::remove_dir_all(&output_dir).await?;
    }
    fs::create_dir_all(&output_dir).await?;

    Ok(output_dir)
}

// ============================================================================
// TEST 1: E2E Training Pipeline - DBN to Checkpoint (RED)
// ============================================================================

#[tokio::test]
#[ignore = "RED phase - will fail until full implementation exists"]
async fn test_e2e_dbn_to_checkpoint() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init()
        .ok();

    info!("🚀 Starting E2E Training Pipeline Test: DBN → Checkpoint → Registry");

    // Skip if test data not available
    if !test_data_available("ES.FUT") {
        warn!("Skipping test - ES.FUT test data not available");
        return Ok(());
    }

    // ARRANGE
    let pool = get_test_db_pool().await;
    let output_dir = create_test_output_dir("dbn_to_checkpoint").await?;

    // Step 1: Load DBN data
    info!("📂 Step 1: Loading DBN data...");
    let dbn_path = get_test_data_path().join("ES.FUT.20240102.dbn");
    let loader = DbnSequenceLoader::new(
        dbn_path.to_str().unwrap(),
        32,
        1,
        Some("ES.FUT".to_string()),
    ).await?;
    info!("✅ DBN loader initialized");

    // Step 2: Configure training
    info!("⚙️  Step 2: Configuring DQN training...");
    let config = TrainingConfig {
        model_type: "DQN".to_string(),
        epochs: 10,
        batch_size: 32,
        learning_rate: 0.001,
        device: Device::Cpu, // Use CPU for E2E test
        checkpoint_dir: output_dir.clone(),
        symbol: "ES.FUT".to_string(),
    };

    // Step 3: Create trainer and train model
    info!("🏋️  Step 3: Training DQN model (10 epochs)...");
    let mut trainer = UnifiedTrainer::new(config)?;

    let start_time = std::time::Instant::now();
    let metrics = trainer.train(&loader).await?;
    let training_duration = start_time.elapsed();

    info!(
        "✅ Training completed in {:.1}s",
        training_duration.as_secs_f64()
    );
    info!("📊 Final loss: {:.6}", metrics.final_loss);
    info!("📊 Epochs trained: {}", metrics.epochs_completed);

    // ACT & ASSERT: Step 4 - Verify checkpoint exists
    info!("💾 Step 4: Verifying checkpoint creation...");
    let checkpoint_path = output_dir.join("dqn_final.safetensors");

    assert!(
        checkpoint_path.exists(),
        "Checkpoint file should exist at {:?}",
        checkpoint_path
    );

    let checkpoint_size = fs::metadata(&checkpoint_path).await?.len();
    info!("✅ Checkpoint created: {} bytes", checkpoint_size);

    assert!(
        checkpoint_size > 1_000,
        "Checkpoint file should be at least 1KB, got {} bytes",
        checkpoint_size
    );

    // Step 5: Verify checkpoint is loadable
    info!("🔄 Step 5: Loading checkpoint for validation...");
    let checkpoint_data = fs::read(&checkpoint_path).await?;

    assert!(
        !checkpoint_data.is_empty(),
        "Checkpoint data should not be empty"
    );
    info!(
        "✅ Checkpoint loaded successfully ({} bytes)",
        checkpoint_data.len()
    );

    // Step 6: Register checkpoint in model registry
    info!("📝 Step 6: Registering checkpoint in model registry...");
    let manager = CheckpointManager::new(pool.clone(), RetentionPolicy::default()).await?;

    // Create checkpoint metadata
    let metadata = CheckpointMetadata {
        checkpoint_id: format!("dqn-v{}", "1.0.0"),
        model_type: ModelType::DQN,
        model_name: "ES.FUT-DQN-e2e".to_string(),
        version: "1.0.0".to_string(),
        created_at: chrono::Utc::now(),
        epoch: Some(10),
        step: Some(100),
        loss: Some(metrics.final_loss),
        accuracy: None,
        hyperparameters: std::collections::HashMap::from([
            ("epochs".to_string(), serde_json::json!(10)),
            ("batch_size".to_string(), serde_json::json!(32)),
            ("learning_rate".to_string(), serde_json::json!(0.001)),
        ]),
        metrics: std::collections::HashMap::from([
            ("final_loss".to_string(), metrics.final_loss),
        ]),
        architecture: std::collections::HashMap::new(),
        format: ml::checkpoint::CheckpointFormat::Binary,
        compression: ml::checkpoint::CompressionType::LZ4,
        file_size: checkpoint_size as usize,
        compressed_size: None,
        checksum: "test_checksum_e2e".to_string(),
        tags: vec!["e2e_test".to_string()],
        custom_metadata: std::collections::HashMap::new(),
        signature: None,
        signature_algorithm: "none".to_string(),
        signing_key_id: "test".to_string(),
        signed_at: None,
    };

    let registration_id = manager.register_checkpoint(metadata).await?;

    info!("✅ Checkpoint registered with ID: {}", registration_id);

    // Verify registration in database
    let record = sqlx::query!(
        r#"
        SELECT model_type, model_id, is_production, is_experimental
        FROM ml_model_versions
        WHERE model_id = $1
        "#,
        registration_id
    )
    .fetch_one(&pool)
    .await?;

    assert_eq!(record.model_type, "DQN", "Model type should be DQN");
    assert_eq!(record.is_experimental, true, "Should be experimental");
    assert_eq!(record.is_production, false, "Should not be production");

    info!("✅ Database registration verified");

    // Cleanup
    fs::remove_dir_all(&output_dir).await?;

    info!("🎉 E2E Training Pipeline Test PASSED!");
    Ok(())
}

// ============================================================================
// TEST 2: E2E All Models Training (RED)
// ============================================================================

#[tokio::test]
#[ignore = "RED phase - will fail until all models implemented"]
async fn test_e2e_all_models_training() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init()
        .ok();

    info!("🚀 Starting E2E All Models Training Test");

    // Skip if test data not available
    if !test_data_available("ES.FUT") {
        warn!("Skipping test - ES.FUT test data not available");
        return Ok(());
    }

    // ARRANGE
    let dbn_path = get_test_data_path().join("ES.FUT.20240102.dbn");
    let models = vec!["DQN", "PPO", "MAMBA2", "TFT"];
    let output_dir = create_test_output_dir("all_models").await?;

    // ACT & ASSERT: Train each model
    for model_name in models {
        info!("\n🏋️  Training {} model...", model_name);

        let model_output_dir = output_dir.join(model_name.to_lowercase());
        fs::create_dir_all(&model_output_dir).await?;

        let config = TrainingConfig {
            model_type: model_name.to_string(),
            epochs: 5, // Shorter for E2E test
            batch_size: 32,
            learning_rate: 0.001,
            device: Device::Cpu,
            checkpoint_dir: model_output_dir.clone(),
            symbol: "ES.FUT".to_string(),
        };

        let loader = DbnSequenceLoader::new(
            dbn_path.to_str().unwrap(),
            32,
            1,
            Some("ES.FUT".to_string()),
        ).await?;

        let mut trainer = UnifiedTrainer::new(config)?;
        let metrics = trainer.train(&loader).await?;

        info!(
            "✅ {} trained: {} epochs, loss={:.6}",
            model_name, metrics.epochs_completed, metrics.final_loss
        );

        // Verify checkpoint created
        let checkpoint_path =
            model_output_dir.join(format!("{}_final.safetensors", model_name.to_lowercase()));

        assert!(
            checkpoint_path.exists(),
            "{} checkpoint should exist",
            model_name
        );

        let checkpoint_size = fs::metadata(&checkpoint_path).await?.len();
        info!("💾 {} checkpoint: {} bytes", model_name, checkpoint_size);

        assert!(
            checkpoint_size > 1_000,
            "{} checkpoint should be at least 1KB",
            model_name
        );
    }

    // Cleanup
    fs::remove_dir_all(&output_dir).await?;

    info!("\n🎉 E2E All Models Training Test PASSED!");
    Ok(())
}

// ============================================================================
// TEST 3: E2E Multi-Symbol Training (RED)
// ============================================================================

#[tokio::test]
#[ignore = "RED phase - will fail until multi-symbol support implemented"]
async fn test_e2e_multi_symbol_training() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init()
        .ok();

    info!("🚀 Starting E2E Multi-Symbol Training Test");

    // ARRANGE
    let symbols = vec!["ES.FUT", "NQ.FUT", "ZN.FUT"];
    let output_dir = create_test_output_dir("multi_symbol").await?;
    let pool = get_test_db_pool().await;
    let manager = CheckpointManager::new(pool, RetentionPolicy::default()).await?;

    let mut trained_symbols = Vec::new();

    // ACT: Train DQN on each available symbol
    for symbol in symbols {
        if !test_data_available(symbol) {
            warn!("Skipping {} - data not available", symbol);
            continue;
        }

        info!("\n🏋️  Training DQN on {}...", symbol);

        let dbn_path = get_test_data_path().join(format!("{}.20240102.dbn", symbol));
        let symbol_output_dir = output_dir.join(symbol.replace(".", "_"));
        fs::create_dir_all(&symbol_output_dir).await?;

        let config = TrainingConfig {
            model_type: "DQN".to_string(),
            epochs: 5,
            batch_size: 32,
            learning_rate: 0.001,
            device: Device::Cpu,
            checkpoint_dir: symbol_output_dir.clone(),
            symbol: symbol.to_string(),
        };

        let loader =
            DbnSequenceLoader::new(dbn_path.to_str().unwrap(), 32, 1, Some(symbol.to_string())).await?;

        let mut trainer = UnifiedTrainer::new(config)?;
        let metrics = trainer.train(&loader).await?;

        info!("✅ {} trained: loss={:.6}", symbol, metrics.final_loss);

        // Register checkpoint
        let checkpoint_path = symbol_output_dir.join("dqn_final.safetensors");
        let checkpoint_size = fs::metadata(&checkpoint_path).await?.len();

        let metadata = CheckpointMetadata {
            checkpoint_id: format!("dqn-{}-v1.0.0", symbol.replace(".", "_")),
            model_type: ModelType::DQN,
            model_name: format!("{}-DQN-multi", symbol),
            version: "1.0.0".to_string(),
            created_at: chrono::Utc::now(),
            epoch: Some(5),
            step: Some(50),
            loss: Some(metrics.final_loss),
            accuracy: None,
            hyperparameters: std::collections::HashMap::from([
                ("epochs".to_string(), serde_json::json!(5)),
                ("batch_size".to_string(), serde_json::json!(32)),
            ]),
            metrics: std::collections::HashMap::from([
                ("final_loss".to_string(), metrics.final_loss),
            ]),
            architecture: std::collections::HashMap::new(),
            format: ml::checkpoint::CheckpointFormat::Binary,
            compression: ml::checkpoint::CompressionType::LZ4,
            file_size: checkpoint_size as usize,
            compressed_size: None,
            checksum: format!("checksum_{}", symbol),
            tags: vec!["e2e_multi_symbol".to_string()],
            custom_metadata: std::collections::HashMap::new(),
            signature: None,
            signature_algorithm: "none".to_string(),
            signing_key_id: "test".to_string(),
            signed_at: None,
        };

        let registration_id = manager.register_checkpoint(metadata).await?;

        info!("📝 {} checkpoint registered: {}", symbol, registration_id);

        trained_symbols.push(symbol);
    }

    // ASSERT: At least one symbol should be trained
    assert!(
        !trained_symbols.is_empty(),
        "At least one symbol should be successfully trained"
    );

    info!(
        "\n✅ Successfully trained on {} symbols: {:?}",
        trained_symbols.len(),
        trained_symbols
    );

    // Cleanup
    fs::remove_dir_all(&output_dir).await?;

    info!("🎉 E2E Multi-Symbol Training Test PASSED!");
    Ok(())
}

// ============================================================================
// TEST 4: E2E Training Metrics Validation (RED)
// ============================================================================

#[tokio::test]
#[ignore = "RED phase - will fail until metrics validation implemented"]
async fn test_e2e_training_metrics_validation() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init()
        .ok();

    info!("🚀 Starting E2E Training Metrics Validation Test");

    if !test_data_available("ES.FUT") {
        warn!("Skipping test - ES.FUT test data not available");
        return Ok(());
    }

    // ARRANGE
    let dbn_path = get_test_data_path().join("ES.FUT.20240102.dbn");
    let output_dir = create_test_output_dir("metrics_validation").await?;

    let config = TrainingConfig {
        model_type: "DQN".to_string(),
        epochs: 10,
        batch_size: 32,
        learning_rate: 0.001,
        device: Device::Cpu,
        checkpoint_dir: output_dir.clone(),
        symbol: "ES.FUT".to_string(),
    };

    let loader = DbnSequenceLoader::new(
        dbn_path.to_str().unwrap(),
        32,
        1,
        Some("ES.FUT".to_string()),
    ).await?;

    // ACT: Train model and collect metrics
    let mut trainer = UnifiedTrainer::new(config)?;
    let metrics = trainer.train(&loader).await?;

    // ASSERT: Validate metrics
    info!("📊 Validating training metrics...");

    // Loss should be finite and positive
    assert!(
        metrics.final_loss.is_finite() && metrics.final_loss >= 0.0,
        "Loss should be finite and non-negative, got: {}",
        metrics.final_loss
    );
    info!("✅ Loss is valid: {:.6}", metrics.final_loss);

    // Epochs completed should match configuration
    assert_eq!(
        metrics.epochs_completed, 10,
        "Should complete 10 epochs, got {}",
        metrics.epochs_completed
    );
    info!("✅ Epochs completed: {}", metrics.epochs_completed);

    // Training time should be reasonable
    assert!(
        metrics.training_time_seconds > 0.0,
        "Training time should be positive, got: {}",
        metrics.training_time_seconds
    );
    info!("✅ Training time: {:.1}s", metrics.training_time_seconds);

    // Convergence metrics
    if let Some(convergence) = metrics.convergence_achieved {
        info!("✅ Convergence achieved: {}", convergence);
    }

    // Loss trajectory should show improvement
    if metrics.loss_history.len() >= 2 {
        let initial_loss = metrics.loss_history.first().unwrap();
        let final_loss = metrics.loss_history.last().unwrap();

        info!("📈 Initial loss: {:.6}", initial_loss);
        info!("📉 Final loss: {:.6}", final_loss);

        // Loss should generally decrease (allowing some fluctuation)
        let improvement_ratio = (initial_loss - final_loss) / initial_loss;
        info!("📊 Improvement: {:.1}%", improvement_ratio * 100.0);
    }

    // Cleanup
    fs::remove_dir_all(&output_dir).await?;

    info!("🎉 E2E Training Metrics Validation Test PASSED!");
    Ok(())
}

// ============================================================================
// TEST 5: E2E Checkpoint Loading and Inference (RED)
// ============================================================================

#[tokio::test]
#[ignore = "RED phase - will fail until inference validation implemented"]
async fn test_e2e_checkpoint_loading_and_inference() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init()
        .ok();

    info!("🚀 Starting E2E Checkpoint Loading and Inference Test");

    if !test_data_available("ES.FUT") {
        warn!("Skipping test - ES.FUT test data not available");
        return Ok(());
    }

    // ARRANGE: Train a model first
    let dbn_path = get_test_data_path().join("ES.FUT.20240102.dbn");
    let output_dir = create_test_output_dir("checkpoint_loading").await?;

    info!("🏋️  Step 1: Training model to create checkpoint...");

    let config = TrainingConfig {
        model_type: "DQN".to_string(),
        epochs: 5,
        batch_size: 32,
        learning_rate: 0.001,
        device: Device::Cpu,
        checkpoint_dir: output_dir.clone(),
        symbol: "ES.FUT".to_string(),
    };

    let loader = DbnSequenceLoader::new(
        dbn_path.to_str().unwrap(),
        32,
        1,
        Some("ES.FUT".to_string()),
    ).await?;

    let mut trainer = UnifiedTrainer::new(config)?;
    trainer.train(&loader).await?;

    let checkpoint_path = output_dir.join("dqn_final.safetensors");
    info!("✅ Checkpoint created: {}", checkpoint_path.display());

    // ACT: Load checkpoint and perform inference
    info!("🔄 Step 2: Loading checkpoint for inference...");

    // Load checkpoint data
    let checkpoint_data = fs::read(&checkpoint_path).await?;
    info!("✅ Checkpoint loaded: {} bytes", checkpoint_data.len());

    // Create inference engine (simplified for test)
    info!("🧠 Step 3: Performing inference...");

    // Generate test features (256 features for MAMBA2)
    let test_features = vec![0.5_f32; 256];

    // In real implementation, this would load model and run inference
    // For TDD RED phase, we just verify the checkpoint is valid format

    // Verify safetensors format (should be valid tensors)
    assert!(
        checkpoint_data.len() > 100,
        "Checkpoint should contain valid model weights"
    );

    info!("✅ Inference validation successful");

    // Cleanup
    fs::remove_dir_all(&output_dir).await?;

    info!("🎉 E2E Checkpoint Loading and Inference Test PASSED!");
    Ok(())
}

// ============================================================================
// TEST 6: E2E GPU Memory Optimization (RED)
// ============================================================================

#[tokio::test]
#[ignore = "RED phase - requires GPU"]
async fn test_e2e_gpu_memory_optimization() -> Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .try_init()
        .ok();

    info!("🚀 Starting E2E GPU Memory Optimization Test");

    // Check if GPU is available
    let device = match Device::cuda_if_available(0) {
        Ok(dev) => dev,
        Err(_) => {
            warn!("GPU not available, skipping test");
            return Ok(());
        },
    };

    info!("✅ GPU detected: {:?}", device);

    if !test_data_available("ES.FUT") {
        warn!("Skipping test - ES.FUT test data not available");
        return Ok(());
    }

    // ARRANGE
    let dbn_path = get_test_data_path().join("ES.FUT.20240102.dbn");
    let output_dir = create_test_output_dir("gpu_memory").await?;

    // ACT: Train with GPU memory constraints
    let config = TrainingConfig {
        model_type: "MAMBA2".to_string(),
        epochs: 5,
        batch_size: 16, // Smaller batch for GPU memory
        learning_rate: 0.001,
        device,
        checkpoint_dir: output_dir.clone(),
        symbol: "ES.FUT".to_string(),
    };

    let loader = DbnSequenceLoader::new(
        dbn_path.to_str().unwrap(),
        16,
        1,
        Some("ES.FUT".to_string()),
    ).await?;

    let mut trainer = UnifiedTrainer::new(config)?;

    info!("🏋️  Training MAMBA2 on GPU with memory optimization...");
    let start_time = std::time::Instant::now();
    let metrics = trainer.train(&loader).await?;
    let training_duration = start_time.elapsed();

    // ASSERT: Training should complete without OOM
    info!(
        "✅ GPU training completed in {:.1}s",
        training_duration.as_secs_f64()
    );
    info!("📊 Final loss: {:.6}", metrics.final_loss);

    assert!(
        metrics.epochs_completed == 5,
        "Should complete all 5 epochs without OOM"
    );

    // Verify checkpoint created
    let checkpoint_path = output_dir.join("mamba2_final.safetensors");
    assert!(checkpoint_path.exists(), "GPU checkpoint should be created");

    // Cleanup
    fs::remove_dir_all(&output_dir).await?;

    info!("🎉 E2E GPU Memory Optimization Test PASSED!");
    Ok(())
}