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(())
}