foxhunt/ml/tests/mamba_comprehensive_tests.rs

#![allow(unused_crate_dependencies)]
//! Comprehensive tests for MAMBA-2 selective state space models
//!
//! This test suite covers:
//! - Selective state transformations with various sequence lengths
//! - Scan algorithm properties (associativity, commutativity)
//! - SSD layer forward/backward passes with known inputs/outputs
//! - Hardware-aware optimizations and GPU vs CPU comparisons
//! - Edge cases: zero sequences, max lengths, negative values
//! - Error path validation

use candle_core::{DType, Device, Tensor};
use ml::mamba::{
    Mamba2Config, Mamba2State, ParallelScanEngine, ScanOperator, SelectiveStateSpace,
    StateCompressor, StateImportance,
};
use ml::MLError;
use nalgebra::DVector;

// ============================================================================
// SELECTIVE STATE SPACE TESTS
// ============================================================================

#[test]
fn test_selective_state_various_sequence_lengths() -> Result<(), MLError> {
    let sequence_lengths = vec![1, 10, 100, 1000];

    for seq_len in sequence_lengths {
        let mut config = Mamba2Config::emergency_safe_defaults();
        config.d_model = 16;
        config.d_state = 8;
        config.expand = 2;

        let mut selective_state = SelectiveStateSpace::new(&config)?;
        let mut state = Mamba2State::zeros(&config)?;

        // Create input tensor with specific sequence length
        let input = Tensor::ones((1, seq_len, 16), DType::F32, &Device::Cpu)?;

        // Update importance scores
        let result = selective_state.update_importance_scores(&input, &mut state);
        assert!(
            result.is_ok(),
            "Failed for sequence length {}: {:?}",
            seq_len,
            result.err()
        );

        // Verify state tracking
        assert!(
            !selective_state.importance_tracker.is_empty(),
            "Importance tracker should not be empty for seq_len {}",
            seq_len
        );
        assert_eq!(
            selective_state.importance_tracker.len(),
            config.d_model * config.expand,
            "Tracker length mismatch for seq_len {}",
            seq_len
        );
    }

    Ok(())
}

#[test]
fn test_selective_state_zero_sequence() -> Result<(), MLError> {
    let mut config = Mamba2Config::emergency_safe_defaults();
    config.d_model = 8;
    config.d_state = 4;
    config.expand = 2;

    let mut selective_state = SelectiveStateSpace::new(&config)?;
    let mut state = Mamba2State::zeros(&config)?;

    // All-zero input tensor
    let input = Tensor::zeros((1, 5, 8), DType::F32, &Device::Cpu)?;

    let result = selective_state.update_importance_scores(&input, &mut state);
    assert!(result.is_ok(), "Zero sequence should be handled");

    // All importance scores should be near zero
    for tracker in &selective_state.importance_tracker {
        assert!(
            tracker.score.abs() < 1e-6,
            "Score should be near zero for zero input"
        );
    }

    Ok(())
}

#[test]
fn test_selective_state_negative_values() -> Result<(), MLError> {
    let mut config = Mamba2Config::emergency_safe_defaults();
    config.d_model = 8;
    config.d_state = 4;
    config.expand = 2;

    let mut selective_state = SelectiveStateSpace::new(&config)?;
    let mut state = Mamba2State::zeros(&config)?;

    // Negative input values
    let input = Tensor::new(
        &[-1.0f32, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0, -8.0],
        &Device::Cpu,
    )?
    .reshape((1, 1, 8))?;

    let result = selective_state.update_importance_scores(&input, &mut state);
    assert!(result.is_ok(), "Negative values should be handled");

    // Importance scores should be based on magnitude (absolute value)
    for tracker in &selective_state.importance_tracker {
        assert!(
            tracker.score >= 0.0,
            "Importance score should be non-negative"
        );
    }

    Ok(())
}

