foxhunt/ml/tests/test_mamba_comprehensive.rs

use foxhunt_ml::mamba::{Mamba2SSM, Mamba2Config, Mamba2State, SSDLayer, SelectiveStateSpace};
use foxhunt_ml::mamba::selective_state::{StateImportance, ImportanceThreshold};
use foxhunt_core::types::{TradingSignal, ModelPerformance};
use candle_core::{Tensor, Device, DType};
use proptest::prelude::*;
use tokio;
use std::collections::{HashMap, BTreeMap};

/// 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 {
        Self {
            config: config.clone(),
            state: Mamba2State::new(&config),
            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))
    }
}

#[tokio::test]
async fn test_mamba2_ssm_creation() {
    let config = Mamba2Config {
        d_model: 512,
        d_state: 64,
        d_conv: 4,
        expand: 2,
        num_layers: 6,
        vocab_size: 10000,
        pad_vocab_size_multiple: 8,
        tie_embeddings: false,
        dt_rank: "auto".to_string(),
        dt_min: 0.001,
        dt_max: 0.1,
        dt_init: "random".to_string(),
        dt_scale: 1.0,
        dt_init_floor: 1e-4,
        conv_bias: true,
        bias: false,
        use_fast_path: true,
    };

    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 = Mamba2Config {
        d_model: 256,
        d_state: 32,
        d_conv: 4,
        expand: 2,
        num_layers: 4,
        vocab_size: 1000,
        pad_vocab_size_multiple: 8,
        tie_embeddings: false,
        dt_rank: "auto".to_string(),
        dt_min: 0.001,
        dt_max: 0.1,
        dt_init: "random".to_string(),
        dt_scale: 1.0,
        dt_init_floor: 1e-4,
        conv_bias: true,
        bias: false,
        use_fast_path: true,
    };

    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 = Mamba2Config {
        d_model: 128,
        d_state: 16,
        d_conv: 4,
        expand: 2,
        num_layers: 2,
        vocab_size: 1000,
        pad_vocab_size_multiple: 8,
        tie_embeddings: false,
        dt_rank: "auto".to_string(),
        dt_min: 0.001,
        dt_max: 0.1,
        dt_init: "random".to_string(),
        dt_scale: 1.0,
        dt_init_floor: 1e-4,
        conv_bias: true,
        bias: false,
        use_fast_path: true,
    };

    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 selective_state = SelectiveStateSpace {
        importance_tracker: Vec::new(),
        active_indices: Vec::new(),
        compressed_states: BTreeMap::new(),
        compression_threshold: ImportanceThreshold {
            min_importance: 0.1,
            max_active_states: 1000,
            compression_ratio: 0.8,
        },
        memory_usage_bytes: 0,
        max_memory_bytes: 1024 * 1024, // 1MB
    };

    assert!(selective_state.importance_tracker.is_empty());
    assert!(selective_state.active_indices.is_empty());
    assert!(selective_state.compressed_states.is_empty());
    assert_eq!(selective_state.compression_threshold.min_importance, 0.1);
    assert_eq!(selective_state.memory_usage_bytes, 0);
}

#[tokio::test]
async fn test_selective_state_importance_scoring() {
    let mut selective_state = SelectiveStateSpace {
        importance_tracker: Vec::new(),
        active_indices: Vec::new(),
        compressed_states: BTreeMap::new(),
        compression_threshold: ImportanceThreshold {
            min_importance: 0.1,
            max_active_states: 5, // Small limit for testing
            compression_ratio: 0.8,
        },
        memory_usage_bytes: 0,
        max_memory_bytes: 1024 * 1024,
    };

    // Add state importance scores
    for i in 0..10 {
        let importance = StateImportance {
            state_index: i,
            importance_score: (i as f64) / 10.0, // 0.0 to 0.9
            last_access: std::time::SystemTime::now(),
            access_count: i,
        };
        selective_state.importance_tracker.push(importance);
    }

    // Only states with importance >= 0.1 and within max_active_states should be active
    selective_state.update_active_indices();

    // Should have top 5 most important states (indices 5,6,7,8,9)
    assert_eq!(selective_state.active_indices.len(), 5);
    assert!(selective_state.active_indices.contains(&9)); // Highest importance
    assert!(selective_state.active_indices.contains(&8));
    assert!(!selective_state.active_indices.contains(&0)); // Lowest importance
}

