foxhunt/ml/tests/e2e_mamba2_training.rs

//! E2E Test: MAMBA-2 Training Pipeline
//!
//! Fast test that validates MAMBA-2 can train for 3 epochs without crashes.
//! Catches shape mismatches, CUDA errors, and data loading issues.
//!
//! ## Why TDD Approach is Faster
//! - ❌ Current: Build (77s) → Run training → Wait for crash (3s) → Debug → Repeat (5+ minutes per cycle)
//! - ✅ TDD: Write test (1 min) → Run test (5-10s) → Fix → Rerun test (5s) → Deploy (30 seconds per cycle)
//!
//! ## Usage
//! ```bash
//! # Run all MAMBA-2 tests
//! cargo test -p ml mamba2 -- --nocapture
//!
//! # Run single test
//! cargo test -p ml test_mamba2_training_3_epochs -- --nocapture
//!
//! # Run with backtrace
//! RUST_BACKTRACE=1 cargo test -p ml test_mamba2_training_3_epochs -- --nocapture
//! ```

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

/// Helper to create default MAMBA-2 config for testing
fn default_mamba2_config() -> Mamba2Config {
    Mamba2Config {
        d_model: 256,
        d_state: 16,
        d_head: 64,
        num_heads: 4,
        expand: 4,
        num_layers: 2, // Small for testing
        dropout: 0.1,
        use_ssd: true,
        use_selective_state: false,
        hardware_aware: true,
        target_latency_us: 5,
        max_seq_len: 60,
        learning_rate: 0.001,
        weight_decay: 0.0001,
        grad_clip: 1.0,
        warmup_steps: 100,
        batch_size: 16,
        seq_len: 60,
    }
}

#[tokio::test]
async fn test_mamba2_simple_forward_pass() -> Result<()> {
    println!("🧪 E2E Test: MAMBA-2 Simple Forward Pass");

    // Initialize device
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("  Device: {:?}", device);

    // Create small config
    let config = default_mamba2_config();
    println!(
        "  Config: d_model={}, layers={}",
        config.d_model, config.num_layers
    );

    // Create model
    let mut model = Mamba2SSM::new(config.clone(), &device)?;
    println!("  Model created");

    // Create dummy input: [batch=8, seq=60, features=256]
    let batch_size = 8;
    let seq_len = 60;
    let input = Tensor::randn(0f64, 1.0, (batch_size, seq_len, config.d_model), &device)?;
    println!("  Input shape: {:?}", input.dims());

    // Forward pass
    let output = model.forward(&input)?;
    println!("  Output shape: {:?}", output.dims());

    // Validate output shape
    let output_dims = output.dims();
    assert_eq!(output_dims.len(), 3, "Output must be 3D");
    assert_eq!(output_dims[0], batch_size, "Batch size must match");
    assert_eq!(output_dims[1], seq_len, "Sequence length must match");

    println!("✅ Simple forward pass PASSED");
    Ok(())
}

#[tokio::test]
async fn test_mamba2_batch_shapes() -> Result<()> {
    println!("🧪 E2E Test: MAMBA-2 Batch Shape Validation");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("  Device: {:?}", device);

    let config = default_mamba2_config();
    let mut model = Mamba2SSM::new(config.clone(), &device)?;

    // Test different batch sizes
    for batch_size in [1, 8, 16, 32] {
        println!("  Testing batch_size={}", batch_size);

        // Create input: [batch, seq, features]
        let input = Tensor::randn(0f64, 1.0, (batch_size, 60, config.d_model), &device)?;
        println!("    Input shape: {:?}", input.dims());

        // Forward pass
        let output = model.forward(&input)?;
        println!("    Output shape: {:?}", output.dims());

        // Validate output shape
        assert_eq!(
            output.dims()[0],
            batch_size,
            "Output batch size {} must match input batch size {}",
            output.dims()[0],
            batch_size
        );
        assert_eq!(
            output.dims()[1],
            60,
            "Output seq length {} must be 60",
            output.dims()[1]
        );

        println!("    ✓ batch_size={} works", batch_size);
    }

    println!("✅ Shape validation PASSED");
    Ok(())
}

#[tokio::test]
async fn test_mamba2_cuda_device() -> Result<()> {
    println!("🧪 E2E Test: MAMBA-2 CUDA Device");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("  Device: {:?}", device);

    let config = default_mamba2_config();
    let mut model = Mamba2SSM::new(config.clone(), &device)?;
    println!("  Model created on device: {:?}", device);

    // Create tensor on device
    let input = Tensor::randn(0f64, 1.0, (16, 60, config.d_model), &device)?;
    println!("  Input tensor created on device: {:?}", input.device());

    // Forward pass
    let output = model.forward(&input)?;
    println!("  Output tensor on device: {:?}", output.device());

    // Verify output is on same device
    match (&device, output.device()) {
        (Device::Cuda(_), Device::Cuda(_)) => {
            println!("  ✓ CUDA device working");
        },
        (Device::Cpu, Device::Cpu) => {
            println!("  ✓ CPU device working (CUDA not available)");
        },
        _ => {
            panic!(
                "Device mismatch: expected {:?}, got {:?}",
                device,
                output.device()
            );
        },
    }

    println!("✅ Device test PASSED");
    Ok(())
}

