foxhunt/ml/tests/mamba2_training_pipeline_test.rs

//! MAMBA-2 Training Pipeline Tests (TDD - Agent 10.6)
//!
//! Test-Driven Development for MAMBA-2 training pipeline targeting 70.6% loss reduction.
//!
//! ## Test Structure (TDD)
//! 1. RED: Write tests first (they should FAIL)
//! 2. GREEN: Implement minimum code to pass tests
//! 3. REFACTOR: Improve quality while keeping tests passing
//!
//! ## Test Coverage
//! - Training on ES.FUT data
//! - Loss reduction >50% (test), >70% (production)
//! - SSM forward pass correctness
//! - B/C matrix shape validation (d_inner)
//! - Checkpoint saving/loading
//! - GPU training compatibility

use anyhow::Result;
use candle_core::{DType, Device, Tensor};
use ml::data_loaders::DbnSequenceLoader;
use ml::mamba::{Mamba2Config, Mamba2SSM};
use std::path::PathBuf;

/// Test configuration for fast unit tests
fn test_config() -> Mamba2Config {
    Mamba2Config {
        d_model: 256,
        d_state: 16,
        d_head: 32,
        num_heads: 8,
        expand: 4,
        num_layers: 2, // Fewer layers for faster tests
        dropout: 0.1,
        use_ssd: true,
        use_selective_state: true,
        hardware_aware: true,
        target_latency_us: 5,
        max_seq_len: 60,
        learning_rate: 0.0001,
        weight_decay: 1e-4,
        grad_clip: 1.0,
        warmup_steps: 10,
        batch_size: 4, // Small batch for tests
        seq_len: 60,
    }
}

/// RED: Test MAMBA-2 trains on ES.FUT data
///
/// This test SHOULD FAIL initially because we haven't implemented
/// the training pipeline yet.
///
/// Success criteria:
/// - Loads ES.FUT data successfully
/// - Trains for 20 epochs
/// - Loss reduction >50%
/// - Best loss tracked correctly
#[tokio::test]
async fn test_mamba2_trains_on_es_fut() -> Result<()> {
    // Arrange: Load ES.FUT data
    let data_dir = PathBuf::from("test_data/real/databento/ml_training_small");

    // Skip if test data not available
    if !data_dir.exists() {
        eprintln!("⚠️  Skipping test: {} not found", data_dir.display());
        return Ok(());
    }

    let mut loader = DbnSequenceLoader::new(60, 256).await?;
    let (train_data, val_data) = loader.load_sequences(&data_dir, 0.8).await?;

    assert!(!train_data.is_empty(), "Training data should not be empty");
    assert!(!val_data.is_empty(), "Validation data should not be empty");

    // Act: Train MAMBA-2 model
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let config = test_config();
    let mut model = Mamba2SSM::new(config, &device)?;

    // Train for 20 epochs (fast test)
    let epochs = 20;
    let training_history = model.train(&train_data, &val_data, epochs, None).await?;

    // Assert: Verify training results
    assert_eq!(
        training_history.len(),
        epochs,
        "Should have 20 training epochs"
    );

    // Loss reduction >50%
    let initial_loss = training_history[0].loss;
    let final_loss = training_history.last().unwrap().loss;
    let loss_reduction = (initial_loss - final_loss) / initial_loss;

    assert!(
        loss_reduction > 0.5,
        "Loss reduction should be >50%, got {:.2}%",
        loss_reduction * 100.0
    );

    // Best loss should be tracked
    let best_loss = training_history
        .iter()
        .map(|e| e.loss)
        .fold(f64::INFINITY, f64::min);
    assert!(
        best_loss < initial_loss,
        "Best loss should improve from initial"
    );

    println!("✅ MAMBA-2 trained on ES.FUT:");
    println!("   Initial loss: {:.6}", initial_loss);
    println!("   Final loss: {:.6}", final_loss);
    println!("   Loss reduction: {:.2}%", loss_reduction * 100.0);

    Ok(())
}

