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;

mod real_data_helpers;
use real_data_helpers::{load_tft_sequences, real_data_available};

/// 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);
    }
}

// ============================================================================
// Real Market Data Tests
// ============================================================================

/// Test selective state importance scoring with real market data
#[tokio::test]
async fn test_selective_state_importance_scoring_real_data() {
    // Skip if no real data available
    if !real_data_available().await {
        eprintln!("Skipping test: real data not available");
        return;
    }

    let config = create_test_config(128, 16, 2);
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut test_state = Mamba2State::zeros(&config).unwrap();

    // Load real market data sequences (10 sequences, 10 timesteps each, padded to 128 features)
    let sequences = load_tft_sequences(10, 10, 128).await.unwrap();
    if sequences.is_empty() {
        eprintln!("Skipping test: no sequences loaded");
        return;
    }

    let device = Device::Cpu;

    // Convert first sequence to tensor [1, 10, 128]
    let sequence = &sequences[0];
    let input_data: Vec<f32> = sequence.clone();
    let input = Tensor::from_vec(input_data, &[1, 10, 128], &device).unwrap();

    // Update importance scores with real data
    let result = selective_state.update_importance_scores(&input, &mut test_state);
    assert!(
        result.is_ok(),
        "Importance score update should succeed with real data"
    );

    // Verify that importance tracking is working
    assert!(
        selective_state.active_indices.len() > 0,
        "Should have some active indices with real data"
    );
}

/// Test selective state full workflow with real market data
#[tokio::test]
async fn test_selective_state_full_workflow_real_data() {
    // Skip if no real data available
    if !real_data_available().await {
        eprintln!("Skipping test: real data not available");
        return;
    }

    let config = create_test_config(128, 16, 2);
    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut state = Mamba2State::zeros(&config).unwrap();

    // Load real market data sequences (5 sequences for 5 timesteps)
    let sequences = load_tft_sequences(5, 10, 128).await.unwrap();
    if sequences.len() < 5 {
        eprintln!("Skipping test: not enough sequences loaded");
        return;
    }

    let device = Device::Cpu;

    // Simulate timesteps with real data
    for i in 0..5 {
        let sequence = &sequences[i];
        let input_data: Vec<f32> = sequence.clone();
        let input = Tensor::from_vec(input_data, &[1, 10, 128], &device).unwrap();

        // Update importance scores
        let result = selective_state.update_importance_scores(&input, &mut state);
        assert!(result.is_ok(), "Step {} should succeed with real data", i);

        // Active indices should be maintained
        assert!(
            selective_state.active_indices.len() > 0,
            "Step {} should have active indices",
            i
        );
    }

    // Verify state accumulated properly over timesteps
    assert!(
        state.selective_state.len() > 0,
        "State should have accumulated over real data timesteps"
    );
}