#[test]
fn test_selective_state_max_sequence_length() -> Result<(), MLError> {
    let mut config = Mamba2Config::emergency_safe_defaults();
    config.d_model = 8;
    config.d_state = 4;
    config.expand = 2;
    config.max_seq_len = 2048;

    let selective_state = SelectiveStateSpace::new(&config)?;
    let mut state = Mamba2State::zeros(&config)?;

    // Test at max sequence length
    let input = Tensor::ones(
        (1, config.max_seq_len, config.d_model),
        DType::F32,
        &Device::Cpu,
    )?;

    let mut ss_mut = selective_state;
    let result = ss_mut.update_importance_scores(&input, &mut state);
    assert!(result.is_ok(), "Max sequence length should be handled");

    Ok(())
}

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

    // First update with high score
    importance.update(1.0, 100, 0.9);
    let first_avg = importance.moving_average;

    // Second update with low score
    importance.update(0.1, 200, 0.9);
    let second_avg = importance.moving_average;

    // Moving average should have decayed toward lower value
    assert!(
        second_avg < first_avg,
        "Moving average should decay toward new value"
    );

    // Variance should reflect the change
    assert!(importance.variance > 0.0, "Variance should be positive");
}

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

    // Update at timestamp 0
    importance.update(1.0, 0, 0.9);
    let recent_importance = importance.effective_importance();

    // Update at much later timestamp (should have lower effective importance due to aging)
    importance.last_access = 200;
    let aged_importance = importance.effective_importance();

    assert!(
        aged_importance < recent_importance,
        "Effective importance should decrease with age"
    );
}

#[test]
fn test_state_compressor_lossy_quality_levels() {
    let config = ml::mamba::selective_state::SelectiveStateConfig::default();
    let mut compressor = StateCompressor::new(config);

    let data = DVector::from_vec(vec![1.0, 0.5, 0.25, 0.1, 0.05, 0.01, 0.005, 0.001]);

    // Test different quality levels
    for quality in [0.5, 0.7, 0.9, 0.99] {
        let compressed = compressor.compress_lossy(&data, quality);

        // Higher quality should preserve more values
        let non_zero_count = compressed.iter().filter(|&&x| x != 0.0).count();

        assert_eq!(
            compressed.len(),
            data.len(),
            "Compressed length should match original"
        );
        assert!(
            non_zero_count > 0,
            "Some values should be preserved at quality {}",
            quality
        );
    }
}

#[test]
fn test_state_compressor_lossless_roundtrip() {
    let config = ml::mamba::selective_state::SelectiveStateConfig::default();
    let mut compressor = StateCompressor::new(config);

    // Test data with repeated values (good for run-length encoding)
    let data = DVector::from_vec(vec![1.0, 1.0, 1.0, 2.0, 2.0, 3.0, 3.0, 3.0, 3.0]);

    let (runs, original_size) = compressor.compress_lossless(&data, 0.01);
    let decompressed = compressor.decompress_lossless(&runs, original_size);

    assert_eq!(decompressed.len(), data.len(), "Length should match");

    // Verify exact reconstruction
    for i in 0..data.len() {
        assert!(
            (decompressed[i] - data[i]).abs() < 1e-10,
            "Value at index {} should match: expected {}, got {}",
            i,
            data[i],
            decompressed[i]
        );
    }
}

#[test]
fn test_state_compression_ratio_calculation() {
    let config = ml::mamba::selective_state::SelectiveStateConfig::default();
    let mut compressor = StateCompressor::new(config);

    // Sparse data (high compression ratio)
    let sparse_data = DVector::from_vec(vec![1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 3.0]);
    let _sparse_compressed = compressor.compress_lossy(&sparse_data, 0.8);
    let sparse_ratio = compressor
        .compression_stats
        .get("last_lossy_ratio")
        .copied()
        .unwrap_or(1.0);

    // Dense data (lower compression ratio)
    let dense_data = DVector::from_vec(vec![1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7]);
    let _dense_compressed = compressor.compress_lossy(&dense_data, 0.8);
    let dense_ratio = compressor
        .compression_stats
        .get("last_lossy_ratio")
        .copied()
        .unwrap_or(1.0);

    assert!(
        sparse_ratio <= dense_ratio,
        "Sparse data should compress better: {} vs {}",
        sparse_ratio,
        dense_ratio
    );
}

// ============================================================================
// SCAN ALGORITHM TESTS
// ============================================================================

