foxhunt/ml/tests/calibration_dataset_test.rs

//! TDD Test for Calibration Dataset Generation
//!
//! This test file is written FIRST following TDD methodology (RED-GREEN-REFACTOR).
//! It should FAIL initially until the calibration module is implemented.
//!
//! Mission: Generate 1,000-sample calibration dataset from ES.FUT data for INT8 quantization.
//!
//! Tests:
//! 1. test_generate_calibration_dataset() - Core functionality
//! 2. test_calibration_json_structure() - JSON format validation
//! 3. test_calibration_statistics() - Per-feature min/max/mean/std
//! 4. test_calibration_feature_count() - 26 features validation
//! 5. test_calibration_sample_count() - 1,000 samples validation
//! 6. test_load_calibration_data() - Load and validate saved JSON

use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;

/// Calibration dataset structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CalibrationDataset {
    /// Total number of samples
    pub sample_count: usize,
    
    /// Number of features per sample
    pub feature_count: usize,
    
    /// Symbol name
    pub symbol: String,
    
    /// Per-feature statistics
    pub feature_stats: Vec<FeatureStats>,
    
    /// Raw sample data (flattened: sample_count * feature_count)
    pub samples: Vec<f32>,
}

/// Per-feature statistics for quantization
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureStats {
    /// Feature index
    pub index: usize,
    
    /// Feature name
    pub name: String,
    
    /// Minimum value
    pub min: f32,
    
    /// Maximum value
    pub max: f32,
    
    /// Mean value
    pub mean: f32,
    
    /// Standard deviation
    pub std: f32,
}

/// Get test data directory
fn get_test_data_dir() -> PathBuf {
    if let Ok(manifest_dir) = std::env::var("CARGO_MANIFEST_DIR") {
        PathBuf::from(manifest_dir).parent().unwrap().join("test_data/real/databento")
    } else {
        PathBuf::from("test_data/real/databento")
    }
}

/// Get calibration output path
fn get_calibration_output_path() -> PathBuf {
    if let Ok(manifest_dir) = std::env::var("CARGO_MANIFEST_DIR") {
        PathBuf::from(manifest_dir).join("calibration")
    } else {
        PathBuf::from("ml/calibration")
    }
}

//
// Test 1: Core functionality - Generate calibration dataset
//
#[tokio::test]
async fn test_generate_calibration_dataset() -> Result<()> {
    println!("🧪 Test 1: Generate calibration dataset (CORE FUNCTIONALITY)\n");
    
    let test_dir = get_test_data_dir();
    let es_fut_file = test_dir.join("ES.FUT_ohlcv-1m_2024-01-02.dbn");
    
    if !es_fut_file.exists() {
        println!("⚠️  ES.FUT data not found at {:?}, skipping test", es_fut_file);
        return Ok(());
    }
    
    // Import the function we're testing (will fail until implemented)
    use ml::data_loaders::calibration::generate_calibration_dataset;
    
    let output_dir = get_calibration_output_path();
    std::fs::create_dir_all(&output_dir)?;
    
    let output_path = output_dir.join("es_fut_calibration.json");
    
    println!("📂 Input:  {:?}", es_fut_file);
    println!("📂 Output: {:?}", output_path);
    println!();
    
    // Generate calibration dataset (THIS WILL FAIL - RED PHASE)
    println!("🔄 Generating calibration dataset...");
    let dataset = generate_calibration_dataset(
        &es_fut_file,
        1000,  // 1,000 samples
        "ES.FUT"
    ).await?;
    
    println!("✅ Generated dataset:");
    println!("   Samples:  {}", dataset.sample_count);
    println!("   Features: {}", dataset.feature_count);
    println!("   Symbol:   {}", dataset.symbol);
    println!();
    
    // Validate basic properties
    assert_eq!(dataset.sample_count, 1000, "Should generate 1,000 samples");
    assert_eq!(dataset.symbol, "ES.FUT", "Symbol should be ES.FUT");
    assert!(dataset.feature_count > 0, "Should have features");
    assert_eq!(dataset.samples.len(), dataset.sample_count * dataset.feature_count,
               "Samples array size should match sample_count * feature_count");
    
    // Save to JSON
    println!("💾 Saving to JSON...");
    let json = serde_json::to_string_pretty(&dataset)?;
    std::fs::write(&output_path, json)?;
    println!("✅ Saved to {:?}", output_path);
    
    // Verify file exists
    assert!(output_path.exists(), "JSON file should exist");
    
    let file_size = std::fs::metadata(&output_path)?.len();
    println!("📊 File size: {} bytes ({:.2} KB)", file_size, file_size as f64 / 1024.0);
    
    Ok(())
}

