foxhunt/ml/tests/mamba_test.rs

#![allow(unused_crate_dependencies)]
use candle_core::{DType, Device, Tensor};
use ml::mamba::{Mamba2Config, Mamba2State, SelectiveStateSpace, StateImportance};
use tokio;

/// Helper function to create test Mamba2Config with reasonable defaults
fn create_test_config(d_model: usize, d_state: usize, num_layers: usize) -> Mamba2Config {
    Mamba2Config {
        d_model,
        d_state,
        d_head: d_state,
        num_heads: 4,
        expand: 2,
        num_layers,
        dropout: 0.1,
        use_ssd: true,
        use_selective_state: true,
        hardware_aware: false, // Disable for tests
        target_latency_us: 100,
        max_seq_len: 1024,
        learning_rate: 1e-4,
        weight_decay: 1e-5,
        grad_clip: 1.0,
        warmup_steps: 100,
        batch_size: 1,
        seq_len: 256,
    }
}

#[tokio::test]
async fn test_mamba2_config_creation() {
    let config = create_test_config(512, 64, 6);
    assert_eq!(config.d_model, 512);
    assert_eq!(config.d_state, 64);
    assert_eq!(config.num_layers, 6);
    assert_eq!(config.expand, 2);
    assert!(config.dropout >= 0.0 && config.dropout <= 1.0);
}

#[tokio::test]
async fn test_mamba2_state_creation() {
    let config = create_test_config(256, 32, 4);
    let state_result = Mamba2State::zeros(&config);

    assert!(state_result.is_ok(), "State creation should succeed");
    let state = state_result.unwrap();

    assert_eq!(state.hidden_states.len(), config.num_layers);
    assert_eq!(state.ssm_states.len(), config.num_layers);
    assert!(!state.selective_state.is_empty());
}

#[tokio::test]
async fn test_selective_state_creation() {
    let config = create_test_config(128, 16, 2);
    let selective_state_result = SelectiveStateSpace::new(&config);

    assert!(
        selective_state_result.is_ok(),
        "Selective state creation should succeed"
    );
    let selective_state = selective_state_result.unwrap();

    // Verify initialization
    assert_eq!(
        selective_state.importance_tracker.len(),
        config.d_model * config.expand
    );
    assert_eq!(selective_state.active_indices.len(), 0); // Initially empty
}

#[tokio::test]
async fn test_selective_state_importance_scoring() {
    let config = create_test_config(128, 16, 2);
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut test_state = Mamba2State::zeros(&config).unwrap();

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

    // Update importance scores
    let result = selective_state.update_importance_scores(&input, &mut test_state);
    if let Err(e) = &result {
        eprintln!("Error in update_importance_scores: {:?}", e);
    }
    assert!(result.is_ok(), "Importance score update should succeed");

    // Verify that importance tracking is working - active_indices should be populated
    assert!(
        selective_state.active_indices.len() > 0,
        "Should have some active indices"
    );
}

#[tokio::test]
async fn test_selective_state_compression() {
    let config = create_test_config(128, 16, 2);
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut test_state = Mamba2State::zeros(&config).unwrap();

    // Compress a state component
    let result = selective_state.compress_state_component(0, &mut test_state);
    assert!(result.is_ok(), "State compression should succeed");

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

#[tokio::test]
async fn test_selective_state_decompression() {
    let config = create_test_config(128, 16, 2);
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut test_state = Mamba2State::zeros(&config).unwrap();

    // First compress
    let _ = selective_state.compress_state_component(0, &mut test_state);

    // Then decompress
    let result = selective_state.decompress_state_component(0, &mut test_state);
    assert!(result.is_ok(), "State decompression should succeed");
}

#[tokio::test]
async fn test_mamba2_state_transitions() {
    let config = create_test_config(64, 8, 2);
    let state = Mamba2State::zeros(&config).unwrap();

    // Test state initialization
    assert_eq!(state.hidden_states.len(), config.num_layers);
    assert_eq!(state.ssm_states.len(), config.num_layers);
    assert!(!state.selective_state.is_empty());

    // Test state structure
    assert!(state.compression_indices.is_empty());
    assert!(state.metrics.is_empty());
}

#[tokio::test]
async fn test_state_importance_new() {
    let importance = StateImportance::new();

    assert_eq!(importance.score, 0.0);
    assert_eq!(importance.usage_count, 0);
    assert_eq!(importance.last_access, 0);
    assert_eq!(importance.moving_average, 0.0);
    assert_eq!(importance.variance, 0.0);
}

#[tokio::test]
async fn test_state_importance_update() {
    let mut importance = StateImportance::new();

    importance.update(1.0, 1, 0.9);
    assert!(importance.score > 0.0);
    assert_eq!(importance.usage_count, 1);
    assert_eq!(importance.last_access, 1);

    importance.update(0.5, 2, 0.9);
    assert_eq!(importance.usage_count, 2);
    assert_eq!(importance.last_access, 2);
}

#[tokio::test]
async fn test_state_importance_effective_importance() {
    let mut importance = StateImportance::new();

    importance.update(1.0, 1, 0.9);
    let effective = importance.effective_importance();

    assert!(effective > 0.0);
    assert!(effective <= 1.0);
}

#[tokio::test]
async fn test_tensor_creation() {
    let device = Device::Cpu;

    // Test various tensor shapes
    let t1 = Tensor::zeros(&[1, 128], DType::F32, &device);
    assert!(t1.is_ok());

    let t2 = Tensor::randn(0.0, 1.0, &[1, 10, 256], &device);
    assert!(t2.is_ok());
}

#[tokio::test]
async fn test_config_validation() {
    let config = create_test_config(256, 32, 4);

    // Verify reasonable defaults
    assert!(config.d_model > 0);
    assert!(config.d_state > 0);
    assert!(config.num_layers > 0);
    assert!(config.expand > 0);
    assert!(config.dropout >= 0.0 && config.dropout <= 1.0);
    assert!(config.learning_rate > 0.0);
    assert!(config.weight_decay >= 0.0);
}

#[tokio::test]
async fn test_multiple_state_operations() {
    let config = create_test_config(64, 8, 2);
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut test_state = Mamba2State::zeros(&config).unwrap();

    // Perform multiple operations in sequence
    let device = Device::Cpu;
    let input = Tensor::randn(0.0, 1.0, &[1, 10, 64], &device).unwrap();

    // Update importance
    let _ = selective_state.update_importance_scores(&input, &mut test_state);

    // Compress some states
    if test_state.selective_state.len() > 2 {
        let _ = selective_state.compress_state_component(0, &mut test_state);
        let _ = selective_state.compress_state_component(1, &mut test_state);
    }

    // Verify operations completed
    assert!(selective_state.active_indices.len() > 0);
}

// Basic integration-style test
#[tokio::test]
async fn test_selective_state_full_workflow() {
    let config = create_test_config(128, 16, 2);
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut state = Mamba2State::zeros(&config).unwrap();

    let device = Device::Cpu;

    // Simulate a few timesteps
    for _i in 0..5 {
        let input = Tensor::randn(0.0, 1.0, &[1, 10, 128], &device).unwrap();

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

        // Active indices should be maintained
        assert!(selective_state.active_indices.len() > 0);
    }
}