foxhunt/ml/tests/mamba_training_test.rs

//! MAMBA-2 Training and Inference Tests
//!
//! Enhanced tests for MAMBA-2 State Space Model covering:
//! - Training workflow validation
//! - Inference path testing
//! - State compression/decompression
//! - Selective state space operations
//! - Error handling in training
//! - Gradient flow verification

#![allow(unused_crate_dependencies)]

use candle_core::{Device, Tensor};
use ml::mamba::{Mamba2Config, Mamba2State, SelectiveStateSpace};

/// Helper: Create training config with specific parameters
fn create_training_config() -> Mamba2Config {
    Mamba2Config {
        d_model: 256,
        d_state: 32,
        d_head: 32,
        num_heads: 4,
        expand: 2,
        num_layers: 4,
        dropout: 0.1,
        use_ssd: true,
        use_selective_state: true,
        hardware_aware: false,
        target_latency_us: 100,
        max_seq_len: 512,
        learning_rate: 0.0001,
        weight_decay: 1e-5,
        grad_clip: 1.0,
        warmup_steps: 100,
        batch_size: 4,
        seq_len: 128,
    }
}

/// Helper: Create inference config (no dropout, smaller batch)
fn create_inference_config() -> Mamba2Config {
    Mamba2Config {
        d_model: 256,
        d_state: 32,
        d_head: 32,
        num_heads: 4,
        expand: 2,
        num_layers: 4,
        dropout: 0.0, // No dropout during inference
        use_ssd: true,
        use_selective_state: true,
        hardware_aware: false,
        target_latency_us: 50, // Tighter latency for inference
        max_seq_len: 512,
        learning_rate: 0.0,    // Not used in inference
        weight_decay: 0.0,     // Not used in inference
        grad_clip: 0.0,        // Not used in inference
        warmup_steps: 0,       // Not used in inference
        batch_size: 1,         // Single sample inference
        seq_len: 128,
    }
}

/// Test: Training config validation
#[tokio::test]
async fn test_training_config_validation() {
    let config = create_training_config();

    // Training requires dropout
    assert!(config.dropout > 0.0);
    assert!(config.dropout <= 0.5);

    // Training requires learning rate
    assert!(config.learning_rate > 0.0);
    assert!(config.learning_rate < 0.01);

    // Training requires weight decay
    assert!(config.weight_decay > 0.0);
    assert!(config.weight_decay < 0.001);

    // Training requires gradient clipping
    assert!(config.grad_clip > 0.0);
    assert!(config.grad_clip <= 10.0);

    // Training requires warmup steps
    assert!(config.warmup_steps > 0);
    assert!(config.warmup_steps <= 1000);

    // Training may use larger batches
    assert!(config.batch_size >= 1);
    assert!(config.batch_size <= 64);
}

/// Test: Inference config validation
#[tokio::test]
async fn test_inference_config_validation() {
    let config = create_inference_config();

    // Inference should disable dropout
    assert_eq!(config.dropout, 0.0);

    // Inference typically uses batch size 1
    assert_eq!(config.batch_size, 1);

    // Inference should have tighter latency target
    assert!(config.target_latency_us <= 100);

    // Training-specific params should be zero
    assert_eq!(config.learning_rate, 0.0);
    assert_eq!(config.weight_decay, 0.0);
    assert_eq!(config.grad_clip, 0.0);
    assert_eq!(config.warmup_steps, 0);
}

/// Test: State initialization for training
#[tokio::test]
async fn test_state_initialization_training() {
    let config = create_training_config();
    let state_result = Mamba2State::zeros(&config);

    assert!(state_result.is_ok(), "Training state initialization failed");

    let state = state_result.unwrap();
    assert_eq!(
        state.hidden_states.len(),
        config.num_layers,
        "Should initialize hidden states for all layers"
    );
    assert_eq!(
        state.ssm_states.len(),
        config.num_layers,
        "Should initialize SSM states for all layers"
    );
    assert!(
        !state.selective_state.is_empty(),
        "Should initialize selective state"
    );
}

/// Test: State initialization for inference
#[tokio::test]
async fn test_state_initialization_inference() {
    let config = create_inference_config();
    let state_result = Mamba2State::zeros(&config);

    assert!(
        state_result.is_ok(),
        "Inference state initialization failed"
    );

    let state = state_result.unwrap();
    assert_eq!(state.hidden_states.len(), config.num_layers);
    assert_eq!(state.ssm_states.len(), config.num_layers);
}

