foxhunt/ml/tests/test_streaming_loader.rs

//! Integration tests for StreamingDbnLoader
//!
//! Tests memory-efficient streaming data loading with real DBN files.

use anyhow::Result;
use ml::data_loaders::{DbnSequenceLoader, StreamingDbnLoader};
use std::path::PathBuf;

/// Test data directory (small dataset with 4 files)
const TEST_DATA_DIR: &str = "test_data/real/databento/ml_training_small";

#[tokio::test]
async fn test_streaming_loader_creation() -> Result<()> {
    let loader = StreamingDbnLoader::new(60, 256).await?;
    println!("✅ StreamingDbnLoader created successfully: {:?}", loader);
    Ok(())
}

#[tokio::test]
async fn test_custom_config() -> Result<()> {
    let loader = StreamingDbnLoader::with_config(60, 256, 5000, 50).await?;
    println!("✅ Custom config applied: {:?}", loader);
    Ok(())
}

#[tokio::test]
async fn test_stream_sequences_small_dataset() -> Result<()> {
    let test_dir = PathBuf::from(TEST_DATA_DIR);

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

    let loader = StreamingDbnLoader::with_config(60, 256, 1000, 10).await?;
    let mut stream = loader.stream_sequences(&test_dir, 0.9).await?;

    let mut total_sequences = 0;
    let mut batch_count = 0;

    // Process all batches
    loop {
        match stream.next_batch().await? {
            Some(batch) => {
                batch_count += 1;
                total_sequences += batch.len();

                // Verify batch contents
                assert!(!batch.is_empty(), "Batch should not be empty");

                for (input, target) in &batch {
                    // Verify tensor shapes
                    assert_eq!(input.dims().len(), 2, "Input should be 2D");
                    assert_eq!(input.dims()[0], 60, "Sequence length should be 60");
                    assert_eq!(input.dims()[1], 256, "Feature dim should be 256");

                    assert_eq!(target.dims().len(), 2, "Target should be 2D");
                    assert_eq!(target.dims()[0], 1, "Target batch size should be 1");
                    assert_eq!(target.dims()[1], 256, "Target dim should be 256");
                }

                println!("   Batch {}: {} sequences", batch_count, batch.len());
            },
            None => break,
        }
    }

    println!(
        "✅ Streamed {} sequences in {} batches",
        total_sequences, batch_count
    );
    assert!(total_sequences > 0, "Should load at least some sequences");
    assert!(batch_count > 0, "Should have at least one batch");

    Ok(())
}

#[tokio::test]
async fn test_streaming_vs_batch_consistency() -> Result<()> {
    let test_dir = PathBuf::from(TEST_DATA_DIR);

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

    // Load with batch loader
    let mut batch_loader = DbnSequenceLoader::with_limits(60, 256, Some(100), 10).await?;
    let (batch_train, batch_val) = batch_loader.load_sequences(&test_dir, 0.9).await?;
    let batch_total = batch_train.len() + batch_val.len();

    println!("   Batch loader: {} sequences", batch_total);

    // Load with streaming loader (same config)
    let streaming_loader = StreamingDbnLoader::with_config(60, 256, 1000, 10).await?;
    let mut stream = streaming_loader.stream_sequences(&test_dir, 0.9).await?;

    let mut streaming_total = 0;
    loop {
        match stream.next_batch().await? {
            Some(batch) => streaming_total += batch.len(),
            None => break,
        }
    }

    println!("   Streaming loader: {} sequences", streaming_total);

    // Should produce similar number of sequences (within 10% due to boundary effects)
    let diff_ratio = (batch_total as f64 - streaming_total as f64).abs() / batch_total as f64;
    assert!(
        diff_ratio < 0.1,
        "Sequence count should be similar (diff: {:.1}%)",
        diff_ratio * 100.0
    );

    println!(
        "✅ Batch and streaming loaders produce consistent results (diff: {:.1}%)",
        diff_ratio * 100.0
    );

    Ok(())
}