/// RED: Test SSM forward pass produces correct shapes
///
/// Tests that SSM state space model forward pass produces
/// expected output dimensions.
#[tokio::test]
async fn test_ssm_forward_pass_shapes() -> Result<()> {
    // Arrange
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let config = test_config();
    let mut model = Mamba2SSM::new(config.clone(), &device)?;

    let batch_size = 2;
    let seq_len = 60;
    let d_model = config.d_model;

    // Create random input [batch, seq, d_model]
    let input = Tensor::randn(0.0f32, 1.0f32, (batch_size, seq_len, d_model), &device)?
        .to_dtype(DType::F64)?;

    // Act: Forward pass
    let output = model.forward(&input)?;

    // Assert: Output shape should be [batch, seq, output_dim=1]
    let output_dims = output.dims();
    assert_eq!(output_dims.len(), 3, "Output should be 3D tensor");
    assert_eq!(output_dims[0], batch_size, "Batch dimension mismatch");
    assert_eq!(output_dims[1], seq_len, "Sequence dimension mismatch");
    assert_eq!(
        output_dims[2], 1,
        "Output dimension should be 1 (regression)"
    );

    println!(
        "✅ SSM forward pass: {:?} → {:?}",
        input.dims(),
        output.dims()
    );

    Ok(())
}

/// RED: Test B/C matrix shapes use d_inner (not d_model)
///
/// Critical fix from Wave 160: B/C matrices should use d_inner
/// dimension after input projection expansion.
#[tokio::test]
async fn test_bc_matrix_shapes_use_d_inner() -> Result<()> {
    // Arrange
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let config = test_config();
    let model = Mamba2SSM::new(config.clone(), &device)?;

    let d_state = config.d_state;
    let d_inner = config.d_model * config.expand; // CRITICAL: d_inner = d_model * expand

    // Assert: B matrix should be [d_state, d_inner]
    let B = &model.state.ssm_states[0].B;
    assert_eq!(B.dims().len(), 2, "B should be 2D matrix");
    assert_eq!(B.dims()[0], d_state, "B first dimension should be d_state");
    assert_eq!(
        B.dims()[1],
        d_inner,
        "B second dimension should be d_inner (NOT d_model)"
    );

    // Assert: C matrix should be [d_inner, d_state]
    let C = &model.state.ssm_states[0].C;
    assert_eq!(C.dims().len(), 2, "C should be 2D matrix");
    assert_eq!(
        C.dims()[0],
        d_inner,
        "C first dimension should be d_inner (NOT d_model)"
    );
    assert_eq!(C.dims()[1], d_state, "C second dimension should be d_state");

    println!("✅ B/C matrix shapes correct:");
    println!("   d_model: {}", config.d_model);
    println!("   d_inner: {} (d_model * expand)", d_inner);
    println!(
        "   B shape: {:?} (expected [{}, {}])",
        B.dims(),
        d_state,
        d_inner
    );
    println!(
        "   C shape: {:?} (expected [{}, {}])",
        C.dims(),
        d_inner,
        d_state
    );

    Ok(())
}

/// RED: Test checkpoint saving and loading
///
/// Verifies that model checkpoints can be saved and restored.
#[tokio::test]
async fn test_checkpoint_save_and_load() -> Result<()> {
    // Arrange
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let config = test_config();
    let mut model = Mamba2SSM::new(config, &device)?;

    let checkpoint_path = "ml/checkpoints/test_mamba2.ckpt";

    // Act: Save checkpoint
    model.save_checkpoint(checkpoint_path).await?;
    assert!(
        model.metadata.last_checkpoint.is_some(),
        "Checkpoint path should be recorded"
    );

    // Load checkpoint
    let mut loaded_model = Mamba2SSM::new(test_config(), &device)?;
    loaded_model.load_checkpoint(checkpoint_path).await?;

    // Assert: Model should be marked as trained
    assert!(
        loaded_model.is_trained,
        "Loaded model should be marked as trained"
    );
    assert_eq!(
        loaded_model.metadata.last_checkpoint.as_deref(),
        Some(checkpoint_path),
        "Checkpoint path should match"
    );

    println!("✅ Checkpoint save/load working");

    Ok(())
}

