foxhunt/ml/tests/mamba2_p0_new_fixes_test.rs

//! # MAMBA-2 P0 New Fixes Test Suite
//!
//! Tests for the 3 new P0 fixes:
//! 1. Sigmoid activation constrains output to [0,1]
//! 2. total_decay_steps from config (not hardcoded)
//! 3. d_state defaults to 64 (Mamba-2 official recommendation)

#![allow(unused_crate_dependencies)]

use candle_core::{Device, Tensor};
use ml::mamba::{Mamba2Config, Mamba2SSM};
use ml::MLError;

/// Test Fix #1: Sigmoid activation constrains output to [0,1]
#[test]
fn test_p0_fix1_sigmoid_activation_output_range() -> Result<(), MLError> {
    println!("\n=== P0 Fix #1: Sigmoid Activation Output Range ===");

    let device = Device::cuda_if_available(0)?;
    let mut model = Mamba2SSM::default_hft(&device)?;

    // Create test input
    let batch_size = 4;
    let seq_len = 16;
    let d_model = model.config.d_model;

    let input_data = vec![0.5f64; batch_size * seq_len * d_model];
    let input = Tensor::from_vec(input_data, (batch_size, seq_len, d_model), &device)?;

    println!("Running forward pass...");
    let output = model.forward(&input)?;

    // Extract output values
    let output_vec = output.flatten_all()?.to_vec1::<f64>()?;

    println!("Output shape: {:?}", output.dims());
    println!(
        "Output samples (first 10): {:?}",
        &output_vec[..10.min(output_vec.len())]
    );

    // ASSERTION: All outputs should be in [0, 1] due to sigmoid activation
    let min_val = output_vec.iter().cloned().fold(f64::INFINITY, f64::min);
    let max_val = output_vec.iter().cloned().fold(f64::NEG_INFINITY, f64::max);

    println!("Output range: [{:.6}, {:.6}]", min_val, max_val);

    assert!(
        min_val >= 0.0 && min_val <= 1.0,
        "FAIL: Min value {:.6} outside [0,1] range",
        min_val
    );
    assert!(
        max_val >= 0.0 && max_val <= 1.0,
        "FAIL: Max value {:.6} outside [0,1] range",
        max_val
    );

    // Check that not all values are exactly 0 or 1 (sigmoid should produce continuous values)
    let mid_range_count = output_vec.iter().filter(|&&v| v > 0.01 && v < 0.99).count();
    println!(
        "Values in (0.01, 0.99): {}/{}",
        mid_range_count,
        output_vec.len()
    );

    assert!(
        mid_range_count > 0,
        "FAIL: No values in mid-range - sigmoid may not be applied"
    );

    println!("\n✅ TEST PASSED: Sigmoid activation constrains output to [0,1]");
    Ok(())
}