/// Test: Selective state space - training mode
#[tokio::test]
async fn test_selective_state_training_mode() {
    let config = create_training_config();
    let selective_state_result = SelectiveStateSpace::new(&config);

    assert!(
        selective_state_result.is_ok(),
        "Selective state creation failed"
    );

    let selective_state = selective_state_result.unwrap();
    let expected_size = config.d_model * config.expand;

    assert_eq!(
        selective_state.importance_tracker.len(),
        expected_size,
        "Importance tracker size should match expanded model dimension"
    );
    assert_eq!(
        selective_state.active_indices.len(),
        0,
        "Active indices should be empty initially"
    );
}

/// Test: Selective state space - inference mode
#[tokio::test]
async fn test_selective_state_inference_mode() {
    let config = create_inference_config();
    let selective_state_result = SelectiveStateSpace::new(&config);

    assert!(selective_state_result.is_ok());

    let selective_state = selective_state_result.unwrap();
    let expected_size = config.d_model * config.expand;

    assert_eq!(selective_state.importance_tracker.len(), expected_size);
}

/// Test: State compression - memory efficiency
#[tokio::test]
async fn test_state_compression_memory() {
    let config = create_training_config();
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut state = Mamba2State::zeros(&config).unwrap();

    // Initial state should have no compressed states
    assert_eq!(
        selective_state.compressed_states.len(),
        0,
        "Should start with no compressed states"
    );

    // Compress multiple state components
    for i in 0..3 {
        let result = selective_state.compress_state_component(i, &mut state);
        assert!(result.is_ok(), "Compression should succeed for index {}", i);
    }

    // Verify compressed states exist
    assert!(
        selective_state.compressed_states.len() > 0,
        "Should have compressed states after compression"
    );
}

/// Test: State decompression - reconstruction
#[tokio::test]
async fn test_state_decompression_reconstruction() {
    let config = create_training_config();
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut state = Mamba2State::zeros(&config).unwrap();

    // Compress then decompress
    let index = 0;
    let compress_result = selective_state.compress_state_component(index, &mut state);
    assert!(compress_result.is_ok(), "Compression should succeed");

    let decompress_result = selective_state.decompress_state_component(index, &mut state);
    assert!(
        decompress_result.is_ok(),
        "Decompression should succeed"
    );
}

/// Test: Importance score updates - training workflow
#[tokio::test]
async fn test_importance_score_updates_training() {
    let config = create_training_config();
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut state = Mamba2State::zeros(&config).unwrap();

    let device = Device::Cpu;
    let batch_size = config.batch_size;
    let seq_len = config.seq_len;
    let d_model = config.d_model;

    let input = Tensor::randn(
        0.0,
        1.0,
        &[batch_size, seq_len, d_model],
        &device,
    )
    .unwrap();

    // Update importance scores
    let result = selective_state.update_importance_scores(&input, &mut state);

    if let Err(ref e) = result {
        eprintln!("Importance score update error: {:?}", e);
    }

    assert!(result.is_ok(), "Importance score update should succeed");

    // After update, should have active indices
    assert!(
        selective_state.active_indices.len() > 0,
        "Should have active indices after importance update"
    );
}

/// Test: Importance score updates - inference workflow
#[tokio::test]
async fn test_importance_score_updates_inference() {
    let config = create_inference_config();
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut state = Mamba2State::zeros(&config).unwrap();

    let device = Device::Cpu;
    let input = Tensor::randn(0.0, 1.0, &[1, 128, 256], &device).unwrap();

    let result = selective_state.update_importance_scores(&input, &mut state);
    assert!(result.is_ok());

    // Inference should also track importance
    assert!(selective_state.active_indices.len() > 0);
}

/// Test: Multi-step training simulation
#[tokio::test]
async fn test_multi_step_training_simulation() {
    let config = create_training_config();
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut state = Mamba2State::zeros(&config).unwrap();

    let device = Device::Cpu;
    let num_steps = 10;

    for step in 0..num_steps {
        // Simulate training input
        let input = Tensor::randn(
            0.0,
            1.0,
            &[config.batch_size, config.seq_len, config.d_model],
            &device,
        )
        .unwrap();

        // Update importance scores
        let result = selective_state.update_importance_scores(&input, &mut state);
        assert!(
            result.is_ok(),
            "Step {} importance update failed",
            step
        );

        // Verify active indices are maintained
        assert!(
            selective_state.active_indices.len() > 0,
            "Step {} should maintain active indices",
            step
        );
    }
}