/// RED: Test GPU training compatibility
///
/// Ensures that model can be trained on CUDA device without errors.
#[tokio::test]
async fn test_gpu_training_compatibility() -> Result<()> {
    // Skip if CUDA not available
    let device = match Device::new_cuda(0) {
        Ok(d) => d,
        Err(_) => {
            eprintln!("⚠️  Skipping GPU test: CUDA not available");
            return Ok(());
        },
    };

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

    // Create small training set on GPU
    let batch_size = 2;
    let seq_len = 60;
    let mut train_data = Vec::new();

    for _ in 0..10 {
        let input = Tensor::randn(0.0f32, 1.0f32, (1, seq_len, config.d_model), &device)?
            .to_dtype(DType::F64)?;
        let target = Tensor::randn(0.0f32, 1.0f32, (1, 1, 1), &device)?.to_dtype(DType::F64)?;
        train_data.push((input, target));
    }

    let val_data = train_data.clone();

    // Act: Train on GPU for 5 epochs
    let training_history = model.train(&train_data, &val_data, 5, None).await?;

    // Assert
    assert_eq!(training_history.len(), 5, "Should complete 5 epochs on GPU");
    assert!(
        training_history[0].loss.is_finite(),
        "Loss should be finite"
    );

    println!(
        "✅ GPU training compatible: {} epochs completed",
        training_history.len()
    );

    Ok(())
}

/// RED: Test loss computation correctness
///
/// Verifies MSE loss is computed correctly for regression.
#[tokio::test]
async fn test_loss_computation() -> Result<()> {
    // Arrange
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let config = test_config();
    let model = Mamba2SSM::new(config, &device)?;

    // Create simple tensors for MSE calculation
    let output = Tensor::new(&[1.0f64, 2.0f64, 3.0f64], &device)?.reshape((1, 3, 1))?;
    let target = Tensor::new(&[1.5f64, 2.5f64, 2.5f64], &device)?.reshape((1, 3, 1))?;

    // Act: Compute loss
    let loss = model.compute_loss(&output, &target)?;
    let loss_value = loss.to_scalar::<f64>()?;

    // Assert: MSE = mean((output - target)^2)
    // Differences: [-0.5, -0.5, 0.5]
    // Squared: [0.25, 0.25, 0.25]
    // Mean: 0.25
    let expected_mse = 0.25;
    let tolerance = 1e-6;

    assert!(
        (loss_value - expected_mse).abs() < tolerance,
        "MSE loss should be {}, got {}",
        expected_mse,
        loss_value
    );

    println!("✅ Loss computation correct: MSE = {:.6}", loss_value);

    Ok(())
}

/// RED: Test gradient flow through SSM layers
///
/// Ensures gradients propagate correctly through state space model.
#[tokio::test]
async fn test_gradient_flow() -> Result<()> {
    // Arrange
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let config = test_config();
    let mut model = Mamba2SSM::new(config.clone(), &device)?;

    // Create single training example
    let input =
        Tensor::randn(0.0f32, 1.0f32, (1, 60, config.d_model), &device)?.to_dtype(DType::F64)?;
    let target = Tensor::randn(0.0f32, 1.0f32, (1, 1, 1), &device)?.to_dtype(DType::F64)?;

    // Act: Forward + backward pass
    model.zero_gradients()?;
    let output = model.forward_with_gradients(&input)?;

    // Extract last timestep for loss
    let seq_len = output.dim(1)?;
    let output_last = output.narrow(1, seq_len - 1, 1)?;

    let loss = model.compute_loss(&output_last, &target)?;
    model.backward_pass(&loss, &input, &target)?;

    // Assert: Gradients should exist for SSM parameters
    assert!(!model.gradients.is_empty(), "Gradients should be computed");

    // Check that A, B, C, delta gradients exist for first layer
    let has_a_grad = model.gradients.contains_key("A_0");
    let has_b_grad = model.gradients.contains_key("B_0");
    let has_c_grad = model.gradients.contains_key("C_0");
    let has_delta_grad = model.gradients.contains_key("delta_0");

    assert!(has_a_grad, "A matrix gradient should exist");
    assert!(has_b_grad, "B matrix gradient should exist");
    assert!(has_c_grad, "C matrix gradient should exist");
    assert!(has_delta_grad, "Delta parameter gradient should exist");

    println!("✅ Gradient flow verified through SSM layers");

    Ok(())
}

