foxhunt/ml/tests/tft_hyperopt_edge_cases.rs

//! TFT-Specific Edge Case Tests for Hyperparameter Optimization
//!
//! This test suite covers TFT-specific edge cases:
//! 1. Attention head constraints (num_heads must divide hidden_size)
//! 2. Discrete parameter quantization
//! 3. Quantile loss edge cases
//! 4. INT8/QAT configuration edge cases
//!
//! Purpose: Ensure TFT adapter handles architectural constraints robustly

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

// ============================================================================
// ATTENTION HEAD CONSTRAINTS
// ============================================================================

#[test]
fn test_num_heads_divides_hidden_size() {
    // Valid combinations: (128, 4), (256, 8), (512, 16)
    let valid_combos = vec![(128, 4), (256, 8), (512, 16)];

    for (hidden_size, num_heads) in valid_combos {
        assert_eq!(
            hidden_size % num_heads,
            0,
            "hidden_size {} must be divisible by num_heads {}",
            hidden_size,
            num_heads
        );
    }
}

#[test]
fn test_invalid_num_heads_returns_penalty() {
    // This test verifies that invalid configurations are caught
    // In practice, TFTParams ensures valid discrete combinations
    let params = TFTParams {
        learning_rate: 1e-4,
        batch_size: 64,
        hidden_size: 128,
        num_heads: 5, // Invalid: 128 % 5 != 0
        dropout: 0.1,
    };

    // Verify parameter space enforces valid combinations
    let continuous = params.to_continuous();
    let recovered = TFTParams::from_continuous(&continuous)
        .expect("Parameter recovery should succeed");

    // Recovered params should have valid num_heads (quantized to 4, 8, or 16)
    assert!(
        recovered.hidden_size % recovered.num_heads == 0,
        "Recovered params should have valid num_heads: {} % {} != 0",
        recovered.hidden_size,
        recovered.num_heads
    );
}

// ============================================================================
// DISCRETE PARAMETER QUANTIZATION
// ============================================================================

#[test]
fn test_hidden_size_quantization() {
    // Test all valid hidden_size values
    let test_cases = vec![
        (0.0, 128),  // Index 0 -> 128
        (1.0, 256),  // Index 1 -> 256
        (2.0, 512),  // Index 2 -> 512
        (0.3, 128),  // Rounds to 0 -> 128
        (1.7, 512),  // Rounds to 2 -> 512
    ];

    for (index, expected_size) in test_cases {
        let continuous = vec![
            (-4.6_f64).ln(), // learning_rate
            64.0,            // batch_size
            index,           // hidden_size_index
            1.0,             // num_heads_index
            0.1,             // dropout
        ];

        let params = TFTParams::from_continuous(&continuous)
            .expect("Should create params from continuous");

        assert_eq!(
            params.hidden_size, expected_size,
            "Index {} should map to hidden_size {}",
            index, expected_size
        );
    }
}

#[test]
fn test_num_heads_quantization() {
    // Test all valid num_heads values
    let test_cases = vec![
        (0.0, 4),   // Index 0 -> 4
        (1.0, 8),   // Index 1 -> 8
        (2.0, 16),  // Index 2 -> 16
        (0.4, 4),   // Rounds to 0 -> 4
        (1.6, 16),  // Rounds to 2 -> 16
    ];

    for (index, expected_heads) in test_cases {
        let continuous = vec![
            (-4.6_f64).ln(), // learning_rate
            64.0,            // batch_size
            1.0,             // hidden_size_index
            index,           // num_heads_index
            0.1,             // dropout
        ];

        let params = TFTParams::from_continuous(&continuous)
            .expect("Should create params from continuous");

        assert_eq!(
            params.num_heads, expected_heads,
            "Index {} should map to num_heads {}",
            index, expected_heads
        );
    }
}

#[test]
fn test_discrete_roundtrip() {
    // Test roundtrip conversion preserves discrete values
    for hidden_size in &[128, 256, 512] {
        for num_heads in &[4, 8, 16] {
            let params = TFTParams {
                learning_rate: 1e-4,
                batch_size: 64,
                hidden_size: *hidden_size,
                num_heads: *num_heads,
                dropout: 0.1,
            };

            let continuous = params.to_continuous();
            let recovered = TFTParams::from_continuous(&continuous)
                .expect("Roundtrip should succeed");

            assert_eq!(
                recovered.hidden_size, *hidden_size,
                "Hidden size should be preserved"
            );
            assert_eq!(
                recovered.num_heads, *num_heads,
                "Num heads should be preserved"
            );
        }
    }
}

// ============================================================================
// PARAMETER BOUNDS
// ============================================================================

#[test]
fn test_tft_params_bounds() {
    let bounds = TFTParams::continuous_bounds();

    assert_eq!(bounds.len(), 5, "TFT should have 5 parameters");

    // Learning rate: log-scale [1e-5, 1e-3]
    let lr_min = bounds[0].0.exp();
    let lr_max = bounds[0].1.exp();
    assert!((lr_min - 1e-5).abs() < 1e-10);
    assert!((lr_max - 1e-3).abs() < 1e-10);

    // Batch size: [16, 128]
    assert_eq!(bounds[1], (16.0, 128.0));

    // Hidden size index: [0, 2]
    assert_eq!(bounds[2], (0.0, 2.0));

    // Num heads index: [0, 2]
    assert_eq!(bounds[3], (0.0, 2.0));

    // Dropout: [0.0, 0.3]
    assert_eq!(bounds[4], (0.0, 0.3));
}