/// Test Fix #2: total_decay_steps from config is used in learning rate scheduler
#[test]
fn test_p0_fix2_total_decay_steps_from_config() -> Result<(), MLError> {
    println!("\n=== P0 Fix #2: total_decay_steps from Config ===");

    let device = Device::cuda_if_available(0)?;

    // Create two models with different total_decay_steps
    let config1 = Mamba2Config {
        d_model: 128,
        d_state: 64,
        d_head: 16,
        num_heads: 4,
        expand: 2,
        num_layers: 2,
        dropout: 0.1,
        use_ssd: true,
        use_selective_state: true,
        hardware_aware: false,
        target_latency_us: 1000,
        max_seq_len: 128,
        learning_rate: 1e-3,
        weight_decay: 1e-4,
        grad_clip: 1.0,
        warmup_steps: 100,
        adam_beta1: 0.9,
        adam_beta2: 0.999,
        adam_epsilon: 1e-8,
        total_decay_steps: 1000, // Short decay
        optimizer_type: ml::mamba::OptimizerType::Adam,
        sgd_momentum: 0.9,
        batch_size: 8,
        seq_len: 32,
        shuffle_batches: false,
        seq_stride: 1,
        use_sinusoidal_position_encoding: false,
        max_position_encoding: 2048,
    };

    let config2 = Mamba2Config {
        total_decay_steps: 5000, // Long decay
        ..config1.clone()
    };

    let mut model1 = Mamba2SSM::new(config1.clone(), &device)?;
    let mut model2 = Mamba2SSM::new(config2.clone(), &device)?;

    println!("Model 1 total_decay_steps: {}", config1.total_decay_steps);
    println!("Model 2 total_decay_steps: {}", config2.total_decay_steps);

    // Create test data
    let batch_size = 8;
    let seq_len = 32;
    let d_model = 128;

    let input_data = vec![0.1f64; batch_size * seq_len * d_model];
    let input = Tensor::from_vec(input_data, (batch_size, seq_len, d_model), &device)?;

    let target_data = vec![0.5f64; batch_size * seq_len];
    let target = Tensor::from_vec(target_data, (batch_size, seq_len, 1), &device)?;

    // Train both models for same number of steps (past warmup)
    let num_steps = 200;
    println!(
        "\nTraining both models for {} steps (warmup: {})...",
        num_steps, config1.warmup_steps
    );

    for step in 0..num_steps {
        // Model 1
        let output1 = model1.forward(&input)?;
        let diff1 = output1.broadcast_sub(&target)?;
        let loss1 = diff1.sqr()?.mean_all()?;
        model1.backward_pass(&loss1, &input, &target)?;
        model1.optimizer_step()?;

        // Model 2
        let output2 = model2.forward(&input)?;
        let diff2 = output2.broadcast_sub(&target)?;
        let loss2 = diff2.sqr()?.mean_all()?;
        model2.backward_pass(&loss2, &input, &target)?;
        model2.optimizer_step()?;

        if step % 50 == 0 {
            println!(
                "  Step {}: LR1={:.6}, LR2={:.6}",
                step, model1.current_lr, model2.current_lr
            );
        }
    }

    println!("\nFinal learning rates:");
    println!("  Model 1 (decay_steps=1000): {:.6}", model1.current_lr);
    println!("  Model 2 (decay_steps=5000): {:.6}", model2.current_lr);

    // ASSERTION: Model 1 should have lower LR than Model 2
    // (faster decay due to shorter total_decay_steps)
    assert!(
        model1.current_lr < model2.current_lr,
        "FAIL: Model 1 LR ({:.6}) should be < Model 2 LR ({:.6}) due to shorter decay_steps",
        model1.current_lr,
        model2.current_lr
    );

    // The difference should be noticeable (>5% different)
    let lr_ratio = model1.current_lr / model2.current_lr;
    println!("LR ratio (model1/model2): {:.3}", lr_ratio);

    assert!(
        lr_ratio < 0.95,
        "FAIL: LR difference too small ({:.3}). Config may not be respected.",
        lr_ratio
    );

    println!("\n✅ TEST PASSED: total_decay_steps from config is respected");
    Ok(())
}

/// Test Fix #3: d_state defaults to 64 (Mamba-2 official recommendation)
#[test]
fn test_p0_fix3_d_state_defaults_to_64() -> Result<(), MLError> {
    println!("\n=== P0 Fix #3: d_state Defaults to 64 ===");

    let device = Device::cuda_if_available(0)?;

    // Test 1: emergency_safe_defaults()
    let config1 = Mamba2Config::emergency_safe_defaults();
    println!("emergency_safe_defaults() d_state: {}", config1.d_state);
    assert_eq!(
        config1.d_state, 64,
        "FAIL: emergency_safe_defaults() should have d_state=64, got {}",
        config1.d_state
    );

    // Test 2: default_hft()
    let model = Mamba2SSM::default_hft(&device)?;
    println!("default_hft() d_state: {}", model.config.d_state);
    assert_eq!(
        model.config.d_state, 64,
        "FAIL: default_hft() should have d_state=64, got {}",
        model.config.d_state
    );

    // Test 3: Verify model actually uses d_state=64 (check SSM state dimensions)
    let batch_size = 4;
    let seq_len = 16;
    let d_model = model.config.d_model;

    let input_data = vec![0.5f64; batch_size * seq_len * d_model];
    let input = Tensor::from_vec(input_data, (batch_size, seq_len, d_model), &device)?;

    // Forward pass to initialize SSM states
    let _ = model.forward(&input)?;

    // Check SSM state dimensions
    if !model.state.ssm_states.is_empty() {
        let ssm_state = &model.state.ssm_states[0];

        // A matrix should be [d_state, d_state]
        let a_dims = ssm_state.A.dims();
        println!("SSM A matrix shape: {:?} (expected: [64, 64])", a_dims);

        assert_eq!(
            a_dims[0], 64,
            "FAIL: A matrix first dim should be 64, got {}",
            a_dims[0]
        );

        // B matrix should be [d_state, d_inner]
        let b_dims = ssm_state.B.dims();
        println!("SSM B matrix shape: {:?} (expected: [64, d_inner])", b_dims);

        assert_eq!(
            b_dims[0], 64,
            "FAIL: B matrix first dim (d_state) should be 64, got {}",
            b_dims[0]
        );

        // C matrix should be [d_inner, d_state]
        let c_dims = ssm_state.C.dims();
        println!("SSM C matrix shape: {:?} (expected: [d_inner, 64])", c_dims);

        assert_eq!(
            c_dims[1], 64,
            "FAIL: C matrix second dim (d_state) should be 64, got {}",
            c_dims[1]
        );
    }

    println!("\n✅ TEST PASSED: d_state defaults to 64 (Mamba-2 official recommendation)");
    Ok(())
}