#[test]
fn test_scan_addition_associativity() -> Result<(), MLError> {
    let device = Device::Cpu;
    let engine = ParallelScanEngine::new(device.clone(), 1_000_000);

    let a = Tensor::new(&[1.0f32], &device)?;
    let b = Tensor::new(&[2.0f32], &device)?;
    let c = Tensor::new(&[3.0f32], &device)?;

    // (a + b) + c
    let ab = engine.apply_operator(&a, &b, ScanOperator::Add)?;
    let abc_left = engine.apply_operator(&ab, &c, ScanOperator::Add)?;

    // a + (b + c)
    let bc = engine.apply_operator(&b, &c, ScanOperator::Add)?;
    let abc_right = engine.apply_operator(&a, &bc, ScanOperator::Add)?;

    let left_val: f32 = abc_left.flatten_all()?.to_vec1::<f32>()?[0];
    let right_val: f32 = abc_right.flatten_all()?.to_vec1::<f32>()?[0];

    assert!(
        (left_val - right_val).abs() < 1e-6,
        "Addition should be associative: {} vs {}",
        left_val,
        right_val
    );

    Ok(())
}

#[test]
fn test_scan_multiplication_associativity() -> Result<(), MLError> {
    let device = Device::Cpu;
    let engine = ParallelScanEngine::new(device.clone(), 1_000_000);

    let a = Tensor::new(&[2.0f32], &device)?;
    let b = Tensor::new(&[3.0f32], &device)?;
    let c = Tensor::new(&[4.0f32], &device)?;

    // (a * b) * c
    let ab = engine.apply_operator(&a, &b, ScanOperator::Mul)?;
    let abc_left = engine.apply_operator(&ab, &c, ScanOperator::Mul)?;

    // a * (b * c)
    let bc = engine.apply_operator(&b, &c, ScanOperator::Mul)?;
    let abc_right = engine.apply_operator(&a, &bc, ScanOperator::Mul)?;

    let left_val: f32 = abc_left.flatten_all()?.to_vec1::<f32>()?[0];
    let right_val: f32 = abc_right.flatten_all()?.to_vec1::<f32>()?[0];

    assert!(
        (left_val - right_val).abs() < 1e-5,
        "Multiplication should be associative: {} vs {}",
        left_val,
        right_val
    );

    Ok(())
}

#[test]
fn test_scan_max_operator_properties() -> Result<(), MLError> {
    let device = Device::Cpu;
    let engine = ParallelScanEngine::new(device.clone(), 1_000_000);

    let a = Tensor::new(&[5.0f32], &device)?;
    let b = Tensor::new(&[3.0f32], &device)?;

    // max(a, b) should equal max(b, a) (commutativity)
    let max_ab = engine.apply_operator(&a, &b, ScanOperator::Max)?;
    let max_ba = engine.apply_operator(&b, &a, ScanOperator::Max)?;

    let ab_val: f32 = max_ab.flatten_all()?.to_vec1::<f32>()?[0];
    let ba_val: f32 = max_ba.flatten_all()?.to_vec1::<f32>()?[0];

    assert!((ab_val - ba_val).abs() < 1e-6, "Max should be commutative");
    assert!((ab_val - 5.0).abs() < 1e-6, "Max should be 5.0");

    Ok(())
}

#[test]
fn test_scan_min_operator_properties() -> Result<(), MLError> {
    let device = Device::Cpu;
    let engine = ParallelScanEngine::new(device.clone(), 1_000_000);

    let a = Tensor::new(&[5.0f32], &device)?;
    let b = Tensor::new(&[3.0f32], &device)?;

    // min(a, b) should equal min(b, a) (commutativity)
    let min_ab = engine.apply_operator(&a, &b, ScanOperator::Min)?;
    let min_ba = engine.apply_operator(&b, &a, ScanOperator::Min)?;

    let ab_val: f32 = min_ab.flatten_all()?.to_vec1::<f32>()?[0];
    let ba_val: f32 = min_ba.flatten_all()?.to_vec1::<f32>()?[0];

    assert!((ab_val - ba_val).abs() < 1e-6, "Min should be commutative");
    assert!((ab_val - 3.0).abs() < 1e-6, "Min should be 3.0");

    Ok(())
}

