foxhunt/ml/examples/generate_calibration_dataset.rs

//! Generate Calibration Dataset for INT8 Quantization
//!
//! Generates 1,000-sample calibration dataset from ES.FUT data for INT8 quantization.
//!
//! Usage:
//! ```bash
//! cargo run -p ml --example generate_calibration_dataset
//! ```

use anyhow::Result;
use ml::data_loaders::calibration::{generate_calibration_dataset, save_calibration_dataset};
use std::path::PathBuf;

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize logging
    tracing_subscriber::fmt()
        .with_max_level(tracing::Level::INFO)
        .init();

    println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
    println!("  Calibration Dataset Generator");
    println!("  INT8 Quantization - ES.FUT Market Data");
    println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
    println!();

    // Input: ES.FUT DBN file
    let es_fut_file = PathBuf::from("test_data/real/databento/ES.FUT_ohlcv-1m_2024-01-02.dbn");
    
    if !es_fut_file.exists() {
        eprintln!("❌ Error: ES.FUT data not found at {:?}", es_fut_file);
        eprintln!("   Please ensure test data is available.");
        return Ok(());
    }

    println!("📂 Input: {:?}", es_fut_file);
    println!();

    // Generate calibration dataset (1,000 samples)
    println!("🔄 Generating calibration dataset...");
    println!("   Target samples: 1,000");
    println!("   Feature dimension: 256 (MAMBA-2)");
    println!();

    let dataset = generate_calibration_dataset(
        &es_fut_file,
        1000,  // 1,000 samples for calibration
        "ES.FUT"
    ).await?;

    println!();
    println!("✅ Dataset generated:");
    println!("   Samples:  {}", dataset.sample_count);
    println!("   Features: {}", dataset.feature_count);
    println!("   Symbol:   {}", dataset.symbol);
    println!();

    // Print sample statistics for first 10 features
    println!("📊 Sample Statistics (first 10 features):");
    println!("   ┌────────┬──────────────────────┬───────────┬───────────┬───────────┬──────────┐");
    println!("   │ Index  │ Name                 │ Min       │ Max       │ Mean      │ Std      │");
    println!("   ├────────┼──────────────────────┼───────────┼───────────┼───────────┼──────────┤");
    
    for stats in dataset.feature_stats.iter().take(10) {
        println!("   │ {:6} │ {:20} │ {:9.4} │ {:9.4} │ {:9.4} │ {:8.4} │",
                 stats.index,
                 stats.name,
                 stats.min,
                 stats.max,
                 stats.mean,
                 stats.std);
    }
    
    println!("   └────────┴──────────────────────┴───────────┴───────────┴───────────┴──────────┘");
    println!();

    // Save to JSON
    let output_dir = PathBuf::from("ml/calibration");
    std::fs::create_dir_all(&output_dir)?;
    
    let output_file = output_dir.join("es_fut_calibration.json");
    println!("💾 Saving to {:?}...", output_file);
    
    save_calibration_dataset(&dataset, &output_file).await?;

    let file_size = std::fs::metadata(&output_file)?.len();
    println!("✅ Saved {} bytes ({:.2} KB, {:.2} MB)",
             file_size,
             file_size as f64 / 1024.0,
             file_size as f64 / 1_048_576.0);
    println!();

    // Validation checks
    println!("🔍 Validation:");
    
    // Check for NaN values
    let nan_count = dataset.samples.iter().filter(|v| v.is_nan()).count();
    if nan_count == 0 {
        println!("   ✅ No NaN values detected");
    } else {
        println!("   ❌ {} NaN values found", nan_count);
    }
    
    // Check for reasonable value ranges
    let mut all_finite = true;
    for stats in &dataset.feature_stats {
        if !stats.min.is_finite() || !stats.max.is_finite() {
            println!("   ❌ Feature {} has non-finite values", stats.index);
            all_finite = false;
        }
    }
    
    if all_finite {
        println!("   ✅ All feature statistics are finite");
    }
    
    // Check sample count
    if dataset.sample_count == 1000 {
        println!("   ✅ Sample count correct (1,000)");
    } else {
        println!("   ⚠️  Sample count: {} (expected 1,000)", dataset.sample_count);
    }
    
    // Check feature count
    if dataset.feature_count == 256 {
        println!("   ✅ Feature count correct (256)");
    } else {
        println!("   ⚠️  Feature count: {} (expected 256)", dataset.feature_count);
    }
    
    println!();
    println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
    println!("  ✅ Calibration Dataset Generation Complete!");
    println!("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━");
    println!();
    println!("📋 Next Steps:");
    println!("   1. Review calibration statistics above");
    println!("   2. Use calibration data for INT8 quantization");
    println!("   3. Apply to TFT model quantization pipeline");
    println!();

    Ok(())
}