foxhunt/ml/tests/test_dbn_parser_fix.rs

//! Test DBN Parser Fix - Verify OHLCV extraction from real DBN files
//!
//! This test validates that the official dbn crate decoder extracts all OHLCV bars
//! from DBN files, not just 2 messages (header metadata).

use anyhow::Result;

#[tokio::test]
async fn test_dqn_dbn_loading() -> Result<()> {
    use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};
    use std::path::Path;

    println!("Testing DQN DBN data loading with official decoder...");

    // Create DQN trainer
    let hyperparams = DQNHyperparameters {
        batch_size: 32, // Small batch for test
        epochs: 1,      // Only 1 epoch for test
        ..Default::default()
    };

    let trainer = DQNTrainer::new(hyperparams)?;
    println!("✓ DQN trainer created");

    // Load one DBN file directly
    let test_file =
        Path::new("test_data/real/databento/ml_training_small/6E.FUT_ohlcv-1m_2024-01-02.dbn");
    if !test_file.exists() {
        println!("⚠ Test file not found, skipping: {:?}", test_file);
        return Ok(());
    }

    // Use the fixed decoder
    let training_data = trainer.convert_dbn_file_to_training_data(test_file)?;

    println!(
        "✓ Loaded {} training samples from DBN file",
        training_data.len()
    );

    // Validate: Should extract 400-500+ OHLCV bars, not just 2 messages
    assert!(
        training_data.len() >= 100,
        "Expected at least 100 OHLCV bars, got {}. Custom parser only extracted 2 messages!",
        training_data.len()
    );

    println!(
        "✅ SUCCESS: DBN parser correctly extracted {} OHLCV bars",
        training_data.len()
    );
    println!("   (Previous custom parser only extracted 2 messages)");

    // Verify feature structure
    if !training_data.is_empty() {
        let (features, target) = &training_data[0];
        println!(
            "   - First bar features: {} prices, {} volumes, {} indicators",
            features.prices.len(),
            features.volumes.len(),
            features.technical_indicators.len()
        );
        println!("   - Target dimensions: {}", target.len());
    }

    Ok(())
}

#[tokio::test]
async fn test_dbn_sequence_loader() -> Result<()> {
    use ml::data_loaders::DbnSequenceLoader;
    use std::path::Path;

    println!("Testing MAMBA-2 DBN sequence loader with official decoder...");

    // Create sequence loader
    let mut loader = DbnSequenceLoader::new(60, 256).await?;
    println!("✓ Sequence loader created (seq_len=60, d_model=256)");

    // Load sequences from directory
    let test_dir = "test_data/real/databento/ml_training_small";
    let test_path = Path::new(test_dir);
    if !test_path.exists() {
        println!("⚠ Test directory not found, skipping: {:?}", test_dir);
        return Ok(());
    }

    let (train_data, val_data) = loader.load_sequences(test_dir, 0.9).await?;

    println!(
        "✓ Loaded {} training sequences, {} validation sequences",
        train_data.len(),
        val_data.len()
    );

    // Validate: Should create many sequences from 400-500+ OHLCV bars per file
    let total_sequences = train_data.len() + val_data.len();
    assert!(
        total_sequences >= 50,
        "Expected at least 50 sequences, got {}. Custom parser only extracted 2 messages per file!",
        total_sequences
    );

    println!(
        "✅ SUCCESS: Sequence loader created {} total sequences",
        total_sequences
    );
    println!("   (Previous custom parser only extracted 2 messages per file)");

    // Verify tensor shapes
    if !train_data.is_empty() {
        let (input, target) = &train_data[0];
        println!("   - Input shape: {:?}", input.shape());
        println!("   - Target shape: {:?}", target.shape());
    }

    Ok(())
}

#[tokio::test]
async fn test_dqn_serialization_fix() -> Result<()> {
    use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};

    println!("Testing DQN model serialization (SafeTensors)...");

    // Create DQN trainer with minimal config
    let hyperparams = DQNHyperparameters {
        learning_rate: 0.0001,
        batch_size: 32,
        gamma: 0.99,
        epsilon_start: 1.0,
        epsilon_end: 0.01,
        epsilon_decay: 0.995,
        buffer_size: 1000,
        epochs: 1,
        checkpoint_frequency: 1,
    };

    let trainer = DQNTrainer::new(hyperparams)?;
    println!("✓ DQN trainer created");

    // Serialize the model
    let checkpoint_data = trainer.serialize_model().await?;
    println!("✓ Model serialized: {} bytes", checkpoint_data.len());

    // CRITICAL VALIDATIONS:

    // 1. Not the old 1024-byte placeholder
    assert!(
        checkpoint_data.len() != 1024,
        "❌ FAIL: Still using hardcoded 1024-byte placeholder!"
    );
    println!("✓ Not the old placeholder");

    // 2. Should be at least 10KB (real Q-network weights)
    assert!(
        checkpoint_data.len() > 10_000,
        "❌ FAIL: Checkpoint too small ({} bytes), expected >10KB for Q-network weights",
        checkpoint_data.len()
    );
    println!(
        "✓ Checkpoint size realistic: {} bytes",
        checkpoint_data.len()
    );

    // 3. Not all zeros
    let is_all_zeros = checkpoint_data.iter().all(|&b| b == 0);
    assert!(!is_all_zeros, "❌ FAIL: Checkpoint is all zeros!");
    println!("✓ Contains non-zero data");

    // 4. SafeTensors format validation (8-byte header + JSON)
    assert!(
        checkpoint_data.len() >= 8,
        "❌ FAIL: Too small for SafeTensors format"
    );

    // SafeTensors starts with 8-byte little-endian header length
    let header_len = u64::from_le_bytes([
        checkpoint_data[0],
        checkpoint_data[1],
        checkpoint_data[2],
        checkpoint_data[3],
        checkpoint_data[4],
        checkpoint_data[5],
        checkpoint_data[6],
        checkpoint_data[7],
    ]);
    println!("✓ SafeTensors header length: {} bytes", header_len);

    assert!(
        header_len > 0 && header_len < checkpoint_data.len() as u64,
        "❌ FAIL: Invalid SafeTensors header length: {}",
        header_len
    );

    // 5. Verify JSON metadata exists
    let json_end = 8 + header_len as usize;
    if json_end <= checkpoint_data.len() {
        let json_bytes = &checkpoint_data[8..json_end];
        let json_str = std::str::from_utf8(json_bytes)?;
        println!("✓ SafeTensors JSON metadata: {} bytes", json_str.len());

        // Should contain tensor info
        assert!(
            json_str.contains("layer") || json_str.contains("weight") || json_str.contains("bias"),
            "❌ FAIL: JSON metadata doesn't contain expected tensor keys"
        );
        println!("✓ JSON contains tensor metadata");
    }

    println!("✅ SUCCESS: DQN serialization produces valid SafeTensors checkpoint");
    println!(
        "   Size: {} bytes ({}KB)",
        checkpoint_data.len(),
        checkpoint_data.len() / 1024
    );
    println!(
        "   Format: Valid SafeTensors with {}-byte JSON header",
        header_len
    );

    Ok(())
}