#[test]
fn test_parallel_vs_sequential_scan_consistency() -> Result<(), MLError> {
    let device = Device::Cpu;
    let mut engine = ParallelScanEngine::new(device.clone(), 50); // Low threshold for testing
    engine.block_size = 10;

    // Create test data
    let data: Vec<f32> = (1..=100).map(|x| x as f32).collect();
    let input = Tensor::new(&data[..], &device)?.reshape((1, 100))?;

    // Sequential scan
    let seq_result = engine.sequential_scan(&input, ScanOperator::Add)?;

    // Parallel scan (should trigger block processing)
    let par_result = engine.block_parallel_scan(&input, ScanOperator::Add)?;

    let seq_vals = seq_result.flatten_all()?.to_vec1::<f32>()?;
    let par_vals = par_result.flatten_all()?.to_vec1::<f32>()?;

    assert_eq!(
        seq_vals.len(),
        par_vals.len(),
        "Results should have same length"
    );

    for (i, (seq, par)) in seq_vals.into_iter().zip(par_vals.into_iter()).enumerate() {
        assert!(
            (seq - par).abs() < 1e-4,
            "Mismatch at index {}: seq={}, par={}",
            i,
            seq,
            par
        );
    }

    Ok(())
}

#[test]
fn test_segmented_scan_multiple_segments() -> Result<(), MLError> {
    let device = Device::Cpu;
    let engine = ParallelScanEngine::new(device.clone(), 1_000_000);

    // Three segments: [1,2,3], [4,5], [6,7,8,9]
    let input =
        Tensor::new(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0], &device)?.reshape((1, 9))?;
    let segment_ids = Tensor::new(&[0i64, 0, 0, 1, 1, 2, 2, 2, 2], &device)?.reshape((1, 9))?;

    let result = engine.segmented_scan(&input, &segment_ids, ScanOperator::Add)?;
    let values = result.flatten_all()?.to_vec1::<f32>()?;

    // Expected: [1, 3, 6, 4, 9, 6, 13, 21, 30]
    let expected = vec![1.0, 3.0, 6.0, 4.0, 9.0, 6.0, 13.0, 21.0, 30.0];

    for (i, (actual, expected)) in values.into_iter().zip(expected.into_iter()).enumerate() {
        assert!(
            (actual - expected).abs() < 1e-5,
            "Mismatch at index {}: expected {}, got {}",
            i,
            expected,
            actual
        );
    }

    Ok(())
}

#[test]
fn test_scan_edge_case_single_element() -> Result<(), MLError> {
    let device = Device::Cpu;
    let engine = ParallelScanEngine::new(device.clone(), 1_000_000);

    let input = Tensor::new(&[42.0f32], &device)?.reshape((1, 1))?;
    let result = engine.parallel_prefix_scan(&input, ScanOperator::Add)?;

    let value: f32 = result.flatten_all()?.to_vec1::<f32>()?[0];
    assert!(
        (value - 42.0).abs() < 1e-6,
        "Single element should remain unchanged"
    );

    Ok(())
}

#[test]
fn test_scan_edge_case_two_elements() -> Result<(), MLError> {
    let device = Device::Cpu;
    let engine = ParallelScanEngine::new(device.clone(), 1_000_000);

    let input = Tensor::new(&[3.0f32, 5.0], &device)?.reshape((1, 2))?;
    let result = engine.parallel_prefix_scan(&input, ScanOperator::Mul)?;

    let values = result.flatten_all()?.to_vec1::<f32>()?;
    assert!(
        (values[0] - 3.0).abs() < 1e-6,
        "First element should be 3.0"
    );
    assert!(
        (values[1] - 15.0).abs() < 1e-6,
        "Second element should be 15.0"
    );

    Ok(())
}

