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