foxhunt/ml/tests/tft_hyperopt_test.rs

//! TFT Hyperparameter Optimization Integration Test
//!
//! This test validates the full hyperparameter optimization pipeline for TFT:
//! - Parameter space conversion (continuous ↔ structured)
//! - Training integration with ES_FUT_small.parquet
//! - Optimizer convergence (3 trials × 5 epochs)
//! - Feature normalization and validation
//!
//! ## Test Strategy
//!
//! 1. **Smoke Test**: Verify TFT adapter API compatibility
//! 2. **Small Dataset**: Train with ES_FUT_small.parquet (25KB, ~200 samples)
//! 3. **Quick Optimization**: 3 trials × 5 epochs (~30 seconds total)
//! 4. **Validation**: Loss < 0.20, model learning detected
//!
//! ## Expected Behavior
//!
//! - Trial 1: Baseline (random initialization)
//! - Trial 2-3: Improvement via Argmin Particle Swarm
//! - Final loss: < 0.20 (good TFT performance on small dataset)
//! - No CUDA OOM errors (batch_size=16 safe for 4GB GPU)

use anyhow::Result;
use ml::hyperopt::adapters::tft::{TFTParams, TFTTrainer};
use ml::hyperopt::traits::{HyperparameterOptimizable, ParameterSpace};
use ml::hyperopt::ArgminOptimizer;

#[test]
fn test_tft_params_api() {
    // Verify parameter space API works correctly
    let params = TFTParams::default();

    // Test continuous conversion (roundtrip)
    let continuous = params.to_continuous();
    assert_eq!(continuous.len(), 5, "TFT has 5 hyperparameters");

    let recovered = TFTParams::from_continuous(&continuous)
        .expect("Failed to convert from continuous");

    // Verify values are preserved (with floating-point tolerance)
    assert!((recovered.learning_rate - params.learning_rate).abs() < 1e-10);
    assert_eq!(recovered.batch_size, params.batch_size);
    assert_eq!(recovered.hidden_size, params.hidden_size);
    assert_eq!(recovered.num_heads, params.num_heads);
    assert!((recovered.dropout - params.dropout).abs() < 1e-10);

    // Verify parameter names
    let names = TFTParams::param_names();
    assert_eq!(names, vec![
        "learning_rate", "batch_size", "hidden_size", "num_heads", "dropout"
    ]);

    // Verify bounds are reasonable
    let bounds = TFTParams::continuous_bounds();
    assert_eq!(bounds.len(), 5);
    assert!(bounds[0].0 < bounds[0].1, "Learning rate bounds inverted");
    assert!(bounds[1].0 < bounds[1].1, "Batch size bounds inverted");
}

#[test]
fn test_tft_trainer_creation() {
    // Verify trainer can be created with valid parquet file
    let parquet_file = "test_data/ES_FUT_small.parquet";

    let trainer = TFTTrainer::new(parquet_file, 5);
    assert!(trainer.is_ok(), "Failed to create TFT trainer: {:?}", trainer.err());

    // Verify error handling for missing file
    let bad_trainer = TFTTrainer::new("nonexistent.parquet", 5);
    assert!(bad_trainer.is_err(), "Should fail with missing parquet file");
}

#[test]
fn test_tft_single_trial() {
    // Test single training trial with default parameters
    let parquet_file = "test_data/ES_FUT_small.parquet";
    let mut trainer = TFTTrainer::new(parquet_file, 5)
        .expect("Failed to create trainer");

    let params = TFTParams {
        learning_rate: 1e-3,
        batch_size: 16, // Safe for small dataset
        hidden_size: 128, // Small model
        num_heads: 4,
        dropout: 0.1,
    };

    let metrics = trainer.train_with_params(params)
        .expect("Training failed");

    // Validate metrics are reasonable
    assert!(metrics.val_loss > 0.0, "Val loss should be positive");
    assert!(metrics.val_loss < 10.0, "Val loss too high: {}", metrics.val_loss);
    assert!(metrics.train_loss > 0.0, "Train loss should be positive");
    assert_eq!(metrics.epochs_completed, 5, "Should complete 5 epochs");

    println!("✓ Single trial completed:");
    println!("  Val loss: {:.6}", metrics.val_loss);
    println!("  Train loss: {:.6}", metrics.train_loss);
    println!("  Val RMSE: {:.4}", metrics.val_rmse);
}