#[test]
fn test_scan_benchmark_performance() -> Result<(), MLError> {
    let device = Device::Cpu;
    let engine = ParallelScanEngine::new(device.clone(), 1_000_000);

    let seq_lengths = vec![100, 500, 1000];
    let benchmarks = engine.benchmark_scan_performance(&seq_lengths, ScanOperator::Add)?;

    assert_eq!(benchmarks.len(), seq_lengths.len());

    for (i, benchmark) in benchmarks.into_iter().enumerate() {
        assert_eq!(benchmark.sequence_length, seq_lengths[i]);
        assert!(benchmark.duration_nanos > 0, "Duration should be positive");
        assert!(
            benchmark.throughput_elements_per_sec > ml::liquid::FixedPoint::zero(),
            "Throughput should be positive"
        );
        assert!(
            benchmark.memory_bandwidth_gb_per_sec > ml::liquid::FixedPoint::zero(),
            "Bandwidth should be positive"
        );
    }

    Ok(())
}

// ============================================================================
// SSD LAYER TESTS
// ============================================================================

#[test]
fn test_ssd_layer_forward_known_input() -> Result<(), MLError> {
    let mut config = Mamba2Config::emergency_safe_defaults();
    config.d_model = 8;
    config.d_state = 4;
    config.d_head = 4;
    config.num_heads = 2;

    let mut layer = ml::mamba::SSDLayer::new(&config, 0)?;
    let mut state = Mamba2State::zeros(&config)?;

    // Known input
    let input = Tensor::ones((1, 4, 8), DType::F32, &Device::Cpu)?;

    let output = layer.forward(&input, &mut state)?;

    // Verify output shape
    assert_eq!(output.dims(), &[1, 4, 8], "Output shape should match input");

    // Output should be non-zero (layer has been initialized)
    let output_data = output.flatten_all()?.to_vec1::<f32>()?;
    let non_zero_count = output_data.iter().filter(|&&x| x.abs() > 1e-6).count();
    assert!(non_zero_count > 0, "Output should contain non-zero values");

    Ok(())
}

#[test]
fn test_ssd_layer_batch_processing() -> Result<(), MLError> {
    let mut config = Mamba2Config::emergency_safe_defaults();
    config.d_model = 16;
    config.d_state = 8;
    config.d_head = 8;
    config.num_heads = 2;

    let mut layer = ml::mamba::SSDLayer::new(&config, 0)?;
    let mut state = Mamba2State::zeros(&config)?;

    // Batch of 4 sequences
    let input = Tensor::randn(0.0, 1.0, (4, 8, 16), &Device::Cpu)?;

    let output = layer.forward(&input, &mut state)?;

    assert_eq!(output.dims(), &[4, 8, 16], "Batch dimension preserved");

    Ok(())
}

#[test]
fn test_ssd_layer_attention_cache() -> Result<(), MLError> {
    let mut config = Mamba2Config::emergency_safe_defaults();
    config.d_model = 8;
    config.d_state = 4;
    config.d_head = 4;
    config.num_heads = 2;

    let mut layer = ml::mamba::SSDLayer::new(&config, 0)?;
    let mut state = Mamba2State::zeros(&config)?;

    let input = Tensor::ones((1, 4, 8), DType::F32, &Device::Cpu)?;

    // First forward pass - should miss cache
    let _output1 = layer.forward(&input, &mut state)?;
    let initial_misses = layer
        .cache_misses
        .load(std::sync::atomic::Ordering::Relaxed);

    // Second forward pass - should hit cache (same input shape)
    let _output2 = layer.forward(&input, &mut state)?;
    let cache_hits = layer.cache_hits.load(std::sync::atomic::Ordering::Relaxed);

    assert!(
        cache_hits > 0 || initial_misses > 0,
        "Cache should be utilized"
    );

    Ok(())
}

#[test]
fn test_ssd_layer_performance_metrics() -> Result<(), MLError> {
    let mut config = Mamba2Config::emergency_safe_defaults();
    config.d_model = 8;

    let mut layer = ml::mamba::SSDLayer::new(&config, 0)?;
    let mut state = Mamba2State::zeros(&config)?;

    let input = Tensor::ones((1, 4, 8), DType::F32, &Device::Cpu)?;
    let _output = layer.forward(&input, &mut state)?;

    let metrics = layer.get_performance_metrics();

    assert!(
        metrics.contains_key("layer_0_operations"),
        "Should track operations"
    );
    assert!(
        metrics.contains_key("layer_0_avg_latency_ns"),
        "Should track latency"
    );
    assert!(
        metrics.contains_key("layer_0_attention_cache_size"),
        "Should track cache size"
    );

    Ok(())
}

