foxhunt/ml/tests/hyperopt_integration_test.rs

//! Integration Tests for Hyperparameter Optimization
//!
//! These tests verify end-to-end optimization workflows with real training data.
//! Fast tests run quickly, while slow tests (marked with #[ignore]) perform
//! full optimization runs.

use anyhow::Result;
use ml::hyperopt::{
    optimize_mamba2, BestHyperparameters, EgoboxOptimizationResult as OptimizationResult,
    EgoboxTrialResult as TrialResult, HyperparameterSpace,
};
use std::path::Path;

// ============================================================================
// HELPER FUNCTIONS
// ============================================================================

/// Verify test data exists
fn ensure_test_data_exists(path: &str) -> Result<()> {
    if !Path::new(path).exists() {
        anyhow::bail!(
            "Test data not found: {}. Run 'cargo test --workspace' to generate test data.",
            path
        );
    }
    Ok(())
}

// ============================================================================
// FAST INTEGRATION TESTS
// ============================================================================

#[test]
fn test_hyperparameter_space_default_creation() {
    let space = HyperparameterSpace::default();

    // Verify default space is reasonable
    assert!(space.learning_rate_log_min < space.learning_rate_log_max);
    assert!(space.batch_size_min < space.batch_size_max);
    assert!(space.dropout_min < space.dropout_max);
    assert!(space.weight_decay_log_min < space.weight_decay_log_max);

    // Verify bounds are production-safe
    assert!(space.learning_rate_log_min >= -6.0, "LR min too small");
    assert!(space.learning_rate_log_max <= -1.0, "LR max too large");
    assert!(space.batch_size_min >= 4, "Batch size too small");
    assert!(space.batch_size_max <= 512, "Batch size too large");
    assert!(space.dropout_min >= 0.0, "Dropout min invalid");
    assert!(space.dropout_max <= 1.0, "Dropout max invalid");
}

#[test]
fn test_custom_search_space_creation() {
    let custom_space = HyperparameterSpace {
        learning_rate_log_min: -4.5,
        learning_rate_log_max: -1.5,
        batch_size_min: 32,
        batch_size_max: 128,
        dropout_min: 0.1,
        dropout_max: 0.4,
        weight_decay_log_min: -5.5,
        weight_decay_log_max: -2.5,
    };

    // Verify custom space is valid
    assert!(custom_space.learning_rate_log_min < custom_space.learning_rate_log_max);
    assert!(custom_space.batch_size_min < custom_space.batch_size_max);
    assert!(custom_space.dropout_min < custom_space.dropout_max);
    assert!(custom_space.weight_decay_log_min < custom_space.weight_decay_log_max);
}

#[test]
fn test_hyperparameter_space_serialization() {
    let space = HyperparameterSpace::default();

    // Serialize to YAML
    let yaml =
        serde_yaml::to_string(&space).expect("Failed to serialize HyperparameterSpace to YAML");
    assert!(yaml.contains("learning_rate_log_min"));
    assert!(yaml.contains("batch_size_min"));

    // Deserialize back
    let deserialized: HyperparameterSpace =
        serde_yaml::from_str(&yaml).expect("Failed to deserialize HyperparameterSpace from YAML");

    assert_eq!(
        deserialized.learning_rate_log_min,
        space.learning_rate_log_min
    );
    assert_eq!(deserialized.batch_size_min, space.batch_size_min);
    assert_eq!(deserialized.dropout_min, space.dropout_min);
}

#[test]
fn test_test_data_available() {
    // Verify we have test data for integration tests
    let small_data = "test_data/ES_FUT_small.parquet";

    if Path::new(small_data).exists() {
        // Check file size is reasonable (should be small)
        let metadata = std::fs::metadata(small_data).expect("Failed to read metadata");
        let size_kb = metadata.len() / 1024;

        assert!(
            size_kb > 0 && size_kb < 1024,
            "Test data size {} KB is unexpected",
            size_kb
        );
        println!("✓ Test data available: {} ({} KB)", small_data, size_kb);
    } else {
        println!("⚠ Test data not found: {}", small_data);
        println!("  Run full workspace tests to generate test data");
    }
}

// ============================================================================
// SLOW INTEGRATION TESTS (marked with #[ignore])
// ============================================================================