//
// Test 2: JSON structure validation
//
#[tokio::test]
async fn test_calibration_json_structure() -> Result<()> {
    println!("🧪 Test 2: JSON structure validation\n");
    
    let output_path = get_calibration_output_path().join("es_fut_calibration.json");
    
    if !output_path.exists() {
        println!("⚠️  Calibration JSON not found, run test_generate_calibration_dataset first");
        return Ok(());
    }
    
    // Load JSON
    let json_str = std::fs::read_to_string(&output_path)?;
    let dataset: CalibrationDataset = serde_json::from_str(&json_str)?;
    
    println!("✅ JSON structure:");
    println!("   sample_count:  {}", dataset.sample_count);
    println!("   feature_count: {}", dataset.feature_count);
    println!("   symbol:        {}", dataset.symbol);
    println!("   feature_stats: {} entries", dataset.feature_stats.len());
    println!("   samples:       {} values", dataset.samples.len());
    println!();
    
    // Validate structure
    assert_eq!(dataset.sample_count, 1000, "Should have 1,000 samples");
    assert_eq!(dataset.feature_stats.len(), dataset.feature_count,
               "Should have stats for each feature");
    
    // Validate feature stats structure
    for (idx, stats) in dataset.feature_stats.iter().take(3).enumerate() {
        println!("   Feature {}: {} (min={:.4}, max={:.4}, mean={:.4}, std={:.4})",
                 stats.index, stats.name, stats.min, stats.max, stats.mean, stats.std);
        
        assert_eq!(stats.index, idx, "Feature index should match position");
        assert!(!stats.name.is_empty(), "Feature name should not be empty");
        assert!(stats.min <= stats.max, "Min should be <= max");
        assert!(stats.std >= 0.0, "Std should be non-negative");
    }
    
    println!("✅ JSON structure valid");
    Ok(())
}

//
// Test 3: Calibration statistics validation
//
#[tokio::test]
async fn test_calibration_statistics() -> Result<()> {
    println!("🧪 Test 3: Calibration statistics validation\n");
    
    let output_path = get_calibration_output_path().join("es_fut_calibration.json");
    
    if !output_path.exists() {
        println!("⚠️  Calibration JSON not found, skipping test");
        return Ok(());
    }
    
    let json_str = std::fs::read_to_string(&output_path)?;
    let dataset: CalibrationDataset = serde_json::from_str(&json_str)?;
    
    println!("📊 Validating per-feature statistics...\n");
    
    // Validate each feature's statistics
    for stats in &dataset.feature_stats {
        // Extract feature values from samples
        let mut values = Vec::new();
        for sample_idx in 0..dataset.sample_count {
            let value_idx = sample_idx * dataset.feature_count + stats.index;
            values.push(dataset.samples[value_idx]);
        }
        
        // Compute actual min/max/mean/std
        let actual_min = values.iter().cloned().fold(f32::INFINITY, f32::min);
        let actual_max = values.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let actual_mean = values.iter().sum::<f32>() / values.len() as f32;
        let actual_var = values.iter()
            .map(|v| (v - actual_mean).powi(2))
            .sum::<f32>() / values.len() as f32;
        let actual_std = actual_var.sqrt();
        
        // Validate (with tolerance for floating point precision)
        let tolerance = 1e-4;
        assert!((stats.min - actual_min).abs() < tolerance,
                "Feature {} min mismatch: stored={}, actual={}", 
                stats.index, stats.min, actual_min);
        assert!((stats.max - actual_max).abs() < tolerance,
                "Feature {} max mismatch: stored={}, actual={}", 
                stats.index, stats.max, actual_max);
        assert!((stats.mean - actual_mean).abs() < tolerance,
                "Feature {} mean mismatch: stored={}, actual={}", 
                stats.index, stats.mean, actual_mean);
        assert!((stats.std - actual_std).abs() < tolerance,
                "Feature {} std mismatch: stored={}, actual={}", 
                stats.index, stats.std, actual_std);
    }
    
    println!("✅ All feature statistics validated");
    println!("   {} features checked", dataset.feature_stats.len());
    println!("   1,000 samples per feature");
    
    Ok(())
}

//
// Test 4: Feature count validation (26 features expected)
//
#[tokio::test]
async fn test_calibration_feature_count() -> Result<()> {
    println!("🧪 Test 4: Feature count validation\n");
    
    let output_path = get_calibration_output_path().join("es_fut_calibration.json");
    
    if !output_path.exists() {
        println!("⚠️  Calibration JSON not found, skipping test");
        return Ok(());
    }
    
    let json_str = std::fs::read_to_string(&output_path)?;
    let dataset: CalibrationDataset = serde_json::from_str(&json_str)?;
    
    println!("📊 Feature count: {}", dataset.feature_count);
    println!();
    
    // Expected: 5 OHLCV + 10 technical indicators + 11 derived = 26 features
    // Or: 256 features if using full MAMBA-2 feature vector
    let valid_counts = vec![26, 256];
    
    assert!(valid_counts.contains(&dataset.feature_count),
            "Feature count should be 26 or 256, got {}", dataset.feature_count);
    
    println!("✅ Feature count valid: {}", dataset.feature_count);
    
    // Print first 10 feature names
    println!("\n📋 First 10 features:");
    for stats in dataset.feature_stats.iter().take(10) {
        println!("   {}: {}", stats.index, stats.name);
    }
    
    Ok(())
}