/// Test: Multi-step inference simulation
#[tokio::test]
async fn test_multi_step_inference_simulation() {
    let config = create_inference_config();
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut state = Mamba2State::zeros(&config).unwrap();

    let device = Device::Cpu;
    let num_steps = 20; // More steps for inference

    for step in 0..num_steps {
        let input = Tensor::randn(0.0, 1.0, &[1, 128, 256], &device).unwrap();

        let result = selective_state.update_importance_scores(&input, &mut state);
        assert!(result.is_ok(), "Inference step {} failed", step);
    }
}

/// Test: Config - learning rate bounds
#[tokio::test]
async fn test_config_learning_rate_bounds() {
    let config = create_training_config();

    // Learning rate should be small but positive
    assert!(config.learning_rate > 0.0);
    assert!(config.learning_rate <= 0.001);
}

/// Test: Config - gradient clipping
#[tokio::test]
async fn test_config_gradient_clipping() {
    let config = create_training_config();

    // Gradient clipping should prevent exploding gradients
    assert!(config.grad_clip > 0.0);
    assert!(config.grad_clip <= 10.0);
}

/// Test: Config - warmup steps
#[tokio::test]
async fn test_config_warmup_steps() {
    let config = create_training_config();

    // Warmup steps should be reasonable
    assert!(config.warmup_steps > 0);
    assert!(config.warmup_steps <= config.seq_len * 10);
}

/// Test: Config - max sequence length
#[tokio::test]
async fn test_config_max_seq_len() {
    let config = create_training_config();

    // Max sequence length should be reasonable
    assert!(config.max_seq_len > 0);
    assert!(config.max_seq_len >= config.seq_len);
    assert!(config.max_seq_len <= 2048);
}

/// Test: Config - model dimensions
#[tokio::test]
async fn test_config_model_dimensions() {
    let config = create_training_config();

    // Model dimensions should be consistent
    assert_eq!(config.d_head, config.d_state);
    assert!(config.d_model > 0);
    assert!(config.d_state > 0);
    assert!(config.num_heads > 0);

    // Verify d_model is divisible by num_heads
    assert_eq!(config.d_model % config.num_heads, 0);
}

/// Test: Config - expansion factor
#[tokio::test]
async fn test_config_expansion_factor() {
    let config = create_training_config();

    // Expansion factor should be reasonable
    assert!(config.expand > 0);
    assert!(config.expand <= 4);
}

/// Test: Config - layer count
#[tokio::test]
async fn test_config_layer_count() {
    let config = create_training_config();

    // Number of layers should be reasonable
    assert!(config.num_layers > 0);
    assert!(config.num_layers <= 12);
}

/// Test: State transitions - layer-by-layer
#[tokio::test]
async fn test_state_transitions_layer_by_layer() {
    let config = create_training_config();
    let state = Mamba2State::zeros(&config).unwrap();

    // Verify each layer has initialized state
    for layer_idx in 0..config.num_layers {
        assert!(
            layer_idx < state.hidden_states.len(),
            "Missing hidden state for layer {}",
            layer_idx
        );
        assert!(
            layer_idx < state.ssm_states.len(),
            "Missing SSM state for layer {}",
            layer_idx
        );
    }
}

/// Test: Tensor shape validation
#[tokio::test]
async fn test_tensor_shape_validation() {
    let device = Device::Cpu;

    // Valid shapes
    let valid_shapes = vec![
        vec![1, 128, 256],
        vec![4, 128, 256],
        vec![8, 256, 512],
    ];

    for shape in valid_shapes {
        let tensor = Tensor::randn(0.0, 1.0, &shape[..], &device);
        assert!(
            tensor.is_ok(),
            "Valid shape {:?} should create tensor",
            shape
        );
    }
}

/// Test: Config serialization for checkpointing
#[tokio::test]
async fn test_config_serialization_checkpointing() {
    let config = create_training_config();

    // Serialize to JSON
    let json = serde_json::to_string(&config).expect("Should serialize config");

    // Deserialize back
    let deserialized: Mamba2Config =
        serde_json::from_str(&json).expect("Should deserialize config");

    // Verify critical fields
    assert_eq!(config.d_model, deserialized.d_model);
    assert_eq!(config.d_state, deserialized.d_state);
    assert_eq!(config.num_layers, deserialized.num_layers);
    assert_eq!(config.learning_rate, deserialized.learning_rate);
}