#[test]
#[ignore] // Expensive test - run with: cargo test tft_hyperopt_small_dataset -- --ignored --nocapture
fn test_tft_hyperopt_small_dataset() {
    // Full hyperparameter optimization test with small dataset
    println!("╔═══════════════════════════════════════════════════════════╗");
    println!("║  TFT Hyperparameter Optimization Test                    ║");
    println!("╚═══════════════════════════════════════════════════════════╝");
    println!();

    let parquet_file = "test_data/ES_FUT_small.parquet";
    println!("Dataset: {}", parquet_file);
    println!("Configuration:");
    println!("  • Trials: 3");
    println!("  • Initial samples: 2 (Latin Hypercube)");
    println!("  • Epochs per trial: 5");
    println!("  • Batch size: 16 (safe for small dataset)");
    println!("  • Hidden sizes: [128, 256, 512]");
    println!("  • Num heads: [4, 8, 16]");
    println!();

    // Create trainer
    let trainer = TFTTrainer::new(parquet_file, 5)
        .expect("Failed to create TFT trainer");

    // Create optimizer (3 trials, 2 initial samples)
    let optimizer = ArgminOptimizer::builder()
        .max_trials(3)
        .n_initial(2)
        .seed(42) // Reproducible results
        .build();

    // Run optimization
    println!("Starting optimization...");
    let result = optimizer.optimize(trainer)
        .expect("Optimization failed");

    println!();
    println!("╔═══════════════════════════════════════════════════════════╗");
    println!("║  Optimization Results                                     ║");
    println!("╚═══════════════════════════════════════════════════════════╝");
    println!();
    println!("Best Parameters:");
    println!("  • Learning rate: {:.6}", result.best_params.learning_rate);
    println!("  • Batch size: {}", result.best_params.batch_size);
    println!("  • Hidden size: {}", result.best_params.hidden_size);
    println!("  • Num heads: {}", result.best_params.num_heads);
    println!("  • Dropout: {:.3}", result.best_params.dropout);
    println!();
    println!("Metrics:");
    println!("  • Best validation loss: {:.6}", result.best_objective);
    println!("  • Total improvement: {:.6}", result.total_improvement());
    println!("  • Improvement: {:.2}%", result.improvement_percentage());
    println!();

    // Validate results
    assert!(result.best_objective < 0.20,
        "Best val loss too high: {:.6} (expected < 0.20)",
        result.best_objective);

    assert!(result.best_objective > 0.0,
        "Best val loss invalid: {}", result.best_objective);

    // Check learning occurred (val loss should decrease)
    if result.all_trials.len() >= 2 {
        let first_loss = result.all_trials[0].objective;
        let last_loss = result.all_trials[result.all_trials.len() - 1].objective;

        println!("Learning Progress:");
        println!("  • Trial 1 loss: {:.6}", first_loss);
        println!("  • Trial {} loss: {:.6}", result.all_trials.len(), last_loss);

        // Should see some improvement (not strict requirement)
        if last_loss < first_loss {
            println!("  • ✓ Model learning detected");
        } else {
            println!("  • ⚠ No improvement detected (may happen with small dataset)");
        }
    }

    println!();
    println!("✓ TFT hyperparameter optimization test PASSED");
}