//
// Test 5: Sample count validation (1,000 samples expected)
//
#[tokio::test]
async fn test_calibration_sample_count() -> Result<()> {
    println!("🧪 Test 5: Sample count validation\n");
    
    let output_path = get_calibration_output_path().join("es_fut_calibration.json");
    
    if !output_path.exists() {
        println!("⚠️  Calibration JSON not found, skipping test");
        return Ok(());
    }
    
    let json_str = std::fs::read_to_string(&output_path)?;
    let dataset: CalibrationDataset = serde_json::from_str(&json_str)?;
    
    println!("📊 Sample count: {}", dataset.sample_count);
    println!("📊 Expected: 1,000 samples");
    println!();
    
    assert_eq!(dataset.sample_count, 1000, "Should have exactly 1,000 samples");
    
    // Validate samples array size
    let expected_size = dataset.sample_count * dataset.feature_count;
    assert_eq!(dataset.samples.len(), expected_size,
               "Samples array should have {} elements (1000 × {}), got {}",
               expected_size, dataset.feature_count, dataset.samples.len());
    
    println!("✅ Sample count valid: 1,000 samples");
    println!("✅ Samples array size: {} elements", dataset.samples.len());
    
    Ok(())
}

//
// Test 6: Load calibration data (integration test)
//
#[tokio::test]
async fn test_load_calibration_data() -> Result<()> {
    println!("🧪 Test 6: Load calibration data (integration)\n");
    
    // Import load function (will fail until implemented)
    use ml::data_loaders::calibration::load_calibration_dataset;
    
    let output_path = get_calibration_output_path().join("es_fut_calibration.json");
    
    if !output_path.exists() {
        println!("⚠️  Calibration JSON not found, skipping test");
        return Ok(());
    }
    
    println!("📖 Loading calibration data from {:?}...", output_path);
    
    // Load dataset using library function
    let dataset = load_calibration_dataset(&output_path).await?;
    
    println!("✅ Loaded dataset:");
    println!("   Samples:  {}", dataset.sample_count);
    println!("   Features: {}", dataset.feature_count);
    println!("   Symbol:   {}", dataset.symbol);
    println!();
    
    // Validate loaded data
    assert_eq!(dataset.sample_count, 1000, "Loaded dataset should have 1,000 samples");
    assert_eq!(dataset.symbol, "ES.FUT", "Symbol should be ES.FUT");
    assert!(dataset.feature_count > 0, "Should have features");
    
    // Test getting min/max for quantization
    println!("📊 Per-feature ranges for quantization:");
    for stats in dataset.feature_stats.iter().take(5) {
        println!("   {}: min={:.6}, max={:.6}, range={:.6}",
                 stats.name, stats.min, stats.max, stats.max - stats.min);
    }
    
    println!("\n✅ Calibration data loaded successfully");
    
    Ok(())
}

//
// Test 7: Integration with DbnSequenceLoader
//
#[tokio::test]
async fn test_calibration_dbn_integration() -> Result<()> {
    println!("🧪 Test 7: Integration with DbnSequenceLoader\n");
    
    use ml::data_loaders::calibration::generate_calibration_dataset;
    
    let test_dir = get_test_data_dir();
    let es_fut_file = test_dir.join("ES.FUT_ohlcv-1m_2024-01-02.dbn");
    
    if !es_fut_file.exists() {
        println!("⚠️  ES.FUT data not found, skipping test");
        return Ok(());
    }
    
    println!("📖 Loading ES.FUT data...");
    
    // Generate small calibration dataset (100 samples for speed)
    let dataset = generate_calibration_dataset(
        &es_fut_file,
        100,  // 100 samples for testing
        "ES.FUT"
    ).await?;
    
    println!("✅ Generated {} samples with {} features",
             dataset.sample_count, dataset.feature_count);
    
    // Validate features are from DbnSequenceLoader
    assert!(dataset.feature_count > 0, "Should have features");
    
    // Check for NaN values
    let nan_count = dataset.samples.iter().filter(|v| v.is_nan()).count();
    assert_eq!(nan_count, 0, "Should have no NaN values, found {}", nan_count);
    
    // Check for reasonable value ranges
    for stats in &dataset.feature_stats {
        assert!(stats.min.is_finite(), "Feature {} min should be finite", stats.name);
        assert!(stats.max.is_finite(), "Feature {} max should be finite", stats.name);
        assert!(stats.mean.is_finite(), "Feature {} mean should be finite", stats.name);
        assert!(stats.std.is_finite(), "Feature {} std should be finite", stats.name);
    }
    
    println!("✅ Integration test passed");
    
    Ok(())
}