#[tokio::test]
async fn test_selective_state_compression() {
    let mut selective_state = SelectiveStateSpace {
        importance_tracker: Vec::new(),
        active_indices: vec![0, 1, 2],
        compressed_states: BTreeMap::new(),
        compression_threshold: ImportanceThreshold {
            min_importance: 0.1,
            max_active_states: 1000,
            compression_ratio: 0.5, // 50% compression
        },
        memory_usage_bytes: 0,
        max_memory_bytes: 1024,
    };

    // Simulate state compression
    let state_data = vec![1u8, 2, 3, 4, 5, 6, 7, 8]; // 8 bytes
    let compressed_data = vec![1u8, 2, 3, 4]; // 4 bytes (50% compression)

    selective_state.compressed_states.insert(0, compressed_data.clone());
    selective_state.memory_usage_bytes += compressed_data.len();

    assert_eq!(selective_state.compressed_states.len(), 1);
    assert_eq!(selective_state.memory_usage_bytes, 4);
    assert!(selective_state.compressed_states.contains_key(&0));
}

#[tokio::test]
async fn test_mamba2_training_step() {
    let config = Mamba2Config {
        d_model: 128,
        d_state: 16,
        d_conv: 4,
        expand: 2,
        num_layers: 2,
        vocab_size: 1000,
        pad_vocab_size_multiple: 8,
        tie_embeddings: false,
        dt_rank: "auto".to_string(),
        dt_min: 0.001,
        dt_max: 0.1,
        dt_init: "random".to_string(),
        dt_scale: 1.0,
        dt_init_floor: 1e-4,
        conv_bias: true,
        bias: false,
        use_fast_path: true,
    };

    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 = Mamba2Config {
        d_model: 64,
        d_state: 8,
        d_conv: 4,
        expand: 2,
        num_layers: 2,
        vocab_size: 100,
        pad_vocab_size_multiple: 8,
        tie_embeddings: false,
        dt_rank: "auto".to_string(),
        dt_min: 0.001,
        dt_max: 0.1,
        dt_init: "random".to_string(),
        dt_scale: 1.0,
        dt_init_floor: 1e-4,
        conv_bias: true,
        bias: false,
        use_fast_path: true,
    };

    let mut state = Mamba2State::new(&config);

    // Test state initialization
    assert_eq!(state.current_position, 0);
    assert!(state.hidden_states.is_empty() || state.hidden_states.len() == config.num_layers);

    // Test state update
    state.update_position(5);
    assert_eq!(state.current_position, 5);
}

#[tokio::test]
async fn test_mamba2_discretization_methods() {
    let config = Mamba2Config {
        d_model: 32,
        d_state: 4,
        d_conv: 4,
        expand: 2,
        num_layers: 1,
        vocab_size: 100,
        pad_vocab_size_multiple: 8,
        tie_embeddings: false,
        dt_rank: "auto".to_string(),
        dt_min: 0.001,
        dt_max: 0.1,
        dt_init: "random".to_string(),
        dt_scale: 1.0,
        dt_init_floor: 1e-4,
        conv_bias: true,
        bias: false,
        use_fast_path: true,
    };

    // Test different dt initialization methods
    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 dt_init
    let mut config2 = config.clone();
    config2.dt_init = "constant".to_string();
    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 = Mamba2Config {
        d_model: 256,
        d_state: 32,
        d_conv: 4,
        expand: 2,
        num_layers: 4,
        vocab_size: 1000,
        pad_vocab_size_multiple: 8,
        tie_embeddings: false,
        dt_rank: "auto".to_string(),
        dt_min: 0.001,
        dt_max: 0.1,
        dt_init: "random".to_string(),
        dt_scale: 1.0,
        dt_init_floor: 1e-4,
        conv_bias: true,
        bias: false,
        use_fast_path: true,
    };

    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 = Mamba2Config {
        d_model: 128,
        d_state: 16,
        d_conv: 4,
        expand: 2,
        num_layers: 2,
        vocab_size: 1000,
        pad_vocab_size_multiple: 8,
        tie_embeddings: false,
        dt_rank: "auto".to_string(),
        dt_min: 0.001,
        dt_max: 0.1,
        dt_init: "random".to_string(),
        dt_scale: 1.0,
        dt_init_floor: 1e-4,
        conv_bias: true,
        bias: false,
        use_fast_path: true, // Hardware optimization enabled
    };

    let mut fast_model = MockMamba2SSM::new(config.clone());
    config.use_fast_path = false; // Disable 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 fast 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,
        dt_min in 0.0001..0.01_f64,
        dt_max in 0.05..0.2_f64,
    ) {
        prop_assume!(dt_max > dt_min);

        let config = Mamba2Config {
            d_model: d_model as usize,
            d_state: d_state as usize,
            d_conv: 4,
            expand: 2,
            num_layers,
            vocab_size: 1000,
            pad_vocab_size_multiple: 8,
            tie_embeddings: false,
            dt_rank: "auto".to_string(),
            dt_min,
            dt_max,
            dt_init: "random".to_string(),
            dt_scale: 1.0,
            dt_init_floor: 1e-4,
            conv_bias: true,
            bias: false,
            use_fast_path: true,
        };

        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.dt_min < model.config.dt_max);
    }
}