foxhunt/ml/tests/mamba_test.rs

use candle_core::{DType, Device, Tensor};
use ml::mamba::selective_state::StateImportance;
use ml::mamba::{Mamba2Config, Mamba2SSM, Mamba2State, SSDLayer, SelectiveStateSpace};
use proptest::prelude::*;
use std::collections::{BTreeMap, HashMap};
use tokio;

/// Mock MAMBA-2 SSM for testing
#[derive(Debug, Clone)]
pub struct MockMamba2SSM {
    pub config: Mamba2Config,
    pub state: Mamba2State,
    pub forward_calls: usize,
    pub training_calls: usize,
}

impl MockMamba2SSM {
    pub fn new(config: Mamba2Config) -> Self {
        // Create state with zeros - handle error by defaulting to a simple state
        let state = Mamba2State::zeros(&config).unwrap_or_else(|_| {
            // Fallback state if creation fails
            Mamba2State {
                hidden_states: Vec::new(),
                selective_state: vec![0.0; config.d_model * config.expand],
                ssm_states: Vec::new(),
                compression_indices: Vec::new(),
                metrics: HashMap::new(),
                last_update: std::time::Instant::now(),
            }
        });
        Self {
            config: config.clone(),
            state,
            forward_calls: 0,
            training_calls: 0,
        }
    }

    pub async fn forward(
        &mut self,
        input: &Tensor,
    ) -> Result<Tensor, Box<dyn std::error::Error + Send + Sync>> {
        self.forward_calls += 1;
        // Mock forward pass - return tensor with same shape
        Ok(input.clone())
    }

    pub async fn train_step(
        &mut self,
        batch: &[Tensor],
    ) -> Result<f64, Box<dyn std::error::Error + Send + Sync>> {
        self.training_calls += 1;
        // Mock training - return decreasing loss
        Ok(1.0 / (self.training_calls as f64 + 1.0))
    }
}

/// 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_ssm_creation() {
    let config = create_test_config(512, 64, 6);
    let model = MockMamba2SSM::new(config.clone());
    assert_eq!(model.config.d_model, 512);
    assert_eq!(model.config.d_state, 64);
    assert_eq!(model.forward_calls, 0);
    assert_eq!(model.training_calls, 0);
}

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

    let mut model = MockMamba2SSM::new(config);
    let device = Device::Cpu;
    let input = Tensor::randn(0.0, 1.0, &[1, 10, 256], &device).unwrap();

    let result = model.forward(&input).await;
    assert!(result.is_ok());
    assert_eq!(model.forward_calls, 1);
}

#[tokio::test]
async fn test_mamba2_linear_attention_complexity() {
    // Test O(n) complexity of linear attention vs O(n²) traditional attention
    let config = create_test_config(128, 16, 2);

    let mut model = MockMamba2SSM::new(config);
    let device = Device::Cpu;

    // Test with different sequence lengths
    let short_seq = Tensor::randn(0.0, 1.0, &[1, 50, 128], &device).unwrap();
    let long_seq = Tensor::randn(0.0, 1.0, &[1, 500, 128], &device).unwrap();

    let start = std::time::Instant::now();
    let _ = model.forward(&short_seq).await;
    let short_duration = start.elapsed();

    let start = std::time::Instant::now();
    let _ = model.forward(&long_seq).await;
    let long_duration = start.elapsed();

    // Linear attention should scale approximately linearly
    let ratio = long_duration.as_nanos() as f64 / short_duration.as_nanos() as f64;
    assert!(
        ratio < 15.0,
        "Attention complexity should be approximately linear, got ratio: {}",
        ratio
    );
}

#[tokio::test]
async fn test_ssd_layer_creation() {
    let ssd_layer = SSDLayer {
        qkv_projection: Default::default(), // Mock linear layer
        attention_cache: HashMap::new(),
        layer_norm: Default::default(),
        output_projection: Default::default(),
        d_model: 256,
        d_state: 32,
        use_cache: true,
    };

    assert_eq!(ssd_layer.d_model, 256);
    assert_eq!(ssd_layer.d_state, 32);
    assert!(ssd_layer.use_cache);
    assert!(ssd_layer.attention_cache.is_empty());
}