#[tokio::test]
async fn test_memory_efficiency() -> Result<()> {
    let test_dir = PathBuf::from(TEST_DATA_DIR);

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

    // Get baseline memory
    let baseline = get_memory_usage_mb()?;
    println!("   Baseline memory: {:.1} MB", baseline);

    // Load with streaming
    let loader = StreamingDbnLoader::with_config(60, 256, 1000, 10).await?;
    let mut stream = loader.stream_sequences(&test_dir, 0.9).await?;

    let mut max_memory = baseline;

    // Process batches and track peak memory
    loop {
        match stream.next_batch().await? {
            Some(_batch) => {
                let current = get_memory_usage_mb()?;
                if current > max_memory {
                    max_memory = current;
                }
            },
            None => break,
        }
    }

    let peak_memory = max_memory - baseline;
    println!("   Peak memory delta: {:.1} MB", peak_memory);

    // For small dataset, peak should be < 100MB
    assert!(
        peak_memory < 100.0,
        "Peak memory should be < 100MB for small dataset, got {:.1} MB",
        peak_memory
    );

    println!("✅ Memory efficiency verified: {:.1} MB peak", peak_memory);

    Ok(())
}

#[tokio::test]
async fn test_train_val_split() -> Result<()> {
    let test_dir = PathBuf::from(TEST_DATA_DIR);

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

    let loader = StreamingDbnLoader::with_config(60, 256, 1000, 10).await?;
    let mut stream = loader.stream_sequences(&test_dir, 0.8).await?;

    // Count training sequences
    let mut train_count = 0;
    loop {
        match stream.next_batch().await? {
            Some(batch) => train_count += batch.len(),
            None => break,
        }
    }

    println!("   Training sequences: {}", train_count);

    // Switch to validation
    stream.switch_to_validation().await?;

    // Count validation sequences
    let mut val_count = 0;
    loop {
        match stream.next_batch().await? {
            Some(batch) => val_count += batch.len(),
            None => break,
        }
    }

    println!("   Validation sequences: {}", val_count);

    // Verify split ratio is approximately correct (within 20% due to boundary effects)
    let total = train_count + val_count;
    let train_ratio = train_count as f64 / total as f64;

    let split_error = (train_ratio - 0.8).abs();
    assert!(
        split_error < 0.2,
        "Train/val split should be approximately 80/20, got {:.1}%/{:.1}%",
        train_ratio * 100.0,
        (1.0 - train_ratio) * 100.0
    );

    println!(
        "✅ Train/val split verified: {:.1}%/{:.1}%",
        train_ratio * 100.0,
        (1.0 - train_ratio) * 100.0
    );

    Ok(())
}

#[tokio::test]
async fn test_different_batch_sizes() -> Result<()> {
    let test_dir = PathBuf::from(TEST_DATA_DIR);

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

    // Test with different batch sizes
    let batch_sizes = vec![1000, 5000, 10000];

    for batch_size in batch_sizes {
        let loader = StreamingDbnLoader::with_config(60, 256, batch_size, 10).await?;
        let mut stream = loader.stream_sequences(&test_dir, 0.9).await?;

        let mut total = 0;
        loop {
            match stream.next_batch().await? {
                Some(batch) => total += batch.len(),
                None => break,
            }
        }

        println!("   Batch size {}: {} total sequences", batch_size, total);
        assert!(
            total > 0,
            "Should load sequences with batch_size={}",
            batch_size
        );
    }

    println!("✅ All batch sizes work correctly");

    Ok(())
}

#[tokio::test]
async fn test_different_strides() -> Result<()> {
    let test_dir = PathBuf::from(TEST_DATA_DIR);

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

    // Test with different strides
    let strides = vec![1, 10, 50, 100];

    for stride in strides {
        let loader = StreamingDbnLoader::with_config(60, 256, 1000, stride).await?;
        let mut stream = loader.stream_sequences(&test_dir, 0.9).await?;

        let mut total = 0;
        loop {
            match stream.next_batch().await? {
                Some(batch) => total += batch.len(),
                None => break,
            }
        }

        println!("   Stride {}: {} total sequences", stride, total);
        assert!(total > 0, "Should load sequences with stride={}", stride);
    }

    println!("✅ All strides work correctly");

    Ok(())
}

/// Get current memory usage in MB
fn get_memory_usage_mb() -> Result<f64> {
    let status = std::fs::read_to_string("/proc/self/status")?;

    for line in status.lines() {
        if line.starts_with("VmRSS:") {
            let kb: usize = line
                .split_whitespace()
                .nth(1)
                .and_then(|s| s.parse().ok())
                .unwrap_or(0);
            return Ok(kb as f64 / 1024.0);
        }
    }

    Ok(0.0)
}