foxhunt/ml/tests/mamba2_hyperopt_edge_cases.rs

//! MAMBA2-Specific Edge Case Tests for Hyperparameter Optimization
//!
//! This test suite covers MAMBA2-specific edge cases:
//! 1. Async data loading edge cases
//! 2. Sequence length and stride edge cases
//! 3. Normalization parameter edge cases
//! 4. SSM-specific numerical stability
//! 5. Batch size clamping with GPU memory
//!
//! Purpose: Ensure MAMBA2 adapter handles all edge cases robustly

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

// ============================================================================
// TEST UTILITIES
// ============================================================================

fn create_test_parquet(temp_dir: &TempDir, num_rows: usize, suffix: &str) -> String {
    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;
    use std::fs::File;
    use std::sync::Arc;

    let schema = Arc::new(Schema::new(vec![
        Field::new("ts_event", DataType::UInt64, false),
        Field::new("rtype", DataType::UInt8, false),
        Field::new("publisher_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(
            "timestamp",
            DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None),
            false,
        ),
    ]));

    let file_path = temp_dir
        .path()
        .join(format!("mamba2_test_{}.parquet", suffix));
    let file = File::create(&file_path).unwrap();
    let props = WriterProperties::builder().build();
    let mut writer = ArrowWriter::try_new(file, schema.clone(), Some(props)).unwrap();

    let base_price = 5000.0;
    let base_timestamp = 1700000000_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<_>>(),
            )),
        ],
    )
    .unwrap();

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

    file_path.to_string_lossy().to_string()
}

// ============================================================================
// PARAMETER ROUNDTRIP TESTS
// ============================================================================

#[test]
fn test_all_12_params_roundtrip() {
    // Verify all 12 MAMBA2 parameters survive roundtrip conversion
    let params = Mamba2Params {
        learning_rate: 5e-5,
        batch_size: 64,
        dropout: 0.15,
        weight_decay: 5e-5,
        grad_clip: 2.0,
        warmup_steps: 500,
        adam_beta1: 0.9,
        adam_beta2: 0.999,
        adam_epsilon: 1e-8,
        lookback_window: 90,
        sequence_stride: 2,
        norm_eps: 1e-5,
    };

    let continuous = params.to_continuous();
    assert_eq!(continuous.len(), 12, "Should have 12 continuous parameters");

    let recovered = Mamba2Params::from_continuous(&continuous).expect("Failed to recover params");

    // Verify all parameters
    assert!((recovered.learning_rate - params.learning_rate).abs() < 1e-10);
    assert_eq!(recovered.batch_size, params.batch_size);
    assert!((recovered.dropout - params.dropout).abs() < 1e-10);
    assert!((recovered.weight_decay - params.weight_decay).abs() < 1e-10);
    assert!((recovered.grad_clip - params.grad_clip).abs() < 1e-6);
    assert_eq!(recovered.warmup_steps, params.warmup_steps);
    assert!((recovered.adam_beta1 - params.adam_beta1).abs() < 1e-10);
    assert!((recovered.adam_beta2 - params.adam_beta2).abs() < 1e-10);
    assert!((recovered.adam_epsilon - params.adam_epsilon).abs() < 1e-12);
    assert_eq!(recovered.lookback_window, params.lookback_window);
    assert_eq!(recovered.sequence_stride, params.sequence_stride);
    assert!((recovered.norm_eps - params.norm_eps).abs() < 1e-12);
}