#[tokio::test]
#[ignore] // Slow test - run with: cargo test --test hyperopt_integration_test --features cuda -- --ignored
async fn test_mamba2_5trial_optimization() -> Result<()> {
    // Initialize tracing for test output
    let _ = tracing_subscriber::fmt()
        .with_env_filter("info")
        .with_test_writer()
        .try_init();

    let test_data = "test_data/ES_FUT_small.parquet";
    ensure_test_data_exists(test_data)?;

    println!("\n╔═══════════════════════════════════════════════════════════╗");
    println!("║  5-Trial MAMBA-2 Optimization Integration Test           ║");
    println!("╚═══════════════════════════════════════════════════════════╝");

    // Run 5-trial optimization
    let space = HyperparameterSpace::default();
    let result = optimize_mamba2(
        space,
        test_data,
        5,  // max_trials
        5,  // epochs_per_trial (fast for testing)
    )
    .await?;

    // Verify optimization completed
    assert_eq!(
        result.best_params.trials_used, 5,
        "Expected 5 trials to complete"
    );
    assert!(
        !result.trial_history.is_empty(),
        "Trial history should not be empty"
    );

    // Verify best loss is reasonable (not NaN, not too high)
    let best_loss = result.best_params.best_validation_loss;
    assert!(best_loss.is_finite(), "Best loss must be finite");
    assert!(best_loss > 0.0, "Best loss must be positive");
    assert!(
        best_loss < 100_000_000.0,
        "Best loss {} is unreasonably high",
        best_loss
    );

    println!("\n✓ Optimization Results:");
    println!("  Best Learning Rate: {:.6}", result.best_params.learning_rate);
    println!("  Best Batch Size: {}", result.best_params.batch_size);
    println!("  Best Dropout: {:.3}", result.best_params.dropout);
    println!("  Best Weight Decay: {:.6}", result.best_params.weight_decay);
    println!("  Best Validation Loss: {:.6}", best_loss);
    println!(
        "  Best Perplexity: {:.4}",
        result.best_params.best_validation_loss.exp()
    );

    // Verify hyperparameters are in valid ranges
    assert!(
        result.best_params.learning_rate >= 1e-6 && result.best_params.learning_rate <= 1e-1,
        "Learning rate {} out of reasonable range",
        result.best_params.learning_rate
    );
    assert!(
        result.best_params.batch_size >= 8 && result.best_params.batch_size <= 512,
        "Batch size {} out of reasonable range",
        result.best_params.batch_size
    );
    assert!(
        result.best_params.dropout >= 0.0 && result.best_params.dropout <= 0.9,
        "Dropout {} out of reasonable range",
        result.best_params.dropout
    );
    assert!(
        result.best_params.weight_decay >= 1e-8 && result.best_params.weight_decay <= 1e-1,
        "Weight decay {} out of reasonable range",
        result.best_params.weight_decay
    );

    println!("\n✓ 5-trial optimization completed successfully");

    Ok(())
}

#[tokio::test]
#[ignore] // Slow test
async fn test_optimization_with_custom_space() -> Result<()> {
    let _ = tracing_subscriber::fmt()
        .with_env_filter("info")
        .with_test_writer()
        .try_init();

    let test_data = "test_data/ES_FUT_small.parquet";
    ensure_test_data_exists(test_data)?;

    println!("\n╔═══════════════════════════════════════════════════════════╗");
    println!("║  Custom Search Space Optimization Test                   ║");
    println!("╚═══════════════════════════════════════════════════════════╝");

    // Narrower search space for faster convergence
    let custom_space = HyperparameterSpace {
        learning_rate_log_min: -4.0,  // 1e-4
        learning_rate_log_max: -2.0,  // 1e-2
        batch_size_min: 32,
        batch_size_max: 64,
        dropout_min: 0.1,
        dropout_max: 0.3,
        weight_decay_log_min: -5.0,  // 1e-5
        weight_decay_log_max: -3.0,  // 1e-3
    };

    let result = optimize_mamba2(
        custom_space,
        test_data,
        5,  // max_trials
        5,  // epochs_per_trial
    )
    .await?;

    // Verify results are within custom space
    assert!(
        result.best_params.learning_rate >= 1e-4 && result.best_params.learning_rate <= 1e-2,
        "Learning rate {} not in custom range [1e-4, 1e-2]",
        result.best_params.learning_rate
    );
    assert!(
        result.best_params.batch_size >= 32 && result.best_params.batch_size <= 64,
        "Batch size {} not in custom range [32, 64]",
        result.best_params.batch_size
    );
    assert!(
        result.best_params.dropout >= 0.1 && result.best_params.dropout <= 0.3,
        "Dropout {} not in custom range [0.1, 0.3]",
        result.best_params.dropout
    );
    assert!(
        result.best_params.weight_decay >= 1e-5 && result.best_params.weight_decay <= 1e-3,
        "Weight decay {} not in custom range [1e-5, 1e-3]",
        result.best_params.weight_decay
    );

    println!("\n✓ Custom space optimization completed successfully");

    Ok(())
}

