foxhunt/ml/tests/mamba2_hyperopt_p0_p1_fixes.rs

//! MAMBA-2 Hyperparameter Optimization P0/P1 Fix Tests
//!
//! Tests for critical issues fixed in mamba2.rs:
//! 1. P0: NaN panic in sorting (line 524)
//! 2. P0: Division by zero tolerance (lines 507, 546)
//! 3. P1: Empty parquet validation (line 414)
//! 4. P1: Validation size check (lines 582-584)
//! 5. P1: CUDA OOM handling (lines 748-800)

use anyhow::Result;
use std::sync::Arc;
use tempfile::NamedTempFile;

use ml::hyperopt::adapters::mamba2::{Mamba2Params, Mamba2Trainer};
use ml::hyperopt::traits::HyperparameterOptimizable;

use arrow::array::{Float64Array, PrimitiveArray, UInt64Array};
use arrow::datatypes::{DataType, Field, Schema, TimestampNanosecondType};
use arrow::record_batch::RecordBatch;
use parquet::arrow::arrow_writer::ArrowWriter;
use parquet::file::properties::WriterProperties;

/// Helper: Create parquet file with market data
fn create_test_parquet(num_rows: usize, base_price: f64) -> Result<NamedTempFile> {
    let temp_file = NamedTempFile::new()?;
    let file = temp_file.as_file().try_clone()?;

    // DBN schema (10 columns)
    let schema = Arc::new(Schema::new(vec![
        Field::new("ts_recv", DataType::UInt64, false),
        Field::new("publisher_id", DataType::UInt8, false),
        Field::new("instrument_id", DataType::UInt16, false),
        Field::new("open", DataType::Float64, false),
        Field::new("high", DataType::Float64, false),
        Field::new("low", DataType::Float64, false),
        Field::new("close", DataType::Float64, false),
        Field::new("volume", DataType::UInt64, false),
        Field::new("symbol", DataType::Utf8, false),
        Field::new(
            "ts_event",
            DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None),
            false,
        ),
    ]));

    let props = WriterProperties::builder().build();
    let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props))?;

    let base_timestamp = 1_700_000_000_000_000_000u64;

    let batch = RecordBatch::try_new(
        schema,
        vec![
            Arc::new(UInt64Array::from(
                (0..num_rows)
                    .map(|i| base_timestamp + i as u64 * 60_000_000_000)
                    .collect::<Vec<_>>(),
            )),
            Arc::new(arrow::array::UInt8Array::from(vec![1u8; num_rows])),
            Arc::new(arrow::array::UInt16Array::from(vec![1u16; num_rows])),
            Arc::new(Float64Array::from(
                (0..num_rows)
                    .map(|i| base_price + (i as f64 * 0.1))
                    .collect::<Vec<_>>(),
            )),
            Arc::new(Float64Array::from(
                (0..num_rows)
                    .map(|i| base_price + (i as f64 * 0.1) + 5.0)
                    .collect::<Vec<_>>(),
            )),
            Arc::new(Float64Array::from(
                (0..num_rows)
                    .map(|i| base_price + (i as f64 * 0.1) - 5.0)
                    .collect::<Vec<_>>(),
            )),
            Arc::new(Float64Array::from(
                (0..num_rows)
                    .map(|i| base_price + (i as f64 * 0.1) + 2.5)
                    .collect::<Vec<_>>(),
            )),
            Arc::new(UInt64Array::from(vec![1000u64; num_rows])),
            Arc::new(arrow::array::StringArray::from(vec!["ES.FUT"; num_rows])),
            Arc::new(PrimitiveArray::<TimestampNanosecondType>::from(
                (0..num_rows)
                    .map(|i| base_timestamp as i64 + i as i64 * 60_000_000_000)
                    .collect::<Vec<_>>(),
            )),
        ],
    )?;

    writer.write(&batch)?;
    writer.close()?;

    Ok(temp_file)
}

#[test]
fn test_p0_nan_panic_in_sorting() {
    // P0: Test that NaN values don't cause panic in sorted_features.sort_by()
    // Fix: Line 524 changed from .unwrap() to .unwrap_or(std::cmp::Ordering::Equal)

    // Create dataset with constant prices (produces zero variance, edge case)
    let temp_file = create_test_parquet(150, 5000.0).unwrap();

    let mut trainer = Mamba2Trainer::new(temp_file.path(), 1)
        .unwrap()
        .with_train_split(0.8);

    let params = Mamba2Params::default();

    // Should NOT panic with NaN values (if they occur)
    let result = trainer.train_with_params(params);

    assert!(
        result.is_ok() || result.is_err(),
        "Training should handle NaN gracefully without panic"
    );
}

#[test]
fn test_p0_division_by_zero_tolerance() {
    // P0: Test that division by zero tolerance is appropriate (1e-6 instead of 1e-10)
    // Fix: Lines 507 and 546 changed tolerance from 1e-10 to 1e-6

    // Create dataset with very small price range
    let temp_file = create_test_parquet(150, 5000.0).unwrap();

    let mut trainer = Mamba2Trainer::new(temp_file.path(), 1)
        .unwrap()
        .with_train_split(0.8);

    let params = Mamba2Params::default();

    let result = trainer.train_with_params(params);

    match result {
        Ok(metrics) => {
            // If succeeds, metrics should be finite
            assert!(
                metrics.val_loss.is_finite(),
                "Val loss should be finite, got: {}",
                metrics.val_loss
            );
        },
        Err(e) => {
            // Should get clear error about zero variance OR validation set too small
            // (validation check happens first if dataset is tiny)
            let msg = e.to_string();
            assert!(
                msg.contains("zero variance")
                    || msg.contains("normalize")
                    || msg.contains("Validation set too small")
                    || msg.contains("Insufficient"),
                "Expected zero variance or validation size error, got: {}",
                msg
            );
        },
    }
}

