foxhunt/ml/examples/download_l2_data.rs

//! Download 90 days of Level 2 order book data from DataBento for TLOB training
//!
//! This downloads MBP-10 (Market By Price, 10 levels) data for multiple futures symbols.
//! MBP-10 provides tick-by-tick order book snapshots with 10 bid and 10 ask price levels.
//!
//! Symbols downloaded:
//! - ES.FUT (E-mini S&P 500) - Stock index futures
//! - NQ.FUT (E-mini NASDAQ) - Tech index futures
//! - ZN.FUT (10-Year Treasury) - Fixed income futures
//! - 6E.FUT (Euro FX) - Currency futures
//!
//! Expected cost: $12-$25 (based on single-day test extrapolation)
//! Expected time: 2-4 hours (network dependent)
//! Expected size: 10-20 GB compressed (30-60 GB uncompressed)
//!
//! Usage:
//!     # Default: 90 days, 4 symbols
//!     cargo run -p ml --example download_l2_data --release
//!
//!     # Custom date range
//!     cargo run -p ml --example download_l2_data --release -- \
//!         --start-date 2024-01-02 --days 30
//!
//!     # Specific symbols only
//!     cargo run -p ml --example download_l2_data --release -- \
//!         --symbols ES.FUT NQ.FUT
//!
//!     # Dry run (preview only)
//!     cargo run -p ml --example download_l2_data --release -- --dry-run

use anyhow::{Context, Result};
use chrono::Datelike;
use chrono::NaiveDate;
use clap::Parser;
use databento::historical::timeseries::GetRangeParams;
use databento::{historical::DateTimeRange, HistoricalClient};
use dbn::{Compression, Schema};
use std::env;
use std::fs;
use std::path::{Path, PathBuf};
use std::str::FromStr;
use tokio::io::AsyncReadExt;

#[derive(Debug, Parser)]
#[command(
    name = "download_l2_data",
    about = "Download Level 2 order book data (MBP-10) for TLOB training"
)]
struct Opts {
    /// Start date (YYYY-MM-DD)
    #[arg(long, default_value = "2024-01-02")]
    start_date: String,

    /// Number of trading days to download
    #[arg(long, default_value = "90")]
    days: i64,

    /// Symbols to download (space-separated)
    #[arg(long, default_value = "ES.FUT")]
    symbols: Vec<String>,

    /// Output directory
    #[arg(long, default_value = "test_data/real/databento/l2_order_book")]
    output_dir: String,

    /// Dry run (preview only, no downloads)
    #[arg(long)]
    dry_run: bool,

    /// Skip confirmation prompt
    #[arg(long)]
    yes: bool,
}

struct DownloadStats {
    successful: usize,
    failed: usize,
    skipped: usize,
    total_bytes: u64,
    total_records: u64,
}

impl DownloadStats {
    fn new() -> Self {
        Self {
            successful: 0,
            failed: 0,
            skipped: 0,
            total_bytes: 0,
            total_records: 0,
        }
    }
}

fn generate_trading_dates(start_date_str: &str, num_days: i64) -> Result<Vec<String>> {
    let start_date = NaiveDate::parse_from_str(start_date_str, "%Y-%m-%d")
        .context("Failed to parse start date. Use format: YYYY-MM-DD")?;

    let mut dates = Vec::new();
    let mut current = start_date;

    while dates.len() < num_days as usize {
        // Skip weekends (Saturday=5, Sunday=6)
        if current.weekday().num_days_from_monday() < 5 {
            // Monday=0, ..., Friday=4
            dates.push(current.format("%Y-%m-%d").to_string());
        }
        current = current.succ_opt().context("Date overflow")?;
    }

    Ok(dates)
}