#[test]
fn test_optimization_result_yaml_export() {
    // Create mock optimization result
    let best_params = BestHyperparameters {
        learning_rate: 0.001,
        batch_size: 64,
        dropout: 0.2,
        weight_decay: 0.0001,
        best_validation_loss: 12.5,
        trials_used: 5,
    };

    let trial_history = vec![
        TrialResult {
            trial_number: 1,
            learning_rate: 0.001,
            batch_size: 64,
            dropout: 0.2,
            weight_decay: 0.0001,
            validation_loss: 15.5,
            training_time_seconds: 18.0,
        },
        TrialResult {
            trial_number: 2,
            learning_rate: 0.002,
            batch_size: 32,
            dropout: 0.3,
            weight_decay: 0.0002,
            validation_loss: 14.2,
            training_time_seconds: 17.5,
        },
        TrialResult {
            trial_number: 5,
            learning_rate: 0.001,
            batch_size: 64,
            dropout: 0.2,
            weight_decay: 0.0001,
            validation_loss: 12.5,
            training_time_seconds: 18.5,
        },
    ];

    let result = OptimizationResult {
        best_params,
        trial_history,
    };

    // Export to YAML
    let yaml = serde_yaml::to_string(&result).expect("Failed to serialize to YAML");

    // Verify YAML structure
    assert!(yaml.contains("best_params:"));
    assert!(yaml.contains("learning_rate:"));
    assert!(yaml.contains("batch_size:"));
    assert!(yaml.contains("trial_history:"));
    assert!(yaml.contains("trial_number:"));
    assert!(yaml.contains("validation_loss:"));

    // Re-import and verify
    let reimported: OptimizationResult =
        serde_yaml::from_str(&yaml).expect("Failed to deserialize from YAML");

    assert_eq!(
        reimported.best_params.learning_rate,
        result.best_params.learning_rate
    );
    assert_eq!(
        reimported.best_params.batch_size,
        result.best_params.batch_size
    );
    assert_eq!(reimported.trial_history.len(), 3);
    assert_eq!(reimported.trial_history[0].trial_number, 1);
    assert_eq!(reimported.trial_history[2].trial_number, 5);

    println!("✓ YAML export/import successful");
    println!("\nSample YAML output:\n{}", yaml);
}

#[test]
fn test_convergence_tracking() {
    // Mock data showing convergence
    let trial_history = vec![
        TrialResult {
            trial_number: 1,
            learning_rate: 0.001,
            batch_size: 64,
            dropout: 0.2,
            weight_decay: 0.0001,
            validation_loss: 20.0,
            training_time_seconds: 18.0,
        },
        TrialResult {
            trial_number: 2,
            learning_rate: 0.002,
            batch_size: 32,
            dropout: 0.3,
            weight_decay: 0.0002,
            validation_loss: 15.0,
            training_time_seconds: 17.0,
        },
        TrialResult {
            trial_number: 3,
            learning_rate: 0.0015,
            batch_size: 48,
            dropout: 0.25,
            weight_decay: 0.00015,
            validation_loss: 12.0,
            training_time_seconds: 18.5,
        },
        TrialResult {
            trial_number: 4,
            learning_rate: 0.0018,
            batch_size: 56,
            dropout: 0.22,
            weight_decay: 0.00012,
            validation_loss: 10.5,
            training_time_seconds: 19.0,
        },
        TrialResult {
            trial_number: 5,
            learning_rate: 0.0016,
            batch_size: 52,
            dropout: 0.23,
            weight_decay: 0.00013,
            validation_loss: 10.0,
            training_time_seconds: 18.8,
        },
    ];

    // Extract convergence data
    let losses: Vec<f64> = trial_history
        .iter()
        .map(|t| t.validation_loss)
        .collect();

    // Verify convergence (losses generally decrease)
    let best_loss = losses.iter().cloned().fold(f64::INFINITY, f64::min);
    assert_eq!(best_loss, 10.0, "Best loss should be 10.0");

    // Check that at least 3 out of 5 trials improved
    let mut improvements = 0;
    for i in 1..losses.len() {
        if losses[i] < losses[i - 1] {
            improvements += 1;
        }
    }

    assert!(
        improvements >= 3,
        "Expected at least 3 improvements, got {}",
        improvements
    );

    println!("✓ Convergence tracking validated");
    println!("  Trial losses: {:?}", losses);
    println!("  Best loss: {}", best_loss);
    println!("  Improvements: {}/4", improvements);
}

// ============================================================================
// ERROR HANDLING TESTS
// ============================================================================

#[tokio::test]
#[ignore] // Requires error condition setup
async fn test_missing_parquet_file_error() {
    let space = HyperparameterSpace::default();

    let result = optimize_mamba2(
        space,
        "nonexistent_file.parquet",
        5,
        5,
    )
    .await;

    assert!(result.is_err(), "Expected error for missing file");

    let err = result.unwrap_err();
    let err_msg = format!("{:?}", err);
    assert!(
        err_msg.contains("not found") || err_msg.contains("No such file"),
        "Error message should mention file not found: {}",
        err_msg
    );

    println!("✓ Missing file error handled correctly");
}

#[tokio::test]
#[ignore] // Requires CUDA
async fn test_cuda_device_requirement() -> Result<()> {
    // This test verifies CUDA is available for training
    // It will fail gracefully if CUDA is not available

    use candle_core::Device;

    match Device::new_cuda(0) {
        Ok(_device) => {
            println!("✓ CUDA device available for optimization");
            Ok(())
        }
        Err(e) => {
            println!("⚠ CUDA not available: {:?}", e);
            println!("  Hyperparameter optimization requires CUDA");
            // Don't fail the test - just warn
            Ok(())
        }
    }
}