#[test]
fn test_full_training_pipeline() {
    // End-to-end test: create data, train, denormalize predictions
    let temp_dir = TempDir::new().unwrap();
    let parquet_file = create_test_parquet(&temp_dir, 200, "full_pipeline");

    let mut trainer = Mamba2Trainer::new(&parquet_file, 10)
        .expect("Failed to create trainer")
        .with_batch_size_bounds(4.0, 32.0)
        .with_async_loading(true, 3)
        .with_train_split(0.8);

    let params = Mamba2Params::default();
    let result = trainer.train_with_params(params);

    assert!(result.is_ok(), "Full training pipeline should succeed");

    let metrics = result.unwrap();

    // Verify metrics are reasonable
    assert!(
        metrics.val_loss.is_finite(),
        "Validation loss should be finite"
    );
    assert!(
        metrics.val_loss >= 0.0,
        "Validation loss should be non-negative"
    );
    assert!(metrics.directional_accuracy >= 0.0 && metrics.directional_accuracy <= 1.0);
    assert!(metrics.mae >= 0.0);
    assert!(metrics.rmse >= 0.0);
    assert!(metrics.r_squared >= -1.0 && metrics.r_squared <= 1.0);
    assert_eq!(metrics.epochs_completed, 10);

    // Test denormalization
    let pred = trainer.denormalize_prediction(0.5);
    assert!(
        pred.is_finite() && pred > 0.0,
        "Denormalized prediction should be valid"
    );
}

// ============================================================================
// CHECKPOINT INTEGRITY TESTS (VarMap Registration)
// ============================================================================

#[test]
fn test_mamba2_checkpoint_saves_all_parameters() {
    use candle_core::{Device, Tensor};
    use ml::mamba::{Mamba2Config, Mamba2SSM};
    use std::collections::HashMap;

    // Small config for fast testing
    let config = Mamba2Config {
        d_model: 8,
        d_state: 4,
        d_head: 4,
        num_heads: 2,
        expand: 2,
        num_layers: 2,
        seq_len: 10,
        batch_size: 1,
        dropout: 0.0,
        norm_eps: 1e-5,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let device = Device::Cpu;
    let mut model = Mamba2SSM::new(config.clone(), &device).expect("Failed to create model");

    // Save checkpoint
    let temp_dir = TempDir::new().unwrap();
    let ckpt_path = temp_dir.path().join("mamba2_params.safetensors");

    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async { model.save_checkpoint(ckpt_path.to_str().unwrap()).await })
        .expect("Failed to save checkpoint");

    // Load and verify SSD layers are present
    let tensors: HashMap<String, Tensor> =
        candle_core::safetensors::load(&ckpt_path, &device).expect("Failed to load checkpoint");

    println!("\nCheckpoint tensors: {}", tensors.len());
    for name in tensors.keys() {
        println!("  {}", name);
    }

    // CRITICAL: Verify SSD layer parameters exist (this is the VarMap bug)
    for i in 0..config.num_layers {
        let qkv_key = format!("ssd_layer_{}.qkv_proj.weight", i);
        assert!(
            tensors.contains_key(&qkv_key),
            "CRITICAL BUG: {} missing from checkpoint! VarMap registration bug detected.",
            qkv_key
        );

        let out_key = format!("ssd_layer_{}.out_proj.weight", i);
        assert!(
            tensors.contains_key(&out_key),
            "CRITICAL BUG: {} missing from checkpoint! VarMap registration bug detected.",
            out_key
        );
    }

    // Count total parameters
    let mut total_params = 0;
    for tensor in tensors.values() {
        let param_count: usize = tensor.shape().dims().iter().product();
        total_params += param_count;
    }

    println!("Total parameters in checkpoint: {}", total_params);

    // Expected parameters (rough estimate)
    // Input proj: 8*16 + 16 = 144
    // Output proj: 16*1 + 1 = 17
    // Per layer: QKV(8*24+24=216) + Out(8*8+8=72) + State(8*4+4=36) + Gate(8*8+8=72) + LN(16*2=32) = 428
    // Total: 144 + 17 + (428*2) = 1017
    let expected_min_params = 800; // Conservative lower bound

    assert!(
        total_params >= expected_min_params,
        "Too few parameters in checkpoint: {} (expected >= {}). VarMap bug likely.",
        total_params,
        expected_min_params
    );
}

