foxhunt/ml/examples/download_l2_test.rs

//! Test DataBento MBP-10 (Level 2) download for TLOB training
//!
//! This tests a single-day download of ES.FUT MBP-10 data to validate:
//! - API connectivity and authentication
//! - MBP-10 schema support
//! - DBN file parsing (Mbp10Msg records)
//! - Cost estimation for full 90-day download
//!
//! Cost: ~$0.01-$0.05 (single day, single symbol)
//!
//! Usage:
//!     # Set API key in .env: DATABENTO_API_KEY=db-95LEt9gtDRPJfc55NVUB5KL3A3uf6
//!     cargo run -p ml --example download_l2_test --release

use anyhow::{Context, Result};
use databento::historical::timeseries::GetRangeParams;
use databento::{HistoricalClient, historical::DateTimeRange};
use dbn::{Compression, Schema};
use dbn::decode::{DbnDecoder, DbnMetadata, DecodeRecordRef};
use dbn::RecordRefEnum;
use std::str::FromStr;
use std::env;
use std::fs::{self, File};
use std::io::BufReader;
use std::path::PathBuf;

#[tokio::main]
async fn main() -> Result<()> {
    println!("================================================================================");
    println!("DataBento MBP-10 Level 2 Order Book Test Download");
    println!("================================================================================\n");

    // Load API key
    dotenv::dotenv().ok();
    let api_key = env::var("DATABENTO_API_KEY")
        .context("DATABENTO_API_KEY not found. Set it in .env file.")?;

    println!("✅ API Key found: {}...{}\n", &api_key[0..10], &api_key[api_key.len() - 10..]);

    // Test parameters
    let symbol = "ES.FUT";
    let date = "2024-01-02"; // Single trading day
    let schema = "mbp-10"; // Level 2 market depth (10 price levels)
    let dataset = "GLBX.MDP3"; // CME Globex

    println!("📋 Test Parameters:");
    println!("  Symbol:      {} (E-mini S&P 500 Futures)", symbol);
    println!("  Date:        {} (single trading day)", date);
    println!("  Schema:      {} (Level 2 Order Book - 10 bid/ask levels)", schema);
    println!("  Dataset:     {} (CME Group MDP 3.0)", dataset);
    println!("  Compression: ZStd (~70% size reduction)");
    println!();

    // Create output directory
    let output_dir = PathBuf::from("test_data/real/databento/l2_test");
    fs::create_dir_all(&output_dir)?;

    let output_file = output_dir.join(format!("{}_mbp-10_{}.dbn", symbol, date));

    println!("📁 Output: {:?}", output_file);
    println!();

    // Initialize DataBento client
    println!("🔌 Initializing DataBento client...");
    let client = HistoricalClient::builder()
        .key(api_key)?
        .build()?;
    println!("✅ Client initialized\n");

    // Build download parameters
    // Parse date and create DateTimeRange
    use time::{PrimitiveDateTime, Date, Time, UtcOffset};
    let date_obj = Date::parse(date, &time::format_description::parse("[year]-[month]-[day]")?)?;
    let start_dt = PrimitiveDateTime::new(date_obj, Time::MIDNIGHT).assume_offset(UtcOffset::UTC);
    let end_dt = start_dt + time::Duration::days(1);
    let date_time_range: DateTimeRange = (start_dt, end_dt).into();

    // Parse schema
    let schema_enum = Schema::from_str(schema)
        .context("Failed to parse schema")?;

    let params = GetRangeParams::builder()
        .dataset(dataset.to_string())
        .symbols(vec![symbol.to_string()])
        .schema(schema_enum)
        .date_time_range(date_time_range)
        .build();

    println!("📥 Downloading MBP-10 data...");
    println!("  This may take 30-60 seconds for a single trading day");
    println!();

    // Download data
    use tokio::io::AsyncReadExt;
    let download_start = std::time::Instant::now();
    let mut decoder = client
        .timeseries()
        .get_range(&params)
        .await
        .context("Failed to download data. Check API key and symbol/date validity.")?;
    
    // Read all data into buffer
    let mut buffer = Vec::new();
    let mut temp_buf = vec![0u8; 8192];
    loop {
        let n = decoder.get_mut().read(&mut temp_buf).await?;
        if n == 0 {
            break;
        }
        buffer.extend_from_slice(&temp_buf[..n]);
    }
    
    let download_duration = download_start.elapsed();

    let size_bytes = buffer.len();
    let size_kb = size_bytes as f64 / 1024.0;
    let size_mb = size_kb / 1024.0;

    println!("✅ Download complete!");
    println!("  Duration:     {:.2}s", download_duration.as_secs_f64());
    println!("  Size:         {} bytes ({:.2} KB, {:.2} MB)", size_bytes, size_kb, size_mb);
    println!();

    // Write to file
    fs::write(&output_file, &buffer)
        .context("Failed to write DBN file")?;
    println!("💾 Saved to: {:?}", output_file);
    println!();

    // Parse DBN file to validate and count records
    println!("🔍 Parsing DBN file...");
    let file = File::open(&output_file)?;
    let reader = BufReader::new(file);
    let mut decoder = DbnDecoder::new(reader)
        .context("Failed to create DBN decoder. File may be corrupted.")?;

    let metadata = decoder.metadata();
    println!("📊 Metadata:");
    println!("  Dataset:      {:?}", metadata.dataset);
    println!("  Schema:       {:?}", metadata.schema);
    println!("  Symbols:      {:?}", metadata.symbols);
    println!("  Start:        {:?}", metadata.start);
    println!("  End:          {:?}", metadata.end);
    println!();

    // Count records by type
    let mut mbp10_count = 0;
    let mut other_count = 0;
    let mut sample_records = Vec::new();

    println!("📈 Decoding records...");
    loop {
        match decoder.decode_record_ref() {
            Ok(Some(record)) => {
                let record_enum = record.as_enum()
                    .context("Failed to convert record to enum")?;

                match record_enum {
                    RecordRefEnum::Mbp10(mbp) => {
                        mbp10_count += 1;

                        // Collect first 3 records as samples
                        if sample_records.len() < 3 {
                            sample_records.push((
                                mbp.hd.ts_event,
                                mbp.levels[0].bid_px,
                                mbp.levels[0].ask_px,
                                mbp.levels[0].bid_sz,
                                mbp.levels[0].ask_sz,
                            ));
                        }
                    }
                    _ => {
                        other_count += 1;
                    }
                }
            }
            Ok(None) => break,
            Err(e) => {
                eprintln!("⚠️  Decode error: {}", e);
                break;
            }
        }
    }

    println!("✅ Parsing complete!");
    println!("  MBP-10 records: {}", mbp10_count);
    println!("  Other records:  {}", other_count);
    println!("  Total:          {}", mbp10_count + other_count);
    println!();

    // Display sample records
    if !sample_records.is_empty() {
        println!("📋 Sample Records (first 3):");
        for (i, (ts, bid_px, ask_px, bid_sz, ask_sz)) in sample_records.iter().enumerate() {
            let bid_f64 = *bid_px as f64 * 1e-9;
            let ask_f64 = *ask_px as f64 * 1e-9;
            let spread = ask_f64 - bid_f64;

            println!("  Record #{}: timestamp={}, bid={:.2}, ask={:.2}, spread={:.4}, bid_sz={}, ask_sz={}",
                     i + 1, ts, bid_f64, ask_f64, spread, bid_sz, ask_sz);
        }
        println!();
    }

    // Cost estimation
    let size_gb = size_bytes as f64 / 1_073_741_824.0;
    let cost_per_gb = 1.0; // Conservative estimate: $1/GB
    let estimated_cost = size_gb * cost_per_gb;

    println!("💰 Cost Estimation:");
    println!("  Single day (1 symbol): ${:.4}", estimated_cost);
    println!();

    // Extrapolate for full download
    let full_download_days = 90;
    let full_download_symbols = 4; // ES.FUT, NQ.FUT, ZN.FUT, 6E.FUT
    let full_size_gb = size_gb * full_download_days as f64 * full_download_symbols as f64;
    let full_cost = full_size_gb * cost_per_gb;

    println!("📊 Extrapolation for Full Download:");
    println!("  Symbols:       {} (ES.FUT, NQ.FUT, ZN.FUT, 6E.FUT)", full_download_symbols);
    println!("  Days:          {} (Jan-Mar 2024)", full_download_days);
    println!("  Estimated GB:  {:.2} GB", full_size_gb);
    println!("  Estimated Cost: ${:.2}", full_cost);
    println!("  Credits Left:  ${:.2} (of $125 available)", 125.0 - full_cost);
    println!();

    // Validation summary
    println!("================================================================================");
    println!("✅ VALIDATION SUMMARY");
    println!("================================================================================");
    println!();

    let mut all_checks_passed = true;

    // Check 1: File exists and non-empty
    let check1 = output_file.exists() && size_bytes > 0;
    println!("[{}] File downloaded and saved", if check1 { "✅" } else { "❌" });
    all_checks_passed &= check1;

    // Check 2: DBN decoder can parse file
    let check2 = mbp10_count > 0;
    println!("[{}] DBN decoder successful (parsed {} MBP-10 records)",
             if check2 { "✅" } else { "❌" }, mbp10_count);
    all_checks_passed &= check2;

    // Check 3: Expected record count (10,000-100,000 for liquid futures)
    let check3 = mbp10_count >= 1_000 && mbp10_count <= 1_000_000;
    println!("[{}] Record count in expected range ({})",
             if check3 { "✅" } else { "⚠️" }, mbp10_count);
    all_checks_passed &= check3;

    // Check 4: Cost within budget
    let check4 = estimated_cost < 0.10; // Single day should be <$0.10
    println!("[{}] Single-day cost acceptable (${:.4})",
             if check4 { "✅" } else { "⚠️" }, estimated_cost);
    all_checks_passed &= check4;

    // Check 5: Full download projected within budget
    let check5 = full_cost < 30.0; // Full download should be <$30
    println!("[{}] Full download projected within budget (${:.2})",
             if check5 { "✅" } else { "⚠️" }, full_cost);
    all_checks_passed &= check5;

    println!();

    if all_checks_passed {
        println!("🎉 SUCCESS! All checks passed.");
        println!();
        println!("📋 NEXT STEPS:");
        println!("1. Review cost estimate (${:.2} for 90 days × 4 symbols)", full_cost);
        println!("2. If acceptable, run full download:");
        println!("   cargo run -p ml --example download_l2_data --release");
        println!("3. Integrate with TLOB training:");
        println!("   See ml/src/data_loaders/tlob_loader.rs");
    } else {
        println!("⚠️  Some checks failed. Review above and debug before full download.");
    }

    println!();
    println!("================================================================================");

    Ok(())
}