#[tokio::test]
async fn test_ssd_layer_caching() {
    let mut ssd_layer = SSDLayer {
        qkv_projection: Default::default(),
        attention_cache: HashMap::new(),
        layer_norm: Default::default(),
        output_projection: Default::default(),
        d_model: 256,
        d_state: 32,
        use_cache: true,
    };

    // Simulate adding cache entries
    let device = Device::Cpu;
    let cache_tensor = Tensor::randn(0.0, 1.0, &[1, 32, 256], &device).unwrap();

    // Mock cache key generation
    let cache_key = "layer_0_step_1".to_string();
    ssd_layer
        .attention_cache
        .insert(cache_key.clone(), cache_tensor);

    assert_eq!(ssd_layer.attention_cache.len(), 1);
    assert!(ssd_layer.attention_cache.contains_key(&cache_key));
}

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

    // Test that selective state is properly initialized
    assert!(selective_state.get_memory_usage() >= 0);
    // Selective state should have reasonable initial state
}

#[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();

    // Create a test state to update importance scores
    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);
    assert!(result.is_ok());

    // Verify that importance tracking is working
    assert!(selective_state.active_indices.len() > 0);
}

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

    // Test compression functionality
    let device = Device::Cpu;
    let test_state = Tensor::randn(0.0, 1.0, &[1, 128], &device).unwrap();

    // Compress and store state
    let result = selective_state.compress_state(0, &test_state);
    assert!(result.is_ok());

    // Verify compression occurred
    assert!(selective_state.compressed_states.len() > 0);
    assert!(selective_state.get_memory_usage() > 0);
}

#[tokio::test]
async fn test_mamba2_training_step() {
    let config = create_test_config(128, 16, 2);

    let mut model = MockMamba2SSM::new(config);
    let device = Device::Cpu;

    let batch = vec![
        Tensor::randn(0.0, 1.0, &[1, 10, 128], &device).unwrap(),
        Tensor::randn(0.0, 1.0, &[1, 10, 128], &device).unwrap(),
    ];

    let loss = model.train_step(&batch).await.unwrap();
    assert!(loss > 0.0);
    assert!(loss <= 1.0);
    assert_eq!(model.training_calls, 1);

    // Second training step should have lower loss
    let loss2 = model.train_step(&batch).await.unwrap();
    assert!(loss2 < loss);
    assert_eq!(model.training_calls, 2);
}

#[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!(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_mamba2_discretization_methods() {
    let config = create_test_config(32, 4, 1);

    // Test SSM discretization
    let mut model1 = MockMamba2SSM::new(config.clone());
    let device = Device::Cpu;
    let input = Tensor::randn(0.0, 1.0, &[1, 5, 32], &device).unwrap();

    let result1 = model1.forward(&input).await;
    assert!(result1.is_ok());

    // Test with different configuration
    let config2 = create_test_config(32, 4, 1);
    let mut model2 = MockMamba2SSM::new(config2);

    let result2 = model2.forward(&input).await;
    assert!(result2.is_ok());
}

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

    let mut model = MockMamba2SSM::new(config);
    let device = Device::Cpu;

    // Test memory usage with long sequences
    let long_input = Tensor::randn(0.0, 1.0, &[1, 1000, 256], &device).unwrap();
    let result = model.forward(&long_input).await;

    assert!(result.is_ok());
    // In a real implementation, we would check that memory usage stays reasonable
    // For mock, we just verify the operation completes
}

#[tokio::test]
async fn test_mamba2_hardware_optimization() {
    let mut config = create_test_config(128, 16, 2);

    let mut fast_model = MockMamba2SSM::new(config.clone());
    config.hardware_aware = false; // Disable hardware optimization
    let mut slow_model = MockMamba2SSM::new(config);

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

    // Both should work, but hardware-aware path should be preferred for performance
    let fast_result = fast_model.forward(&input).await;
    let slow_result = slow_model.forward(&input).await;

    assert!(fast_result.is_ok());
    assert!(slow_result.is_ok());
}

// Property-based tests using proptest
proptest! {
    #[test]
    fn test_mamba2_config_properties(
        d_model in 32..512_u32,
        d_state in 8..64_u32,
        num_layers in 1..8_usize,
    ) {
        let config = create_test_config(d_model as usize, d_state as usize, num_layers);

        let model = MockMamba2SSM::new(config.clone());
        prop_assert_eq!(model.config.d_model, d_model as usize);
        prop_assert_eq!(model.config.d_state, d_state as usize);
        prop_assert_eq!(model.config.num_layers, num_layers);
        prop_assert!(model.config.expand > 0);
        prop_assert!(model.config.dropout >= 0.0 && model.config.dropout <= 1.0);
    }
}