/// Integration test: All 3 P0 fixes work together
#[test]
fn test_p0_integration_all_three_fixes() -> Result<(), MLError> {
    println!("\n=== P0 Integration: All 3 Fixes Together ===");

    let device = Device::cuda_if_available(0)?;

    // Use default_hft which should have all fixes
    let mut model = Mamba2SSM::default_hft(&device)?;

    println!("Config d_state: {}", model.config.d_state);
    println!(
        "Config total_decay_steps: {}",
        model.config.total_decay_steps
    );

    // Create test data
    let batch_size = 8;
    let seq_len = 32;
    let d_model = model.config.d_model;

    let input_data = vec![0.3f64; batch_size * seq_len * d_model];
    let input = Tensor::from_vec(input_data, (batch_size, seq_len, d_model), &device)?;

    let target_data = vec![0.7f64; batch_size * seq_len];
    let target = Tensor::from_vec(target_data, (batch_size, seq_len, 1), &device)?;

    // Train for a few steps
    let mut initial_lr = 0.0;
    let mut final_lr = 0.0;
    let mut output_ranges = Vec::new();

    println!("\nTraining for 50 steps...");
    for step in 0..50 {
        let output = model.forward(&input)?;
        let diff = output.broadcast_sub(&target)?;
        let loss = diff.sqr()?.mean_all()?;

        model.backward_pass(&loss, &input, &target)?;
        model.optimizer_step()?;

        if step == 0 {
            initial_lr = model.current_lr;
        }
        if step == 49 {
            final_lr = model.current_lr;
        }

        // Check output range
        let output_vec = output.flatten_all()?.to_vec1::<f64>()?;
        let min_val = output_vec.iter().cloned().fold(f64::INFINITY, f64::min);
        let max_val = output_vec.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
        output_ranges.push((min_val, max_val));

        if step % 10 == 0 {
            let loss_val = loss.to_scalar::<f64>()?;
            println!(
                "  Step {}: loss={:.6}, LR={:.6}, out_range=[{:.3}, {:.3}]",
                step, loss_val, model.current_lr, min_val, max_val
            );
        }
    }

    // VERIFICATION
    println!("\n--- Verification ---");

    // Fix #1: Sigmoid - all outputs in [0,1]
    let mut all_in_range = true;
    for (min_val, max_val) in &output_ranges {
        if *min_val < 0.0 || *max_val > 1.0 {
            all_in_range = false;
            break;
        }
    }
    assert!(all_in_range, "FAIL: Some outputs outside [0,1] range");
    println!("✓ Fix #1: All outputs in [0,1] range (sigmoid working)");

    // Fix #2: Learning rate changed (using config total_decay_steps)
    assert_ne!(
        initial_lr, final_lr,
        "FAIL: Learning rate should change over training"
    );
    println!(
        "✓ Fix #2: Learning rate changed: {:.6} → {:.6}",
        initial_lr, final_lr
    );

    // Fix #3: d_state=64
    assert_eq!(model.config.d_state, 64, "FAIL: d_state should be 64");
    println!(
        "✓ Fix #3: d_state={} (Mamba-2 official)",
        model.config.d_state
    );

    println!("\n✅ TEST PASSED: All 3 P0 fixes working correctly");
    Ok(())
}