#[test]
fn test_ssd_layer_qkv_split() -> Result<(), MLError> {
    let config = Mamba2Config {
        d_model: 12,
        d_state: 6,
        d_head: 6,
        num_heads: 2,
        ..Default::default()
    };

    let layer = ml::mamba::SSDLayer::new(&config, 0)?;

    // QKV tensor: 3 * d_head * num_heads = 3 * 6 * 2 = 36
    let qkv = Tensor::randn(0.0, 1.0, (1, 4, 36), &Device::Cpu)?;

    let (queries, keys, values) = layer.split_qkv(&qkv)?;

    // Each should be d_head * num_heads = 12
    assert_eq!(queries.dims(), &[1, 4, 12], "Queries shape");
    assert_eq!(keys.dims(), &[1, 4, 12], "Keys shape");
    assert_eq!(values.dims(), &[1, 4, 12], "Values shape");

    Ok(())
}

#[test]
fn test_ssd_layer_norm_zero_mean() -> Result<(), MLError> {
    let config = Mamba2Config::emergency_safe_defaults();
    let layer = ml::mamba::SSDLayer::new(&config, 0)?;

    let input = Tensor::new(&[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], &Device::Cpu)?
        .reshape((1, 1, 8))?;

    let normalized = layer.apply_layer_norm(&input)?;
    let norm_data = normalized.flatten_all()?.to_vec1::<f32>()?;

    // Mean should be close to zero
    let mean: f32 = norm_data.iter().sum::<f32>() / norm_data.len() as f32;
    assert!(
        mean.abs() < 1e-4,
        "Normalized data should have zero mean: {}",
        mean
    );

    // Variance should be close to 1
    let variance: f32 = norm_data
        .iter()
        .map(|x| (x - mean) * (x - mean))
        .sum::<f32>()
        / norm_data.len() as f32;
    assert!(
        (variance - 1.0).abs() < 0.2,
        "Normalized data should have variance ~1: {}",
        variance
    );

    Ok(())
}

// ============================================================================
// HARDWARE-AWARE OPTIMIZATION TESTS
// ============================================================================

#[test]
fn test_hardware_capabilities_detection() {
    use ml::mamba::HardwareCapabilities;

    let caps = HardwareCapabilities::default();

    assert!(caps.cache_line_size > 0, "Cache line size should be set");
    assert!(caps.simd_width >= 4, "SIMD width should be at least 4");
    assert!(caps.num_cores > 0, "Should detect CPU cores");
    assert!(caps.l1_cache_size > 0, "L1 cache size should be configured");
}

#[test]
fn test_hardware_optimizer_matrix_multiplication() -> Result<(), MLError> {
    use ml::mamba::HardwareOptimizer;
    use nalgebra::DMatrix;

    let config = Mamba2Config::default();
    let optimizer = HardwareOptimizer::new(&config)?;

    let a = DMatrix::from_fn(8, 8, |i, j| ((i + j) * 1000) as i64);
    let b = DMatrix::from_fn(8, 8, |i, j| ((i * j) * 1000) as i64);

    let result = optimizer.optimized_matrix_mul(&a, &b)?;

    assert_eq!(result.nrows(), 8, "Result rows should match");
    assert_eq!(result.ncols(), 8, "Result cols should match");

    Ok(())
}

