foxhunt/ml/tests/mamba2_early_stopping_test.rs

//! MAMBA-2 Early Stopping Tests (TDD Implementation)
//!
//! This test suite validates the early stopping implementation for MAMBA-2
//! following the TFT reference pattern at `ml/src/trainers/tft.rs:1702-1727`.
//!
//! ## Test Coverage
//!
//! 1. **State Tracking**: Verify best_val_loss, patience_counter initialization
//! 2. **Plateau Detection**: Verify training stops when no improvement
//! 3. **Reset on Improvement**: Verify patience_counter resets when val_loss improves
//! 4. **Minimum Epochs**: Verify min_epochs prevents premature stopping
//! 5. **Hyperopt Integration**: Verify hyperopt trials use early stopping

use candle_core::{Device, Tensor};
use ml::mamba::{Mamba2Config, Mamba2SSM, OptimizerType};

/// Helper function to create a test Mamba2Config with early stopping
fn create_test_config(
    patience: usize,
    min_epochs: usize,
    seq_len: usize,
    batch_size: usize,
) -> Mamba2Config {
    Mamba2Config {
        d_model: 225,
        d_state: 16,
        d_head: 28,
        num_heads: 8,
        expand: 2,
        num_layers: 2,
        dropout: 0.0, // No dropout for deterministic tests
        use_ssd: false,
        use_selective_state: false,
        hardware_aware: false,
        target_latency_us: 5,
        max_seq_len: 120,
        learning_rate: 1e-3,
        weight_decay: 0.0,
        grad_clip: 1.0,
        warmup_steps: 0,
        adam_beta1: 0.9,
        adam_beta2: 0.999,
        adam_epsilon: 1e-8,
        total_decay_steps: 100,
        batch_size,
        seq_len,
        shuffle_batches: false,
        optimizer_type: OptimizerType::Adam,
        sgd_momentum: 0.9,
        sequence_stride: 1,
        norm_eps: 1e-5,
        early_stopping_enabled: true,
        early_stopping_patience: patience,
        early_stopping_min_delta: 1e-4,
        early_stopping_min_epochs: min_epochs,
    }
}

/// Test 1: Early Stopping State Tracking
///
/// Verifies that MAMBA-2 early stopping state is correctly initialized
/// and tracked during training.
#[test]
fn test_mamba2_early_stopping_state() {
    // Create MAMBA-2 model with early stopping config
    let config = create_test_config(5, 5, 10, 4);
    let device = Device::Cpu;
    let model = Mamba2SSM::new(config, &device).expect("Failed to create MAMBA-2 model");

    // Verify early stopping state fields exist
    assert_eq!(
        model.state.best_val_loss,
        f64::INFINITY,
        "best_val_loss should initialize to INFINITY"
    );
    assert_eq!(
        model.state.patience_counter, 0,
        "patience_counter should initialize to 0"
    );
    assert!(
        !model.state.stopped,
        "stopped flag should initialize to false"
    );
    assert_eq!(
        model.state.stopped_at_epoch, None,
        "stopped_at_epoch should be None initially"
    );

    println!("✅ Early stopping state correctly initialized");
}

/// Test 2: check_early_stopping Method Behavior
///
/// Verifies that check_early_stopping correctly implements the TFT pattern
#[test]
fn test_check_early_stopping_method() {
    let config = create_test_config(5, 10, 10, 4);
    let device = Device::Cpu;
    let mut model = Mamba2SSM::new(config, &device).expect("Failed to create MAMBA-2 model");

    // Test 1: Before min_epochs, should NOT stop
    for epoch in 0..10 {
        let should_stop = model.check_early_stopping(epoch, 1.0);
        assert!(
            !should_stop,
            "Should not stop before min_epochs ({})",
            epoch
        );
    }

    // Test 2: After min_epochs, with improvement, should NOT stop
    model.check_early_stopping(10, 0.5); // Improvement
    assert_eq!(
        model.state.patience_counter, 0,
        "Patience should reset on improvement"
    );
    assert_eq!(
        model.state.best_val_loss, 0.5,
        "Best val loss should update"
    );

    // Test 3: After min_epochs, without improvement, should increment patience
    for i in 1..=4 {
        model.check_early_stopping(10 + i, 0.6); // No improvement
        assert_eq!(
            model.state.patience_counter, i,
            "Patience should increment without improvement"
        );
    }

    // Test 4: After patience exhausted, should stop
    let should_stop = model.check_early_stopping(15, 0.6);
    assert!(should_stop, "Should stop after patience exhausted");
    assert!(model.state.stopped, "Stopped flag should be set");
    assert_eq!(
        model.state.stopped_at_epoch,
        Some(15),
        "Stopped epoch should be recorded"
    );

    println!("✅ check_early_stopping method works correctly");
}

/// Test 3: Verify early stopping default config
#[test]
fn test_early_stopping_default_config() {
    let config = Mamba2Config::default();

    assert!(
        config.early_stopping_enabled,
        "Early stopping should be enabled by default"
    );
    assert_eq!(
        config.early_stopping_patience, 20,
        "Default patience should be 20"
    );
    assert_eq!(
        config.early_stopping_min_delta, 1e-4,
        "Default min_delta should be 1e-4"
    );
    assert_eq!(
        config.early_stopping_min_epochs, 20,
        "Default min_epochs should be 20"
    );

    println!("✅ Default early stopping config validated");
}

/// Test 4: Integration test - verify early stopping actually stops training
///
/// This test creates a small training scenario and verifies that training
/// stops early when validation loss plateaus.
#[tokio::test]
async fn test_early_stopping_integration() {
    let seq_len = 10;
    let d_model = 225;
    let batch_size = 4;

    // Create minimal training data
    let mut train_data = Vec::new();
    for _ in 0..20 {
        let input = Tensor::ones((1, seq_len, d_model), candle_core::DType::F64, &Device::Cpu)
            .expect("Failed to create input");
        let target = Tensor::ones((1, 1, 1), candle_core::DType::F64, &Device::Cpu)
            .expect("Failed to create target");

        train_data.push((input, target));
    }

    let val_data = train_data[15..20].to_vec();

    // Create config with VERY short patience for fast testing
    let config = create_test_config(3, 2, seq_len, batch_size);
    let mut model = Mamba2SSM::new(config, &Device::Cpu).expect("Failed to create MAMBA-2 model");

    // Train with early stopping
    let max_epochs = 50;
    let history = model
        .train(&train_data, &val_data, max_epochs, None)
        .await
        .expect("Training failed");

    // Verify training stopped early
    assert!(
        history.len() < max_epochs,
        "Training should stop early, but ran {} epochs (max: {})",
        history.len(),
        max_epochs
    );

    println!(
        "✅ Early stopping integration test passed: stopped at {} epochs (max: {})",
        history.len(),
        max_epochs
    );
}