#[tokio::test]
async fn test_mamba2_sequence_lengths() -> Result<()> {
    println!("🧪 E2E Test: MAMBA-2 Sequence Length Validation");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("  Device: {:?}", device);

    let config = default_mamba2_config();
    let mut model = Mamba2SSM::new(config.clone(), &device)?;

    // Test different sequence lengths
    for seq_len in [10, 30, 60, 120] {
        println!("  Testing seq_len={}", seq_len);

        // Create input: [batch, seq, features]
        let input = Tensor::randn(0f64, 1.0, (16, seq_len, config.d_model), &device)?;
        println!("    Input shape: {:?}", input.dims());

        // Forward pass
        let output = model.forward(&input)?;
        println!("    Output shape: {:?}", output.dims());

        // Validate output shape
        assert_eq!(output.dims()[0], 16, "Output batch size must be 16");
        assert_eq!(
            output.dims()[1],
            seq_len,
            "Output seq length {} must match input seq length {}",
            output.dims()[1],
            seq_len
        );

        println!("    ✓ seq_len={} works", seq_len);
    }

    println!("✅ Sequence length validation PASSED");
    Ok(())
}

#[tokio::test]
async fn test_mamba2_gradient_flow() -> Result<()> {
    println!("🧪 E2E Test: MAMBA-2 Gradient Flow");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("  Device: {:?}", device);

    // Create model
    let config = default_mamba2_config();
    let mut model = Mamba2SSM::new(config.clone(), &device)?;
    println!("  Model created");

    // Create input and target
    let input = Tensor::randn(0f64, 1.0, (8, 60, config.d_model), &device)?;
    let target = Tensor::randn(0f64, 1.0, (8, 60, 1), &device)?; // Output is [batch, seq, 1]
    println!("  Input/target created");

    // Forward pass
    let output = model.forward(&input)?;
    println!("  Forward pass complete");
    println!(
        "  Output shape: {:?}, Target shape: {:?}",
        output.dims(),
        target.dims()
    );

    // Compute loss (MSE)
    let diff = output.sub(&target)?;
    let squared = diff.sqr()?;
    let loss = squared.mean_all()?;

    let loss_value = loss.to_scalar::<f64>()?;
    println!("  Loss: {:.6}", loss_value);

    // Validate loss is reasonable
    assert!(
        loss_value.is_finite(),
        "Loss must be finite, got {}",
        loss_value
    );
    assert!(
        loss_value >= 0.0,
        "Loss must be non-negative, got {}",
        loss_value
    );

    println!("✅ Gradient flow test PASSED");
    Ok(())
}

#[tokio::test]
async fn test_mamba2_training_loop_simple() -> Result<()> {
    println!("🧪 E2E Test: MAMBA-2 Simple Training Loop (3 batches)");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("  Device: {:?}", device);

    let config = default_mamba2_config();
    let mut model = Mamba2SSM::new(config.clone(), &device)?;
    println!("  Model created");

    // Simulate 3 batches
    for batch_idx in 1..=3 {
        println!("  Batch {}/3", batch_idx);

        // Generate synthetic batch
        let input = Tensor::randn(0f64, 1.0, (16, 60, config.d_model), &device)?;
        let target = Tensor::randn(0f64, 1.0, (16, 60, 1), &device)?;

        // Forward pass
        let output = model.forward(&input)?;
        println!("    Output shape: {:?}", output.dims());

        // Compute loss
        let diff = output.sub(&target)?;
        let squared = diff.sqr()?;
        let loss = squared.mean_all()?;
        let loss_value = loss.to_scalar::<f64>()?;

        println!("    Loss: {:.6}", loss_value);
        assert!(loss_value.is_finite(), "Loss must be finite");
    }

    println!("✅ Training loop test PASSED");
    Ok(())
}

#[tokio::test]
async fn test_mamba2_config_variations() -> Result<()> {
    println!("🧪 E2E Test: MAMBA-2 Config Variations");

    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    println!("  Device: {:?}", device);

    // Test different configurations
    let configs = vec![("Small", 128, 2), ("Medium", 256, 4), ("Large", 512, 6)];

    for (name, d_model, num_layers) in configs {
        println!(
            "  Testing {} config: d_model={}, layers={}",
            name, d_model, num_layers
        );

        let mut config = default_mamba2_config();
        config.d_model = d_model;
        config.num_layers = num_layers;

        let mut model = Mamba2SSM::new(config.clone(), &device)?;
        let input = Tensor::randn(0f64, 1.0, (8, 60, d_model), &device)?;
        let output = model.forward(&input)?;

        assert_eq!(
            output.dims()[2],
            1,
            "Output should have 1 feature (regression)"
        );
        println!("    ✓ {} config works", name);
    }

    println!("✅ Config variation test PASSED");
    Ok(())
}