#[test]
fn test_p1_empty_parquet_validation() {
    // P1: Test that empty parquet files are rejected early
    // Fix: Lines 486-491 added validation before feature extraction

    // Create empty parquet
    let temp_file = create_test_parquet(0, 5000.0).unwrap();

    let mut trainer = Mamba2Trainer::new(temp_file.path(), 1)
        .unwrap()
        .with_train_split(0.8);

    let params = Mamba2Params::default();

    let result = trainer.train_with_params(params);

    assert!(result.is_err(), "Empty parquet should be rejected");

    let error_msg = result.unwrap_err().to_string();
    assert!(
        error_msg.contains("Insufficient data")
            || error_msg.contains("No features")
            || error_msg.contains("Empty"),
        "Expected clear error about empty data, got: {}",
        error_msg
    );
}

#[test]
fn test_p1_validation_size_check() {
    // P1: Test that datasets too small for validation split are rejected
    // Fix: Lines 574-577 added validation set size check

    // Create dataset with only 5 rows (too small for 60-bar sequence + validation)
    let temp_file = create_test_parquet(5, 5000.0).unwrap();

    let mut trainer = Mamba2Trainer::new(temp_file.path(), 1)
        .unwrap()
        .with_train_split(0.8);

    let params = Mamba2Params::default();

    let result = trainer.train_with_params(params);

    // Should fail with clear error or return penalty metrics
    assert!(
        result.is_err() || (result.is_ok() && result.as_ref().unwrap().val_loss >= 1000.0),
        "Tiny dataset should be rejected or return penalty metrics"
    );

    if let Err(e) = result {
        let msg = e.to_string();
        assert!(
            msg.contains("Insufficient")
                || msg.contains("too small")
                || msg.contains("Validation set"),
            "Expected error about insufficient data, got: {}",
            msg
        );
    }
}

#[test]
fn test_p1_cuda_oom_handling() {
    // P1: Test that CUDA OOM errors are handled gracefully (return penalty, not panic)
    // Fix: Lines 748-800 wrap training in catch_unwind

    // Create small dataset
    let temp_file = create_test_parquet(100, 5000.0).unwrap();

    let mut trainer = Mamba2Trainer::new(temp_file.path(), 1)
        .unwrap()
        .with_train_split(0.8);

    // Use massive batch size to potentially trigger OOM
    let mut params = Mamba2Params::default();
    params.batch_size = 100000; // Impossibly large

    // Should either clamp batch size or return penalty (NOT panic)
    let result = trainer.train_with_params(params);

    assert!(
        result.is_ok() || result.is_err(),
        "OOM should be handled gracefully without panic"
    );

    if let Ok(metrics) = result {
        // If succeeds (batch size clamped), metrics should be valid
        assert!(
            metrics.val_loss.is_finite(),
            "Val loss should be finite, got: {}",
            metrics.val_loss
        );
    }
}

#[test]
fn test_constant_prices_zero_variance() {
    // Edge case: Constant prices (zero variance) should be rejected
    // This tests the tolerance fixes (lines 507, 546)

    // All prices identical (zero variance)
    let temp_file = create_test_parquet(100, 5000.0).unwrap();

    let mut trainer = Mamba2Trainer::new(temp_file.path(), 1)
        .unwrap()
        .with_train_split(0.8);

    let params = Mamba2Params::default();

    let result = trainer.train_with_params(params);

    // Should be rejected with clear error
    assert!(result.is_err(), "Constant prices should be rejected");

    let error_msg = result.unwrap_err().to_string();
    assert!(
        error_msg.contains("zero variance")
            || error_msg.contains("normalize")
            || error_msg.contains("Validation set too small")
            || error_msg.contains("Insufficient"),
        "Expected zero variance or validation size error, got: {}",
        error_msg
    );
}

#[test]
fn test_batch_size_clamping() {
    // Test that batch_size is clamped to configured bounds (prevents OOM)

    let temp_file = create_test_parquet(100, 5000.0).unwrap();

    // Set tight batch size bounds
    let mut trainer = Mamba2Trainer::new(temp_file.path(), 1)
        .unwrap()
        .with_batch_size_bounds(8.0, 16.0);

    // Try with batch size outside bounds
    let mut params = Mamba2Params::default();
    params.batch_size = 256; // Above max

    let result = trainer.train_with_params(params);

    // Should clamp and succeed (or fail gracefully)
    assert!(
        result.is_ok() || result.is_err(),
        "Batch size clamping should handle out-of-bounds values"
    );
}

#[test]
fn test_normalized_targets_bounds() {
    // Test that normalized targets stay in [0,1] range (critical for model)

    let prices: Vec<f64> = vec![
        5000.0, 5100.0, 5200.0, 5300.0, 5400.0, 5500.0, 5600.0, 5700.0, 5800.0, 5900.0, 6000.0,
    ];

    let min_price = prices.iter().copied().fold(f64::INFINITY, f64::min);
    let max_price = prices.iter().copied().fold(f64::NEG_INFINITY, f64::max);

    for &price in &prices {
        let normalized = (price - min_price) / (max_price - min_price);

        assert!(
            normalized >= 0.0 && normalized <= 1.0,
            "Normalized target {} out of [0,1] range for price {}",
            normalized,
            price
        );

        // Test denormalization
        let denormalized = normalized * (max_price - min_price) + min_price;
        assert!(
            (denormalized - price).abs() < 1e-6,
            "Denormalization failed: {} -> {} -> {}",
            price,
            normalized,
            denormalized
        );
    }
}