async fn download_symbol_day(
    client: &mut HistoricalClient,
    symbol: &str,
    date: &str,
    output_dir: &Path,
) -> Result<Option<(u64, u64)>> {
    let output_file = output_dir.join(format!("{}_mbp-10_{}.dbn", symbol, date));

    // Skip if file already exists
    if output_file.exists() {
        let size = fs::metadata(&output_file)?.len();
        // Estimate record count (avg 480 bytes per MBP-10 record)
        let estimated_records = size / 480;
        return Ok(Some((size, estimated_records)));
    }

    // Parse date range (full trading day UTC)
    use time::{Date, PrimitiveDateTime, 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();

    let schema_enum = Schema::from_str("mbp-10").context("Failed to parse schema")?;

    // Build download parameters
    let params = GetRangeParams::builder()
        .dataset("GLBX.MDP3".to_string()) // CME Globex
        .symbols(vec![symbol.to_string()])
        .schema(schema_enum) // Level 2: 10 bid/ask levels
        .date_time_range(date_time_range)
        .build();

    // Download data with retry logic
    let mut retries = 0;
    let max_retries = 3;

    loop {
        match client.timeseries().get_range(&params).await {
            Ok(mut decoder) => {
                // 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 size = buffer.len() as u64;

                // Validate minimum size (should be >1 KB for a trading day)
                if size < 1024 {
                    return Ok(None); // Likely no data (holiday/no trading)
                }

                // Write to file
                fs::write(&output_file, &buffer).context("Failed to write data file")?;

                // Estimate record count
                let estimated_records = size / 480;

                return Ok(Some((size, estimated_records)));
            },
            Err(e) => {
                retries += 1;
                if retries >= max_retries {
                    return Err(anyhow::anyhow!("Max retries exceeded: {}", e));
                }

                eprintln!("  ⚠️  Retry {}/{}: {}", retries, max_retries, e);
                tokio::time::sleep(tokio::time::Duration::from_secs(2_u64.pow(retries))).await;
            },
        }
    }
}

#[tokio::main]
async fn main() -> Result<()> {
    let opts = Opts::parse();

    println!("================================================================================");
    println!("DataBento MBP-10 Level 2 Order Book Download");
    println!("TLOB Neural Network Training Data Acquisition");
    println!("================================================================================\n");

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

    // Generate trading dates
    let dates = generate_trading_dates(&opts.start_date, opts.days)?;

    // Estimate cost based on single-day test ($0.03-$0.08 per symbol per day)
    let cost_per_symbol_day = 0.05; // Conservative midpoint
    let estimated_cost = dates.len() as f64 * opts.symbols.len() as f64 * cost_per_symbol_day;

    // Estimate size based on single-day test (~50-150 MB per symbol per day compressed)
    let mb_per_symbol_day = 100.0; // Conservative midpoint
    let estimated_mb = dates.len() as f64 * opts.symbols.len() as f64 * mb_per_symbol_day;
    let estimated_gb = estimated_mb / 1024.0;

    println!("📊 Download Configuration:");
    println!("  Start date:    {}", opts.start_date);
    println!("  Trading days:  {}", dates.len());
    println!(
        "  Symbols:       {} ({})",
        opts.symbols.len(),
        opts.symbols.join(", ")
    );
    println!("  Schema:        mbp-10 (Level 2 Order Book - 10 bid/ask levels)");
    println!("  Dataset:       GLBX.MDP3 (CME Globex)");
    println!("  Compression:   ZStd (~70% size reduction)");
    println!("  Output:        {}", opts.output_dir);
    println!();
    println!(
        "📦 Total Downloads: {} files",
        dates.len() * opts.symbols.len()
    );
    println!("💾 Estimated Size: {:.2} GB compressed", estimated_gb);
    println!("💰 Estimated Cost: ${:.2}", estimated_cost);
    println!(
        "⏱️  Estimated Time: {:.1}-{:.1} hours (network dependent)",
        estimated_gb / 10.0,
        estimated_gb / 5.0
    ); // 5-10 MB/s throughput
    println!();

    if opts.dry_run {
        println!("🔍 DRY RUN: Preview complete. Remove --dry-run to execute.");
        println!();
        println!("First 5 dates to download:");
        for date in dates.iter().take(5) {
            println!("  • {}", date);
        }
        if dates.len() > 5 {
            println!("  ... ({} more dates)", dates.len() - 5);
        }
        return Ok(());
    }

    // Confirm before proceeding
    if !opts.yes {
        println!("⚠️  This will download Level 2 order book data and incur costs:");
        println!("   • Estimated cost: ${:.2}", estimated_cost);
        println!("   • Estimated size: {:.2} GB", estimated_gb);
        println!(
            "   • Estimated time: {:.1}-{:.1} hours",
            estimated_gb / 10.0,
            estimated_gb / 5.0
        );
        println!();
        print!("Proceed with download? (yes/no): ");
        std::io::Write::flush(&mut std::io::stdout())?;

        let mut input = String::new();
        std::io::stdin().read_line(&mut input)?;
        if !input.trim().eq_ignore_ascii_case("yes") && !input.trim().eq_ignore_ascii_case("y") {
            println!("Download cancelled.");
            return Ok(());
        }
        println!();
    }

    // Create output directory
    let output_path = PathBuf::from(&opts.output_dir);
    fs::create_dir_all(&output_path)?;
    println!("📁 Created output directory: {}", opts.output_dir);
    println!();

    // Initialize DataBento client
    let mut client = HistoricalClient::builder().key(api_key)?.build()?;
    println!("✅ DataBento client initialized");
    println!();

    // Track statistics
    let mut stats = DownloadStats::new();
    let total_files = dates.len() * opts.symbols.len();
    let start_time = std::time::Instant::now();

    // Download all combinations
    let mut current_file = 0;

    for symbol in &opts.symbols {
        println!("{:-<80}", "");
        println!("📥 Downloading: {}", symbol);
        println!("{:-<80}", "");
        println!();

        for date in &dates {
            current_file += 1;
            let progress = (current_file as f64 / total_files as f64) * 100.0;
            let elapsed = start_time.elapsed().as_secs_f64();
            let eta = if current_file > 1 {
                elapsed / (current_file - 1) as f64 * (total_files - current_file) as f64
            } else {
                0.0
            };

            print!(
                "[{}/{} - {:.1}%] {} @ {} (ETA: {:.0}m)... ",
                current_file,
                total_files,
                progress,
                symbol,
                date,
                eta / 60.0
            );
            std::io::Write::flush(&mut std::io::stdout())?;

            match download_symbol_day(&mut client, symbol, date, &output_path).await {
                Ok(Some((size, records))) => {
                    stats.successful += 1;
                    stats.total_bytes += size;
                    stats.total_records += records;
                    println!(
                        "✅ {:.1} MB ({} records)",
                        size as f64 / 1_048_576.0,
                        records
                    );
                },
                Ok(None) => {
                    stats.skipped += 1;
                    println!("⏭️  Skipped (no data - holiday/no trading)");
                },
                Err(e) => {
                    stats.failed += 1;
                    println!("❌ Error: {}", e);
                },
            }

            // Rate limit: Max 10 requests per minute (6 second delay)
            if current_file % 10 == 0 && current_file < total_files {
                println!("  ⏸️  Rate limit pause (10 req/min limit)...");
                tokio::time::sleep(tokio::time::Duration::from_secs(6)).await;
            }
        }
        println!();
    }

    let total_duration = start_time.elapsed();

    // Summary
    println!();
    println!("================================================================================");
    println!("📊 DOWNLOAD SUMMARY");
    println!("================================================================================");
    println!();
    println!("✅ Successful:    {}/{}", stats.successful, total_files);
    println!("⏭️  Skipped:       {}/{}", stats.skipped, total_files);
    println!("❌ Failed:        {}/{}", stats.failed, total_files);
    println!();
    println!(
        "💾 Total Size:    {:.2} GB",
        stats.total_bytes as f64 / 1_073_741_824.0
    );
    println!(
        "📈 Total Records: {:.1}M order book updates",
        stats.total_records as f64 / 1_000_000.0
    );
    println!(
        "⏱️  Duration:      {:.1} minutes",
        total_duration.as_secs_f64() / 60.0
    );
    println!("💰 Estimated Cost: ${:.2}", estimated_cost);
    println!();

    let success_rate = (stats.successful as f64 / total_files as f64) * 100.0;

    println!("📋 NEXT STEPS:");
    println!("1. Validate downloaded data:");
    println!("   cargo run -p ml --example validate_l2_data --release");
    println!();
    println!("2. Create TLOB data loader:");
    println!("   See ml/src/data_loaders/tlob_loader.rs");
    println!();
    println!("3. Run TLOB training:");
    println!("   tli train --model TLOB --epochs 10");
    println!();

    if success_rate >= 95.0 {
        println!(
            "✅ SUCCESS: Downloaded {:.1}% of requested data!",
            success_rate
        );
        println!(
            "   {} order book updates ready for TLOB training",
            stats.total_records
        );
    } else if success_rate >= 80.0 {
        println!(
            "⚠️  PARTIAL SUCCESS: Downloaded {:.1}% of data",
            success_rate
        );
        println!("   May be sufficient for training, but consider re-downloading missing files");
    } else {
        println!("❌ ERROR: Only downloaded {:.1}% of data", success_rate);
        println!("   Check errors above and retry missing files");
    }

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

    Ok(())
}