foxhunt/ml/examples/download_training_data.rs

//! Download 90 days of real market data from Databento using Rust
//!
//! This uses the official Databento Rust client to download OHLCV-1m data
//! for multiple futures symbols for ML model training.
//!
//! Symbols downloaded:
//! - ES.FUT (E-mini S&P 500) - Stock index
//! - NQ.FUT (E-mini NASDAQ) - Tech index
//! - ZN.FUT (10-Year Treasury) - Fixed income
//! - 6E.FUT (Euro FX) - Currency
//!
//! Usage:
//!     # Set API key in .env file: DATABENTO_API_KEY=your-key
//!     cargo run -p ml --example download_training_data --release
//!
//!     # Custom date range
//!     cargo run -p ml --example download_training_data --release -- \
//!         --start-date 2024-01-02 --days 90
//!
//!     # Specific symbols only
//!     cargo run -p ml --example download_training_data --release -- \
//!         --symbols ES.FUT NQ.FUT

use anyhow::{Context, Result};
use chrono::{Duration, NaiveDate, Utc};
use clap::Parser;
use databento::historical::timeseries::GetRangeParams;
use databento::{Compression, HistoricalClient};
use std::env;
use std::fs;
use std::path::{Path, PathBuf};

#[derive(Debug, Parser)]
#[command(
    name = "download_training_data",
    about = "Download ML training data from Databento"
)]
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 (comma-separated)
    #[arg(
        long,
        value_delimiter = ',',
        default_value = "ES.FUT,NQ.FUT,ZN.FUT,6E.FUT"
    )]
    symbols: Vec<String>,

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

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

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

impl DownloadStats {
    fn new() -> Self {
        Self {
            successful: 0,
            failed: 0,
            skipped: 0,
            total_bytes: 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")?;

    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 {
            dates.push(current.format("%Y-%m-%d").to_string());
        }
        current = current.succ_opt().context("Date overflow")?;
    }

    Ok(dates)
}

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

    // Skip if file already exists
    if output_file.exists() {
        let size = fs::metadata(&output_file)?.len();
        return Ok(Some(size));
    }

    println!("  Downloading {} @ {}...", symbol, date);

    // Parse date range (full trading day UTC)
    let start_str = format!("{}T00:00:00Z", date);
    let end_str = format!("{}T23:59:59Z", date);

    // Build download parameters
    let params = GetRangeParams::builder()
        .dataset("GLBX.MDP3".to_string())
        .symbols(vec![symbol.to_string()])
        .schema("ohlcv-1m".to_string())
        .start(start_str)
        .end(end_str)
        .compression(Compression::ZStd)
        .build();

    // Download data
    let data = client
        .timeseries()
        .get_range(&params)
        .await
        .context("Failed to download data")?;

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

    let size = data.len() as u64;
    println!("    ✅ {} bytes written", size);

    Ok(Some(size))
}

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

    println!("================================================================================");
    println!("ML Training Data Download - Databento (Rust)");
    println!("================================================================================\n");

    // Load API key from environment or .env file
    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 ($0.12 per symbol per day)
    let estimated_cost = dates.len() as f64 * opts.symbols.len() as f64 * 0.12;

    println!("📊 Download Configuration:");
    println!("  Start date: {}", opts.start_date);
    println!("  Trading days: {}", dates.len());
    println!(
        "  Symbols: {} ({})",
        opts.symbols.len(),
        opts.symbols.join(", ")
    );
    println!("  Schema: ohlcv-1m");
    println!("  Dataset: GLBX.MDP3");
    println!("  Output: {}", opts.output_dir);
    println!();
    println!(
        "📦 Total Downloads: {} files",
        dates.len() * opts.symbols.len()
    );
    println!("💰 Estimated Cost: ${:.2}", estimated_cost);
    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
    println!(
        "⚠️  This will download data and incur costs (~${:.2})",
        estimated_cost
    );
    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 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();

    // 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;
            print!(
                "[{}/{} - {:.1}%] {} @ {}... ",
                current_file, total_files, progress, symbol, date
            );
            std::io::Write::flush(&mut std::io::stdout())?;

            match download_symbol_day(&client, symbol, date, &output_path).await {
                Ok(Some(size)) => {
                    if output_path
                        .join(format!("{}_ohlcv-1m_{}.dbn", symbol, date))
                        .exists()
                    {
                        stats.successful += 1;
                        stats.total_bytes += size;
                        println!("✅ {} KB", size / 1024);
                    } else {
                        stats.skipped += 1;
                        println!("⏭️  Skipped (already exists)");
                    }
                },
                Ok(None) => {
                    stats.failed += 1;
                    println!("⚠️  No data (holiday/no trading)");
                },
                Err(e) => {
                    stats.failed += 1;
                    println!("❌ Error: {}", e);
                },
            }
        }
        println!();
    }

    // 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: {:.1} MB",
        stats.total_bytes as f64 / 1_048_576.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. Run ML readiness validation with new data:");
    println!("   cargo test -p ml --test ml_readiness_validation_tests");
    println!();
    println!("2. Run training time benchmarks:");
    println!("   cargo run -p ml --example benchmark_training_time --release");
    println!();

    if success_rate >= 80.0 {
        println!(
            "✅ SUCCESS: Downloaded {:.1}% of requested data!",
            success_rate
        );
        println!("   Ready for ML training benchmarks on RTX 3050 Ti");
    } else if success_rate >= 50.0 {
        println!(
            "⚠️  PARTIAL SUCCESS: Downloaded {:.1}% of data",
            success_rate
        );
        println!(
            "   May be sufficient for benchmarking, but consider re-downloading missing files"
        );
    } else {
        println!("❌ ERROR: Only downloaded {:.1}% of data", success_rate);
        println!("   Check errors above and retry");
    }

    Ok(())
}