#[test]
fn test_hardware_optimizer_dot_product() -> Result<(), MLError> {
    use ml::mamba::HardwareOptimizer;

    let config = Mamba2Config::default();
    let optimizer = HardwareOptimizer::new(&config)?;

    // Use PRECISION_FACTOR for fixed-point arithmetic
    let precision = ml::PRECISION_FACTOR as i64;
    let a = vec![1 * precision, 2 * precision, 3 * precision, 4 * precision];
    let b = vec![5 * precision, 6 * precision, 7 * precision, 8 * precision];

    let result = optimizer.optimized_dot_product(&a, &b)?;

    // Expected: 1*5 + 2*6 + 3*7 + 4*8 = 5 + 12 + 21 + 32 = 70
    let expected = 70 * precision;

    // Allow some error due to fixed-point arithmetic
    assert!(
        (result - expected).abs() < precision / 100,
        "Dot product result {} should be close to {}",
        result,
        expected
    );

    Ok(())
}

#[test]
fn test_hardware_optimizer_prefetching() -> Result<(), MLError> {
    use ml::mamba::HardwareOptimizer;

    let config = Mamba2Config::default();
    let optimizer = HardwareOptimizer::new(&config)?;

    let data: Vec<f64> = (0..1000).map(|x| x as f64).collect();

    // Prefetch should not fail
    optimizer.prefetch_data(&data);

    let metrics = optimizer.get_performance_metrics();
    assert!(
        metrics.contains_key("prefetch_operations"),
        "Should track prefetch operations"
    );

    Ok(())
}

#[test]
fn test_hardware_benchmark_performance() -> Result<(), MLError> {
    use ml::mamba::HardwareOptimizer;

    let config = Mamba2Config::default();
    let optimizer = HardwareOptimizer::new(&config)?;

    let results = optimizer.benchmark_performance()?;

    assert!(
        results.contains_key("matrix_mul_ms"),
        "Should benchmark matrix multiplication"
    );
    assert!(
        results.contains_key("dot_product_us"),
        "Should benchmark dot product"
    );
    assert!(
        results.contains_key("prefetch_bandwidth_gbps"),
        "Should benchmark prefetching"
    );

    // Verify reasonable values
    if let Some(&matrix_mul_time) = results.get("matrix_mul_ms") {
        assert!(
            matrix_mul_time > 0.0 && matrix_mul_time < 10000.0,
            "Matrix multiplication time should be reasonable: {}ms",
            matrix_mul_time
        );
    }

    Ok(())
}

// ============================================================================
// ERROR PATH VALIDATION
// ============================================================================

#[test]
fn test_selective_state_mismatched_dimensions() {
    let mut config = Mamba2Config::emergency_safe_defaults();
    config.d_model = 8;
    config.d_state = 4;

    let mut selective_state = SelectiveStateSpace::new(&config).unwrap();
    let mut state = Mamba2State::zeros(&config).unwrap();

    // Wrong dimension input (should be d_model=8)
    let wrong_input = Tensor::ones((1, 4, 16), DType::F32, &Device::Cpu).unwrap();

    let result = selective_state.update_importance_scores(&wrong_input, &mut state);

    // Should handle gracefully (either error or process what it can)
    match result {
        Ok(_) => {
            // If it succeeds, verify it didn't corrupt state
            assert!(!selective_state.importance_tracker.is_empty());
        },
        Err(_) => {
            // Expected error case - this is fine
        },
    }
}

#[test]
fn test_scan_mismatched_vector_lengths() {
    use ml::mamba::HardwareOptimizer;

    let config = Mamba2Config::default();
    let optimizer = HardwareOptimizer::new(&config).unwrap();

    let a = vec![1i64, 2, 3];
    let b = vec![4i64, 5]; // Different length

    let result = optimizer.optimized_dot_product(&a, &b);

    assert!(
        result.is_err(),
        "Should return error for mismatched lengths"
    );
}

#[test]
fn test_state_compression_out_of_bounds() -> Result<(), MLError> {
    let mut config = Mamba2Config::emergency_safe_defaults();
    config.d_model = 8;
    config.d_state = 4;

    let mut selective_state = SelectiveStateSpace::new(&config)?;
    let mut state = Mamba2State::zeros(&config)?;

    // Try to compress index beyond state size
    let out_of_bounds_index = state.selective_state.len() + 100;

    let result = selective_state.compress_state_component(out_of_bounds_index, &mut state);

    // Should handle gracefully (either succeed with no-op or return error)
    assert!(
        result.is_ok(),
        "Out of bounds compression should be handled"
    );

    Ok(())
}