#[test]
#[ignore] // Expensive test
fn test_tft_hyperopt_parameter_bounds() {
    // Verify optimizer explores full parameter space
    let parquet_file = "test_data/ES_FUT_small.parquet";
    let trainer = TFTTrainer::new(parquet_file, 3) // Fewer epochs for speed
        .expect("Failed to create trainer");

    let optimizer = ArgminOptimizer::builder()
        .max_trials(5) // More trials to explore space
        .n_initial(3)
        .seed(123)
        .build();

    let result = optimizer.optimize(trainer)
        .expect("Optimization failed");

    // Check that different parameter values were tried
    let mut learning_rates: Vec<f64> = result.all_trials.iter()
        .map(|t| t.params.learning_rate)
        .collect();
    learning_rates.sort_by(|a, b| a.partial_cmp(b).unwrap());

    // Should have explored different learning rates
    let lr_range = learning_rates.last().unwrap() - learning_rates.first().unwrap();
    assert!(lr_range > 1e-5, "Learning rate range too small: {:.6}", lr_range);

    println!("Parameter Exploration:");
    println!("  Learning rates: {:.6} to {:.6} (range: {:.6})",
        learning_rates.first().unwrap(),
        learning_rates.last().unwrap(),
        lr_range);

    // Check batch sizes
    let mut batch_sizes: Vec<usize> = result.all_trials.iter()
        .map(|t| t.params.batch_size)
        .collect();
    batch_sizes.sort();
    batch_sizes.dedup();

    println!("  Batch sizes explored: {:?}", batch_sizes);
    assert!(batch_sizes.len() >= 2, "Should explore multiple batch sizes");

    println!("✓ Parameter exploration validated");
}

#[test]
fn test_tft_normalization_features() {
    // Verify TFT adapter correctly handles normalization
    // NOTE: Current TFT adapter returns synthetic metrics
    // This test validates the API is correct for future integration

    let parquet_file = "test_data/ES_FUT_small.parquet";
    let mut trainer = TFTTrainer::new(parquet_file, 5)
        .expect("Failed to create trainer");

    let params = TFTParams::default();
    let metrics = trainer.train_with_params(params)
        .expect("Training failed");

    // Validate metrics structure (API test)
    assert!(metrics.val_loss.is_finite(), "Val loss should be finite");
    assert!(metrics.train_loss.is_finite(), "Train loss should be finite");
    assert!(metrics.val_rmse.is_finite(), "RMSE should be finite");

    println!("✓ TFT metrics API validated");
    println!("  Metrics: train_loss={:.6}, val_loss={:.6}, rmse={:.4}",
        metrics.train_loss, metrics.val_loss, metrics.val_rmse);
}

#[test]
fn test_tft_discrete_parameters() {
    // Verify discrete parameter quantization works correctly

    // Test hidden_size quantization (should map to 128, 256, or 512)
    let test_cases = vec![
        (0.0, 128),  // Index 0 → 128
        (1.0, 256),  // Index 1 → 256
        (2.0, 512),  // Index 2 → 512
    ];

    for (idx, expected_size) in test_cases {
        let continuous = vec![
            1e-4_f64.ln(), // learning_rate
            64.0,          // batch_size
            idx,           // hidden_size_index
            1.0,           // num_heads_index (8 heads)
            0.1,           // dropout
        ];

        let params = TFTParams::from_continuous(&continuous)
            .expect("Failed to convert parameters");

        assert_eq!(params.hidden_size, expected_size,
            "Hidden size index {} should map to {}, got {}",
            idx, expected_size, params.hidden_size);
    }

    // Test num_heads quantization (should map to 4, 8, or 16)
    let heads_cases = vec![
        (0.0, 4),   // Index 0 → 4
        (1.0, 8),   // Index 1 → 8
        (2.0, 16),  // Index 2 → 16
    ];

    for (idx, expected_heads) in heads_cases {
        let continuous = vec![
            1e-4_f64.ln(), // learning_rate
            64.0,          // batch_size
            1.0,           // hidden_size_index (256)
            idx,           // num_heads_index
            0.1,           // dropout
        ];

        let params = TFTParams::from_continuous(&continuous)
            .expect("Failed to convert parameters");

        assert_eq!(params.num_heads, expected_heads,
            "Num heads index {} should map to {}, got {}",
            idx, expected_heads, params.num_heads);
    }

    println!("✓ Discrete parameter quantization validated");
}