#[test]
fn test_param_clamping() {
    // Test extreme values are clamped properly
    let extreme_continuous = vec![
        1000.0,  // learning_rate (should clamp)
        10000.0, // batch_size (should clamp to 128)
        100.0,   // hidden_size_index (should clamp to 2)
        100.0,   // num_heads_index (should clamp to 2)
        10.0,    // dropout (should clamp to 0.3)
    ];

    let params = TFTParams::from_continuous(&extreme_continuous)
        .expect("Should handle extreme values");

    assert!(params.learning_rate < 1.0, "LR should be reasonable");
    assert!(params.batch_size <= 128, "Batch size should be clamped");
    assert!(params.hidden_size <= 512, "Hidden size should be clamped");
    assert!(params.num_heads <= 16, "Num heads should be clamped");
    assert!(params.dropout <= 0.3, "Dropout should be clamped");
}

// ============================================================================
// PARAMETER SPACE VALIDATION
// ============================================================================

#[test]
fn test_param_names() {
    let names = TFTParams::param_names();

    assert_eq!(names.len(), 5);
    assert_eq!(names[0], "learning_rate");
    assert_eq!(names[1], "batch_size");
    assert_eq!(names[2], "hidden_size");
    assert_eq!(names[3], "num_heads");
    assert_eq!(names[4], "dropout");
}

#[test]
fn test_all_valid_configurations() {
    // Test all valid (hidden_size, num_heads) combinations
    let valid_configs = vec![
        (128, 4),
        (128, 8),
        (256, 4),
        (256, 8),
        (256, 16),
        (512, 4),
        (512, 8),
        (512, 16),
    ];

    for (hidden_size, num_heads) in valid_configs {
        let params = TFTParams {
            learning_rate: 1e-4,
            batch_size: 64,
            hidden_size,
            num_heads,
            dropout: 0.1,
        };

        // Verify divisibility
        assert_eq!(
            hidden_size % num_heads,
            0,
            "Configuration ({}, {}) should be valid",
            hidden_size,
            num_heads
        );

        // Verify roundtrip
        let continuous = params.to_continuous();
        let recovered = TFTParams::from_continuous(&continuous)
            .expect("Valid config should roundtrip");

        assert_eq!(recovered.hidden_size, hidden_size);
        assert_eq!(recovered.num_heads, num_heads);
    }
}

#[test]
fn test_default_params_valid() {
    let params = TFTParams::default();

    // Default params should be valid
    assert_eq!(params.hidden_size % params.num_heads, 0);
    assert!(params.learning_rate > 0.0);
    assert!(params.batch_size > 0);
    assert!(params.dropout >= 0.0 && params.dropout <= 1.0);
}

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

#[test]
fn test_tft_trainer_creation() {
    // Test that TFTTrainer rejects invalid paths
    let result = TFTTrainer::new("nonexistent.parquet", 10);

    assert!(
        result.is_err(),
        "Should error on nonexistent parquet file"
    );

    let err_msg = format!("{:?}", result.unwrap_err());
    assert!(
        err_msg.contains("not found") || err_msg.contains("Config"),
        "Should mention file not found"
    );
}

#[test]
fn test_parameter_space_coverage() {
    // Verify parameter space covers production requirements
    let bounds = TFTParams::continuous_bounds();

    // Sample 10 random points in parameter space
    for _ in 0..10 {
        let continuous: Vec<f64> = bounds
            .iter()
            .map(|(min, max)| (min + max) / 2.0) // Use midpoint
            .collect();

        let params = TFTParams::from_continuous(&continuous)
            .expect("Midpoint should be valid");

        // Verify all params are in valid ranges
        assert!(params.learning_rate > 0.0 && params.learning_rate < 1.0);
        assert!(params.batch_size >= 16 && params.batch_size <= 128);
        assert!(vec![128, 256, 512].contains(&params.hidden_size));
        assert!(vec![4, 8, 16].contains(&params.num_heads));
        assert!(params.dropout >= 0.0 && params.dropout <= 0.3);
        assert_eq!(params.hidden_size % params.num_heads, 0);
    }
}

#[test]
fn test_extreme_learning_rates() {
    // Test very small and very large learning rates
    let small_lr = TFTParams {
        learning_rate: 1e-6,
        ..Default::default()
    };

    let large_lr = TFTParams {
        learning_rate: 1e-2,
        ..Default::default()
    };

    // Both should be valid
    assert!(small_lr.learning_rate > 0.0);
    assert!(large_lr.learning_rate < 1.0);

    // Verify roundtrip
    let small_continuous = small_lr.to_continuous();
    let large_continuous = large_lr.to_continuous();

    assert!(TFTParams::from_continuous(&small_continuous).is_ok());
    assert!(TFTParams::from_continuous(&large_continuous).is_ok());
}

#[test]
fn test_batch_size_boundaries() {
    // Test min and max batch sizes
    let min_batch = TFTParams {
        batch_size: 16,
        ..Default::default()
    };

    let max_batch = TFTParams {
        batch_size: 128,
        ..Default::default()
    };

    // Verify roundtrip
    let min_continuous = min_batch.to_continuous();
    let max_continuous = max_batch.to_continuous();

    let recovered_min = TFTParams::from_continuous(&min_continuous)
        .expect("Min batch size should be valid");
    let recovered_max = TFTParams::from_continuous(&max_continuous)
        .expect("Max batch size should be valid");

    assert_eq!(recovered_min.batch_size, 16);
    assert_eq!(recovered_max.batch_size, 128);
}