#[test]
fn test_mamba2_checkpoint_restore_determinism() {
    use candle_core::{Device, Tensor};
    use ml::mamba::{Mamba2Config, Mamba2SSM};

    // Small config for fast testing
    let config = Mamba2Config {
        d_model: 8,
        d_state: 4,
        d_head: 4,
        num_heads: 2,
        expand: 2,
        num_layers: 2,
        seq_len: 10,
        batch_size: 1,
        dropout: 0.0, // Disable for determinism
        norm_eps: 1e-5,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let device = Device::Cpu;
    let mut model1 = Mamba2SSM::new(config.clone(), &device).expect("Failed to create model");

    // Create test input
    let input_data: Vec<f64> = (0..80).map(|i| (i as f64) * 0.01).collect();
    let input = Tensor::from_vec(input_data, (1, 10, 8), &device).expect("Failed to create tensor");

    // Run inference BEFORE saving
    let output1 = model1.forward(&input).expect("Forward pass 1 failed");
    let output1_vec = output1
        .flatten_all()
        .expect("Flatten failed")
        .to_vec1::<f64>()
        .expect("to_vec1 failed");

    // Save checkpoint
    let temp_dir = TempDir::new().unwrap();
    let ckpt_path = temp_dir.path().join("mamba2_restore.safetensors");

    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async { model1.save_checkpoint(ckpt_path.to_str().unwrap()).await })
        .expect("Failed to save checkpoint");

    // Create NEW model and load checkpoint
    let mut model2 = Mamba2SSM::new(config.clone(), &device).expect("Failed to create model 2");

    rt.block_on(async { model2.load_checkpoint(ckpt_path.to_str().unwrap()).await })
        .expect("Failed to load checkpoint");

    // Run inference AFTER loading
    let output2 = model2.forward(&input).expect("Forward pass 2 failed");
    let output2_vec = output2
        .flatten_all()
        .expect("Flatten failed")
        .to_vec1::<f64>()
        .expect("to_vec1 failed");

    // CRITICAL: Outputs must be identical
    assert_eq!(
        output1_vec.len(),
        output2_vec.len(),
        "Output length mismatch"
    );

    let max_diff = output1_vec
        .iter()
        .zip(output2_vec.iter())
        .map(|(a, b)| (a - b).abs())
        .fold(0.0, f64::max);

    println!("Max output difference: {:.10e}", max_diff);

    assert!(
        max_diff < 1e-6,
        "Output mismatch after checkpoint restore! Max diff: {:.10e}\n\
         This indicates weights were not fully restored (VarMap bug).",
        max_diff
    );
}

#[test]
fn test_mamba2_checkpoint_size_reasonable() {
    use candle_core::Device;
    use ml::mamba::{Mamba2Config, Mamba2SSM};

    // Small config for fast testing
    let config = Mamba2Config {
        d_model: 8,
        d_state: 4,
        d_head: 4,
        num_heads: 2,
        expand: 2,
        num_layers: 2,
        seq_len: 10,
        batch_size: 1,
        dropout: 0.0,
        norm_eps: 1e-5,
        learning_rate: 1e-4,
        ..Default::default()
    };

    let device = Device::Cpu;
    let mut model = Mamba2SSM::new(config.clone(), &device).expect("Failed to create model");

    // Save checkpoint
    let temp_dir = TempDir::new().unwrap();
    let ckpt_path = temp_dir.path().join("mamba2_size.safetensors");

    let rt = tokio::runtime::Runtime::new().unwrap();
    rt.block_on(async { model.save_checkpoint(ckpt_path.to_str().unwrap()).await })
        .expect("Failed to save checkpoint");

    // Check file size
    let metadata = std::fs::metadata(&ckpt_path).expect("Failed to get metadata");
    let size_kb = metadata.len() as f64 / 1024.0;

    println!("Checkpoint size: {:.2} KB", size_kb);

    // File should be at least 5 KB (800+ parameters * 8 bytes = 6.4KB minimum)
    // If it's smaller, layers are missing
    assert!(
        size_kb > 5.0,
        "Checkpoint suspiciously small: {:.2} KB. VarMap bug likely (layers not saved).",
        size_kb
    );

    // Sanity check: shouldn't be huge either (max 100KB for this small model)
    assert!(
        size_kb < 100.0,
        "Checkpoint unexpectedly large: {:.2} KB. May indicate duplicate parameters.",
        size_kb
    );
}