/// RED: Test optimizer updates SSM parameters
///
/// Verifies Adam optimizer correctly updates A, B, C, delta parameters.
#[tokio::test]
async fn test_optimizer_updates_parameters() -> Result<()> {
    // Arrange
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);
    let config = test_config();
    let mut model = Mamba2SSM::new(config.clone(), &device)?;

    // Initialize optimizer
    model.initialize_optimizer()?;

    // Store original parameters
    let A_original = model.state.ssm_states[0].A.clone();
    let B_original = model.state.ssm_states[0].B.clone();

    // Create gradients (simulated) - use broadcast_mul for scalar multiplication
    let scale_scalar = Tensor::new(&[0.01f64], &device)?.reshape(&[])?; // 0-D scalar

    let A_grad = Tensor::ones((config.d_state, config.d_state), DType::F64, &device)?
        .broadcast_mul(&scale_scalar)?;
    model.gradients.insert("A_0".to_string(), A_grad);

    let B_grad = Tensor::ones(
        (config.d_state, config.d_model * config.expand),
        DType::F64,
        &device,
    )?
    .broadcast_mul(&scale_scalar)?;
    model.gradients.insert("B_0".to_string(), B_grad);

    let C_grad = Tensor::ones(
        (config.d_model * config.expand, config.d_state),
        DType::F64,
        &device,
    )?
    .broadcast_mul(&scale_scalar)?;
    model.gradients.insert("C_0".to_string(), C_grad);

    let delta_grad =
        Tensor::ones((config.d_model,), DType::F64, &device)?.broadcast_mul(&scale_scalar)?;
    model.gradients.insert("delta_0".to_string(), delta_grad);

    // Act: Run optimizer step
    model.optimizer_step()?;

    // Assert: Parameters should have changed
    let A_updated = &model.state.ssm_states[0].A;
    let B_updated = &model.state.ssm_states[0].B;

    // Check that parameters differ (optimizer applied updates)
    let A_diff = (A_updated - &A_original)?
        .abs()?
        .sum_all()?
        .to_scalar::<f64>()?;
    let B_diff = (B_updated - &B_original)?
        .abs()?
        .sum_all()?
        .to_scalar::<f64>()?;

    assert!(A_diff > 1e-8, "A matrix should be updated by optimizer");
    assert!(B_diff > 1e-8, "B matrix should be updated by optimizer");

    println!("✅ Optimizer updates SSM parameters:");
    println!("   A parameter change: {:.6}", A_diff);
    println!("   B parameter change: {:.6}", B_diff);

    Ok(())
}

/// Production test: Full 200-epoch training (marked as ignored by default)
///
/// Run with: cargo test -p ml --test mamba2_training_pipeline_test -- --ignored
#[tokio::test]
#[ignore = "Long-running test - run explicitly with --ignored"]
async fn test_mamba2_production_training_200_epochs() -> Result<()> {
    // Load ES.FUT data
    let data_dir = PathBuf::from("test_data/real/databento/ml_training_small");
    if !data_dir.exists() {
        eprintln!(
            "⚠️  Skipping production test: {} not found",
            data_dir.display()
        );
        return Ok(());
    }

    let mut loader = DbnSequenceLoader::new(60, 256).await?;
    let (train_data, val_data) = loader.load_sequences(&data_dir, 0.8).await?;

    // Production configuration
    let config = Mamba2Config {
        d_model: 256,
        d_state: 16,
        d_head: 32,
        num_heads: 8,
        expand: 4,
        num_layers: 6, // Full model
        dropout: 0.1,
        use_ssd: true,
        use_selective_state: true,
        hardware_aware: true,
        target_latency_us: 5,
        max_seq_len: 60,
        learning_rate: 0.0001,
        weight_decay: 1e-4,
        grad_clip: 1.0,
        warmup_steps: 1000,
        batch_size: 32,
        seq_len: 60,
    };

    let device = Device::new_cuda(0).expect("CUDA required for production training");
    let mut model = Mamba2SSM::new(config, &device)?;

    // Train for 200 epochs
    println!("🚀 Starting 200-epoch production training...");
    let training_history = model.train(&train_data, &val_data, 200, None).await?;

    // Assert: Loss reduction >70% (Wave 160 benchmark)
    let initial_loss = training_history[0].loss;
    let final_loss = training_history.last().unwrap().loss;
    let loss_reduction = (initial_loss - final_loss) / initial_loss;

    assert!(
        loss_reduction > 0.70,
        "Production training should achieve >70% loss reduction, got {:.2}%",
        loss_reduction * 100.0
    );

    // Save final checkpoint
    model
        .save_checkpoint("ml/checkpoints/mamba2_es_fut_v1.safetensors")
        .await?;

    println!("✅ Production training complete:");
    println!("   Initial loss: {:.6}", initial_loss);
    println!("   Final loss: {:.6}", final_loss);
    println!("   Loss reduction: {:.2}%", loss_reduction * 100.0);
    println!("   Checkpoint: ml/checkpoints/mamba2_es_fut_v1.safetensors");

    Ok(())
}