foxhunt/ml/tests/test_dbn_sequence_256_features.rs

//! Test DbnSequenceLoader produces correct 256-dimensional features
//!
//! Validates that the fixed DbnSequenceLoader correctly extracts and pads
//! features to exactly 256 dimensions for MAMBA-2 training.

use anyhow::Result;
use candle_core::IndexOp;
use ml::data_loaders::DbnSequenceLoader;
use std::env;
use std::path::PathBuf;

/// Get test data directory path
fn get_test_data_dir() -> PathBuf {
    // Try CARGO_MANIFEST_DIR first (works in tests)
    if let Ok(manifest_dir) = env::var("CARGO_MANIFEST_DIR") {
        PathBuf::from(manifest_dir)
            .parent()
            .unwrap()
            .join("test_data/real/databento/ml_training_small")
    } else {
        // Fallback to relative path from project root
        PathBuf::from("test_data/real/databento/ml_training_small")
    }
}

#[tokio::test]
async fn test_feature_dimension_256() -> Result<()> {
    println!("🔍 Testing DbnSequenceLoader 256-dimensional features...\n");

    let test_dir = get_test_data_dir();

    if !test_dir.exists() {
        println!("⚠️  Test data not found at {:?}, skipping test", test_dir);
        return Ok(());
    }

    // Create loader with d_model=256
    let mut loader = DbnSequenceLoader::with_limits(60, 256, Some(10), 10).await?;
    println!("✅ Created DbnSequenceLoader (seq_len=60, d_model=256, max=10, stride=10)\n");

    // Load sequences
    println!("📖 Loading sequences from {:?}...", test_dir);
    let (train_data, val_data) = loader.load_sequences(&test_dir, 0.9).await?;

    let total_sequences = train_data.len() + val_data.len();
    println!(
        "✅ Loaded {} sequences ({} train, {} val)\n",
        total_sequences,
        train_data.len(),
        val_data.len()
    );

    // Verify at least some data was loaded
    assert!(total_sequences > 0, "Should load at least some sequences");

    // Test 1: Verify input tensor dimensions
    println!("📊 Test 1: Verifying input tensor dimensions...");
    for (idx, (input, _target)) in train_data.iter().take(5).enumerate() {
        let input_dims = input.dims();

        println!("   Sequence {}: input shape = {:?}", idx, input_dims);

        // Input should be [batch=1, seq_len=60, d_model=256]
        assert_eq!(
            input_dims.len(),
            3,
            "Input should be 3D (batch, seq_len, features), got {:?}",
            input_dims
        );
        assert_eq!(
            input_dims[0], 1,
            "Batch dimension should be 1, got {}",
            input_dims[0]
        );
        assert_eq!(
            input_dims[1], 60,
            "Sequence length should be 60, got {}",
            input_dims[1]
        );
        assert_eq!(
            input_dims[2], 256,
            "Feature dimension should be 256, got {}",
            input_dims[2]
        );
    }
    println!("✅ All input tensors have correct shape [1, 60, 256]\n");

    // Test 2: Verify target tensor dimensions
    println!("📊 Test 2: Verifying target tensor dimensions...");
    for (idx, (_input, target)) in train_data.iter().take(5).enumerate() {
        let target_dims = target.dims();

        println!("   Sequence {}: target shape = {:?}", idx, target_dims);

        // Target should be [batch=1, timesteps=1, d_model=256]
        assert_eq!(
            target_dims.len(),
            3,
            "Target should be 3D, got {:?}",
            target_dims
        );
        assert_eq!(
            target_dims[0], 1,
            "Target batch should be 1, got {}",
            target_dims[0]
        );
        assert_eq!(
            target_dims[1], 1,
            "Target timesteps should be 1, got {}",
            target_dims[1]
        );
        assert_eq!(
            target_dims[2], 256,
            "Target feature dim should be 256, got {}",
            target_dims[2]
        );
    }
    println!("✅ All target tensors have correct shape [1, 1, 256]\n");

    // Test 3: Verify feature values are normalized (not all zeros/NaN)
    println!("📊 Test 3: Verifying feature normalization...");
    let (first_input, _) = &train_data[0];
    let flattened = first_input.flatten_all()?;
    let values = flattened.to_vec1::<f64>()?;

    // Check for NaN values
    let nan_count = values.iter().filter(|v| v.is_nan()).count();
    assert_eq!(nan_count, 0, "Found {} NaN values in features", nan_count);
    println!("   ✅ No NaN values detected");

    // Check for all-zero sequences (should have some variation)
    let non_zero_count = values.iter().filter(|v| v.abs() > 1e-6).count();
    let non_zero_ratio = non_zero_count as f64 / values.len() as f64;
    println!(
        "   ✅ Non-zero values: {}/{} ({:.1}%)",
        non_zero_count,
        values.len(),
        non_zero_ratio * 100.0
    );

    assert!(
        non_zero_ratio > 0.01,
        "Features appear to be all zeros (only {:.1}% non-zero)",
        non_zero_ratio * 100.0
    );

    // Check value range (normalized features should be roughly in [-5, 5] range)
    let min_val = values.iter().cloned().fold(f64::INFINITY, f64::min);
    let max_val = values.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
    let mean = values.iter().sum::<f64>() / values.len() as f64;

    println!(
        "   ✅ Value range: [{:.4}, {:.4}], mean: {:.4}",
        min_val, max_val, mean
    );

    // Normalized features should not have extreme outliers
    assert!(
        min_val > -100.0 && max_val < 100.0,
        "Feature values seem unnormalized: range [{:.2}, {:.2}]",
        min_val,
        max_val
    );

    println!("✅ Features are properly normalized\n");

    // Test 4: Verify validation data has same properties
    println!("📊 Test 4: Verifying validation data...");
    if !val_data.is_empty() {
        let (val_input, val_target) = &val_data[0];

        assert_eq!(
            val_input.dims(),
            &[1, 60, 256],
            "Validation input should be [1, 60, 256]"
        );
        assert_eq!(
            val_target.dims(),
            &[1, 1, 256],
            "Validation target should be [1, 1, 256]"
        );

        println!("   ✅ Validation data shapes correct");
        println!("   ✅ {} validation sequences verified", val_data.len());
    } else {
        println!("   ⚠️  No validation data (split ratio may be too high)");
    }

    println!("\n✅ ALL TESTS PASSED!");
    println!("   - Feature dimension: ✅ 256");
    println!("   - Input shape: ✅ [1, 60, 256]");
    println!("   - Target shape: ✅ [1, 1, 256]");
    println!("   - Normalization: ✅ Valid");
    println!(
        "   - Total sequences: {} ({} train, {} val)",
        total_sequences,
        train_data.len(),
        val_data.len()
    );

    Ok(())
}

