foxhunt/ml/tests/mamba2_p1_metrics_test.rs

//! P1 Metrics Tests for MAMBA-2
//!
//! Tests for the new directional accuracy, MAE, RMSE, and R² metrics.
//! This ensures that batch size bounds are correct and metrics are calculated properly.

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

/// Test directional accuracy calculation
///
/// Verifies that:
/// - Perfect predictions (100% correct direction) = 100% accuracy
/// - Random predictions (50% correct direction) = ~50% accuracy
/// - Inverse predictions (0% correct direction) = 0% accuracy
#[test]
fn test_directional_accuracy_perfect() {
    // Test data: predictions that perfectly match actual direction
    let predictions: Vec<f64> = vec![105.0, 98.0, 103.0, 97.0, 101.0];
    let targets: Vec<f64> = vec![105.5, 98.2, 103.3, 97.1, 101.4];
    let prev_prices: Vec<f64> = vec![100.0, 100.0, 100.0, 100.0, 100.0];

    let correct = predictions
        .iter()
        .zip(&targets)
        .zip(&prev_prices)
        .filter(|((&pred, &tgt), &prev)| {
            let pred_direction = ((pred - prev) as f64).signum();
            let actual_direction = ((tgt - prev) as f64).signum();
            (pred_direction - actual_direction).abs() < 0.01
        })
        .count();

    let directional_accuracy = correct as f64 / predictions.len() as f64;

    assert_eq!(
        directional_accuracy, 1.0,
        "Perfect predictions should have 100% directional accuracy"
    );
}

#[test]
fn test_directional_accuracy_inverse() {
    // Test data: predictions that are exactly opposite of actual direction
    let predictions: Vec<f64> = vec![95.0, 105.0, 97.0, 103.0, 99.0];
    let targets: Vec<f64> = vec![105.5, 98.2, 103.3, 97.1, 101.4];
    let prev_prices: Vec<f64> = vec![100.0, 100.0, 100.0, 100.0, 100.0];

    let correct = predictions
        .iter()
        .zip(&targets)
        .zip(&prev_prices)
        .filter(|((&pred, &tgt), &prev)| {
            let pred_direction = ((pred - prev) as f64).signum();
            let actual_direction = ((tgt - prev) as f64).signum();
            (pred_direction - actual_direction).abs() < 0.01
        })
        .count();

    let directional_accuracy = correct as f64 / predictions.len() as f64;

    assert_eq!(
        directional_accuracy, 0.0,
        "Inverse predictions should have 0% directional accuracy"
    );
}

#[test]
fn test_directional_accuracy_mixed() {
    // Test data: 4 correct, 1 incorrect
    // Pred 105→up, Target 105.5→up ✓
    // Pred 95→down, Target 98.2→down ✓
    // Pred 103→up, Target 103.3→up ✓
    // Pred 97→down, Target 103.1→up ✗
    // Pred 101→up, Target 101.4→up ✓
    let predictions: Vec<f64> = vec![105.0, 95.0, 103.0, 97.0, 101.0];
    let targets: Vec<f64> = vec![105.5, 98.2, 103.3, 103.1, 101.4];
    let prev_prices: Vec<f64> = vec![100.0, 100.0, 100.0, 100.0, 100.0];

    let correct = predictions
        .iter()
        .zip(&targets)
        .zip(&prev_prices)
        .filter(|((&pred, &tgt), &prev)| {
            let pred_direction = ((pred - prev) as f64).signum();
            let actual_direction = ((tgt - prev) as f64).signum();
            (pred_direction - actual_direction).abs() < 0.01
        })
        .count();

    let directional_accuracy = correct as f64 / predictions.len() as f64;

    assert_eq!(
        directional_accuracy, 0.8,
        "4/5 correct should be 80% directional accuracy"
    );
}

/// Test MAE calculation
#[test]
fn test_mae_calculation() {
    let predictions: Vec<f64> = vec![100.0, 105.0, 110.0, 95.0, 102.0];
    let targets: Vec<f64> = vec![102.0, 103.0, 108.0, 97.0, 100.0];

    let mae = predictions
        .iter()
        .zip(&targets)
        .map(|(&p, &t)| ((p - t) as f64).abs())
        .sum::<f64>()
        / predictions.len() as f64;

    // Expected: (2 + 2 + 2 + 2 + 2) / 5 = 2.0
    assert_eq!(mae, 2.0, "MAE should be 2.0");
}

#[test]
fn test_mae_zero() {
    // Perfect predictions should have MAE = 0
    let predictions: Vec<f64> = vec![100.0, 105.0, 110.0, 95.0, 102.0];
    let targets = predictions.clone();

    let mae = predictions
        .iter()
        .zip(&targets)
        .map(|(&p, &t)| ((p - t) as f64).abs())
        .sum::<f64>()
        / predictions.len() as f64;

    assert_eq!(mae, 0.0, "Perfect predictions should have MAE = 0");
}

