foxhunt/crates/ml/tests/mamba2_early_stopping_test.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! 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

// candle eliminated — test uses native cudarc StreamTensor APIs
use std::sync::OnceLock;
use ml::mamba::{Mamba2Config, Mamba2SSM, OptimizerType};
use ml_core::cuda_autograd::stream_ops::StreamTensor;
use ml_core::device::MlDevice;
use tracing::info;

static SHARED_CUDA: OnceLock<MlDevice> = OnceLock::new();

fn cuda_device() -> MlDevice {
    SHARED_CUDA.get_or_init(|| MlDevice::cuda(0).expect("CUDA required")).clone()
}

/// 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: 54,
        d_state: 16,
        d_head: 7,
        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,
        spsa_epsilon: 0.01,
    }
}

/// 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 = cuda_device();
    let stream = device.cuda_stream().expect("CUDA stream required");
    let model = Mamba2SSM::new(config, stream).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"
    );

    info!("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 = cuda_device();
    let stream = device.cuda_stream().expect("CUDA stream required");
    let mut model = Mamba2SSM::new(config, stream).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"
    );

    info!("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"
    );

    info!("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 = 54;
    let batch_size = 4;

    let device = cuda_device();
    let stream = device.cuda_stream().expect("CUDA stream required");

    // Create minimal training data as (StreamTensor, StreamTensor) pairs
    let mut train_data = Vec::new();
    for _ in 0..20 {
        let input_data = vec![1.0_f32; 1 * seq_len * d_model];
        let target_data = vec![1.0_f32; 1 * 1 * 1];

        let input = StreamTensor::from_vec(input_data, &[1, seq_len, d_model], stream)
            .expect("Failed to create input");
        let target = StreamTensor::from_vec(target_data, &[1, 1, 1], stream)
            .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, stream).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
    );

    info!(
        stopped_at = history.len(),
        max_epochs,
        "Early stopping integration test passed"
    );
}