#[tokio::test]
async fn test_extract_features_dimension() -> Result<()> {
    println!("🔍 Testing extract_features() returns 256 dimensions...\n");

    let test_dir = get_test_data_dir();

    if !test_dir.exists() {
        println!("⚠️  Test data not found, skipping test");
        return Ok(());
    }

    // Create loader
    let mut loader = DbnSequenceLoader::new(60, 256).await?;
    println!("✅ Created DbnSequenceLoader\n");

    // Load sequences
    let (train_data, _) = loader.load_sequences(&test_dir, 0.9).await?;

    assert!(!train_data.is_empty(), "Should have training data");

    // Get first sequence and verify it was created with 256-dim features
    let (input, _) = &train_data[0];

    // Input is [1, 60, 256] where 256 is the feature dimension
    let feature_dim = input.dims()[2];

    println!("📊 Feature dimension from tensor: {}", feature_dim);
    assert_eq!(
        feature_dim, 256,
        "Feature dimension should be 256, got {}",
        feature_dim
    );

    println!("✅ extract_features() correctly produces 256-dimensional features\n");

    Ok(())
}

#[tokio::test]
async fn test_different_d_model_values() -> Result<()> {
    println!("🔍 Testing different d_model values (128, 256, 512)...\n");

    let test_dir = get_test_data_dir();

    if !test_dir.exists() {
        println!("⚠️  Test data not found, skipping test");
        return Ok(());
    }

    // Test different d_model values
    let d_models = vec![128, 256, 512];

    for d_model in d_models {
        println!("📊 Testing d_model={}...", d_model);

        let mut loader = DbnSequenceLoader::with_limits(60, d_model, Some(5), 10).await?;
        let (train_data, _) = loader.load_sequences(&test_dir, 0.9).await?;

        if !train_data.is_empty() {
            let (input, target) = &train_data[0];

            // Verify input shape
            assert_eq!(
                input.dims()[2],
                d_model,
                "Input feature dim should be {}",
                d_model
            );

            // Verify target shape
            assert_eq!(
                target.dims()[2],
                d_model,
                "Target feature dim should be {}",
                d_model
            );

            println!(
                "   ✅ d_model={}: input={:?}, target={:?}",
                d_model,
                input.dims(),
                target.dims()
            );
        }
    }

    println!("\n✅ All d_model values produce correct dimensions\n");

    Ok(())
}