/// Test RMSE calculation
#[test]
fn test_rmse_calculation() {
    let predictions: Vec<f64> = vec![100.0, 105.0, 110.0, 95.0, 102.0];
    let targets: Vec<f64> = vec![102.0, 103.0, 108.0, 97.0, 100.0];

    let mse = predictions
        .iter()
        .zip(&targets)
        .map(|(&p, &t)| ((p - t) as f64).powi(2))
        .sum::<f64>()
        / predictions.len() as f64;
    let rmse = mse.sqrt();

    // Expected: sqrt((4 + 4 + 4 + 4 + 4) / 5) = sqrt(4) = 2.0
    assert_eq!(rmse, 2.0, "RMSE should be 2.0");
}

#[test]
fn test_rmse_vs_mae() {
    // RMSE should always be >= MAE
    let predictions: Vec<f64> = vec![100.0, 110.0, 90.0, 105.0, 95.0];
    let targets: Vec<f64> = vec![102.0, 108.0, 92.0, 107.0, 97.0];

    let mae = predictions
        .iter()
        .zip(&targets)
        .map(|(&p, &t)| ((p - t) as f64).abs())
        .sum::<f64>()
        / predictions.len() as f64;

    let mse = predictions
        .iter()
        .zip(&targets)
        .map(|(&p, &t)| ((p - t) as f64).powi(2))
        .sum::<f64>()
        / predictions.len() as f64;
    let rmse = mse.sqrt();

    assert!(rmse >= mae, "RMSE ({}) should be >= MAE ({})", rmse, mae);
}

/// Test R² calculation
#[test]
fn test_r_squared_perfect() {
    // Perfect predictions should have R² = 1.0
    let predictions: Vec<f64> = vec![100.0, 105.0, 110.0, 95.0, 102.0];
    let targets = predictions.clone();

    let target_mean = targets.iter().sum::<f64>() / targets.len() as f64;
    let ss_tot: f64 = targets.iter().map(|t| (t - target_mean).powi(2)).sum();
    let ss_res: f64 = predictions
        .iter()
        .zip(&targets)
        .map(|(p, t)| (t - p).powi(2))
        .sum();

    let r_squared = 1.0 - (ss_res / ss_tot);

    assert_eq!(r_squared, 1.0, "Perfect predictions should have R² = 1.0");
}

#[test]
fn test_r_squared_mean_model() {
    // Predicting the mean should give R² = 0.0
    let targets: Vec<f64> = vec![100.0, 105.0, 110.0, 95.0, 102.0];
    let target_mean = targets.iter().sum::<f64>() / targets.len() as f64;
    let predictions: Vec<f64> = vec![target_mean; targets.len()];

    let ss_tot: f64 = targets.iter().map(|t| (t - target_mean).powi(2)).sum();
    let ss_res: f64 = predictions
        .iter()
        .zip(&targets)
        .map(|(p, t)| (t - p).powi(2))
        .sum();

    let r_squared = 1.0 - (ss_res / ss_tot);

    assert!(
        (r_squared - 0.0).abs() < 1e-10,
        "Mean predictions should have R² ≈ 0.0"
    );
}

#[test]
fn test_r_squared_worse_than_mean() {
    // Terrible predictions should have R² < 0
    let targets: Vec<f64> = vec![100.0, 105.0, 110.0, 95.0, 102.0];
    let predictions: Vec<f64> = vec![200.0, 250.0, 210.0, 195.0, 202.0];

    let target_mean = targets.iter().sum::<f64>() / targets.len() as f64;
    let ss_tot: f64 = targets.iter().map(|t| (t - target_mean).powi(2)).sum();
    let ss_res: f64 = predictions
        .iter()
        .zip(&targets)
        .map(|(p, t)| (t - p).powi(2))
        .sum();

    let r_squared = 1.0 - (ss_res / ss_tot);

    assert!(r_squared < 0.0, "Terrible predictions should have R² < 0");
}

/// Test batch size validation
#[test]
fn test_batch_size_bounds() {
    use ml::hyperopt::adapters::mamba2::Mamba2Params;
    use ml::hyperopt::traits::ParameterSpace;

    let bounds = Mamba2Params::continuous_bounds();
    let batch_size_bounds = bounds[1]; // batch_size is second parameter

    assert_eq!(
        batch_size_bounds,
        (4.0, 64.0),
        "Batch size bounds should be (4, 64)"
    );
}

#[test]
fn test_batch_size_max_vs_dataset_size() {
    // Max batch size (64) should be <= 60% of typical dataset size (108)
    let typical_dataset_size = 108;
    let max_batch_size = 64;
    let ratio = max_batch_size as f64 / typical_dataset_size as f64;

    assert!(
        ratio <= 0.6,
        "Max batch size should be <= 60% of dataset size"
    );
    assert!(max_batch_size >= 4, "Min batch size should be >= 4");
}

#[test]
fn test_batch_size_allows_multiple_batches() {
    // With dataset size 108 and batch size 64, we should have at least 1 full batch
    let dataset_size = 108;
    let batch_size = 64;
    let num_batches = (dataset_size + batch_size - 1) / batch_size;

    assert!(num_batches >= 1, "Should have at least 1 batch");

    // With min batch size 4, we should have many batches
    let min_batch_size = 4;
    let num_batches_min = dataset_size / min_batch_size;
    assert!(
        num_batches_min >= 27,
        "Min batch size should allow 27+ batches per epoch"
    );
}

/// Integration test: Full metric calculation with MAMBA-2
#[tokio::test]
async fn test_mamba2_metrics_integration() -> Result<(), Box<dyn std::error::Error>> {
    use ml::mamba::{Mamba2Config, OptimizerType};

    let device = Device::Cpu;

    // Small model for fast testing
    let config = Mamba2Config {
        d_model: 10,
        d_state: 4,
        d_head: 2,
        num_heads: 2,
        expand: 2,
        num_layers: 1,
        dropout: 0.0,
        use_ssd: false,
        use_selective_state: false,
        hardware_aware: false,
        target_latency_us: 1000,
        max_seq_len: 10,
        learning_rate: 0.001,
        weight_decay: 0.0,
        grad_clip: 1.0,
        warmup_steps: 0,
        adam_beta1: 0.9,
        adam_beta2: 0.999,
        adam_epsilon: 1e-8,
        total_decay_steps: 1000,
        batch_size: 4, // Use new minimum
        seq_len: 5,
        shuffle_batches: false,
        optimizer_type: OptimizerType::Adam,
        sgd_momentum: 0.9,
        sequence_stride: 1,
        norm_eps: 1e-5,
    };

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

    // Create dummy training data
    let mut train_data = Vec::new();
    let mut val_data = Vec::new();

    for i in 0..20 {
        let input = Tensor::zeros((1, 5, 10), candle_core::DType::F64, &device)?;
        let target = Tensor::new(&[100.0 + i as f64], &device)?.reshape((1, 1, 1))?;

        if i < 16 {
            train_data.push((input.clone(), target.clone()));
        } else {
            val_data.push((input.clone(), target.clone()));
        }
    }

    // Train for 2 epochs
    let history = model.train(&train_data, &val_data, 2, None).await?;

    // Verify metrics exist and are reasonable
    assert_eq!(history.len(), 2, "Should have 2 epochs");

    for epoch in &history {
        assert!(epoch.train_loss >= 0.0, "Train loss should be non-negative");
        assert!(epoch.val_loss >= 0.0, "Val loss should be non-negative");
        assert!(
            epoch.directional_accuracy >= 0.0 && epoch.directional_accuracy <= 1.0,
            "Directional accuracy should be in [0, 1]"
        );
        assert!(epoch.mae >= 0.0, "MAE should be non-negative");
        assert!(epoch.rmse >= 0.0, "RMSE should be non-negative");
        assert!(epoch.rmse >= epoch.mae, "RMSE should be >= MAE");
        // R² can be negative for bad models, so just check it's not NaN
        assert!(!epoch.r_squared.is_nan(), "R² should not be NaN");
    }

    Ok(())
}

/// Test that train_loss and val_loss are tracked separately
#[tokio::test]
async fn test_separate_train_val_loss() -> Result<(), Box<dyn std::error::Error>> {
    use ml::mamba::{Mamba2Config, OptimizerType};

    let device = Device::Cpu;

    let config = Mamba2Config {
        d_model: 10,
        d_state: 4,
        d_head: 2,
        num_heads: 2,
        expand: 2,
        num_layers: 1,
        dropout: 0.0,
        use_ssd: false,
        use_selective_state: false,
        hardware_aware: false,
        target_latency_us: 1000,
        max_seq_len: 10,
        learning_rate: 0.001,
        weight_decay: 0.0,
        grad_clip: 1.0,
        warmup_steps: 0,
        adam_beta1: 0.9,
        adam_beta2: 0.999,
        adam_epsilon: 1e-8,
        total_decay_steps: 1000,
        batch_size: 4,
        seq_len: 5,
        shuffle_batches: false,
        optimizer_type: OptimizerType::Adam,
        sgd_momentum: 0.9,
        sequence_stride: 1,
        norm_eps: 1e-5,
    };

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

    // Create different train/val data to ensure losses differ
    let mut train_data = Vec::new();
    let mut val_data = Vec::new();

    for i in 0..20 {
        let input = Tensor::zeros((1, 5, 10), candle_core::DType::F64, &device)?;
        let target = Tensor::new(&[100.0 + i as f64], &device)?.reshape((1, 1, 1))?;

        if i < 16 {
            train_data.push((input, target));
        } else {
            // Make validation targets different
            let val_target = Tensor::new(&[200.0 + i as f64], &device)?.reshape((1, 1, 1))?;
            val_data.push((input, val_target));
        }
    }

    let history = model.train(&train_data, &val_data, 2, None).await?;

    // Verify train_loss and val_loss are both present and tracked
    for epoch in &history {
        assert!(epoch.train_loss.is_finite(), "Train loss should be finite");
        assert!(epoch.val_loss.is_finite(), "Val loss should be finite");

        // They should be different (though not necessarily for all models)
        // Just verify they're both being calculated
        println!(
            "Epoch {}: Train Loss = {:.6}, Val Loss = {:.6}",
            epoch.epoch, epoch.train_loss, epoch.val_loss
        );
    }

    Ok(())
}