#[tokio::test]
async fn test_sequence_temporal_ordering() -> Result<()> {
    println!("🔍 Testing temporal ordering of sequences...\n");

    let test_dir = get_test_data_dir();

    if !test_dir.exists() {
        println!("⚠️  Test data not found, skipping test");
        return Ok(());
    }

    // Create loader with stride=1 to get consecutive sequences
    let mut loader = DbnSequenceLoader::with_limits(10, 256, Some(3), 1).await?;
    let (train_data, _) = loader.load_sequences(&test_dir, 0.9).await?;

    if train_data.len() >= 2 {
        println!("📊 Comparing consecutive sequences...");

        let (seq1_input, _) = &train_data[0];
        let (seq2_input, _) = &train_data[1];

        // With stride=1, the second sequence should be a shifted version of the first
        // seq1: [t0, t1, t2, ..., t9]
        // seq2: [t1, t2, t3, ..., t10]

        // Extract last 9 timesteps from seq1
        let seq1_last_9 = seq1_input.i((0, 1..10, ..))?;

        // Extract first 9 timesteps from seq2
        let seq2_first_9 = seq2_input.i((0, 0..9, ..))?;

        // These should be identical (temporal ordering)
        let diff = (seq1_last_9 - seq2_first_9)?;
        let diff_flat = diff.abs()?.flatten_all()?;
        let diff_vec = diff_flat.to_vec1::<f64>()?;
        let max_diff = diff_vec.iter().cloned().fold(f64::NEG_INFINITY, f64::max);

        println!(
            "   ✅ Max difference between overlapping windows: {:.6}",
            max_diff
        );

        assert!(
            max_diff < 1e-6,
            "Consecutive sequences should overlap with stride=1, max_diff={}",
            max_diff
        );
    }

    println!("✅ Temporal ordering verified\n");

    Ok(())
}

#[tokio::test]
async fn test_batch_processing() -> Result<()> {
    println!("🔍 Testing batch processing with multiple sequences...\n");

    let test_dir = get_test_data_dir();

    if !test_dir.exists() {
        println!("⚠️  Test data not found, skipping test");
        return Ok(());
    }

    // Load multiple sequences
    let mut loader = DbnSequenceLoader::with_limits(60, 256, Some(100), 10).await?;
    let (train_data, val_data) = loader.load_sequences(&test_dir, 0.8).await?;

    let total = train_data.len() + val_data.len();
    println!("📊 Loaded {} sequences", total);

    // Verify all sequences have consistent dimensions
    let mut valid_count = 0;
    for (input, target) in train_data.iter().chain(val_data.iter()) {
        if input.dims() == &[1, 60, 256] && target.dims() == &[1, 1, 256] {
            valid_count += 1;
        }
    }

    println!(
        "   ✅ {}/{} sequences have correct dimensions",
        valid_count, total
    );
    assert_eq!(
        valid_count, total,
        "All sequences should have correct dimensions"
    );

    println!("✅ Batch processing verified\n");

    Ok(())
}