foxhunt/crates/ml/examples/precompute_features.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! Precompute Features — DBN to .fxcache pipeline
//!
//! Reads OHLCV + MBP-10 + trades DBN data, runs the full feature extraction
//! pipeline (standalone, no GPU required), and writes a `.fxcache` binary file for
//! zero-overhead GPU loading during training.
//!
//! Usage:
//!     cargo run -p ml --example precompute_features --release -- \
//!         --data-dir test_data/futures-baseline \
//!         --mbp10-data-dir test_data/futures-baseline-mbp10 \
//!         --trades-data-dir test_data/futures-baseline-trades \
//!         --output-dir test_data/feature-cache \
//!         --symbol ES.FUT \
//!         --bf16 \
//!         --yes

use std::path::{Path, PathBuf};
use std::time::Instant;

use anyhow::{Context, Result};
use clap::Parser;
use tracing::info;

use ml::trainers::dqn::{
    collect_dbn_files_recursive, extract_features_from_bars, extract_ohlcv_bars_from_dbn,
};
use ml::features::extraction::OHLCVBar;

// ---------------------------------------------------------------------------
// CLI
// ---------------------------------------------------------------------------

#[derive(Debug, Parser)]
#[command(
    name = "precompute_features",
    about = "Precompute features from DBN data and write .fxcache binary"
)]
struct Opts {
    /// Directory containing OHLCV .dbn/.dbn.zst files (e.g. test_data/futures-baseline)
    #[arg(long)]
    data_dir: String,

    /// Directory containing MBP-10 .dbn/.dbn.zst files (optional)
    #[arg(long)]
    mbp10_data_dir: Option<String>,

    /// Directory containing trades .dbn/.dbn.zst files (optional)
    #[arg(long)]
    trades_data_dir: Option<String>,

    /// Output directory for .fxcache files (default: sibling of --data-dir)
    #[arg(long)]
    output_dir: Option<String>,

    /// Symbol name (used in output filename)
    #[arg(long, default_value = "ES.FUT")]
    symbol: String,

    /// Write version 2 (bf16) instead of version 1 (f64)
    #[arg(long)]
    bf16: bool,

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

// ---------------------------------------------------------------------------
// Cache key computation
// ---------------------------------------------------------------------------

/// Compute a SHA256 cache key from all DBN files across the given directories.
///
// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

#[tokio::main]
async fn main() -> Result<()> {
    // Initialize tracing
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
        )
        .init();

    let opts = Opts::parse();

    let data_dir = PathBuf::from(&opts.data_dir);
    // Default output_dir: sibling of data_dir (e.g. test_data/futures-baseline → test_data/feature-cache)
    let output_dir = match opts.output_dir {
        Some(ref d) => PathBuf::from(d),
        None => data_dir.parent().unwrap_or(&data_dir).join("feature-cache"),
    };
    let mbp10_dir = opts.mbp10_data_dir.as_ref().map(PathBuf::from);
    let trades_dir = opts.trades_data_dir.as_ref().map(PathBuf::from);

    // ── Validate inputs ──────────────────────────────────────────────────────
    if !data_dir.exists() {
        anyhow::bail!("Data directory not found: {}", data_dir.display());
    }

    // ── Print config summary ─────────────────────────────────────────────────
    println!("================================================================================");
    println!("Precompute Features — DBN to .fxcache Pipeline");
    println!("================================================================================");
    println!();
    println!("Data dir:      {}", data_dir.display());
    if let Some(ref d) = mbp10_dir {
        println!("MBP-10 dir:    {}", d.display());
    }
    if let Some(ref d) = trades_dir {
        println!("Trades dir:    {}", d.display());
    }
    println!("Output dir:    {}", output_dir.display());
    println!("Symbol:        {}", opts.symbol);
    println!("Format:        {}", if opts.bf16 { "bf16 (v2)" } else { "f64 (v1)" });
    println!();

    // ── Confirmation ─────────────────────────────────────────────────────────
    if !opts.yes {
        print!("Proceed with feature extraction? (yes/no): ");
        std::io::Write::flush(&mut std::io::stdout())?;

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

    let t0 = Instant::now();

    // ── Step 1: Load OHLCV bars from DBN files ──────────────────────────────
    info!("Loading OHLCV bars from {}...", data_dir.display());
    let mut dbn_files = collect_dbn_files_recursive(&data_dir);
    dbn_files.sort();
    info!("Found {} DBN files", dbn_files.len());

    let mut all_bars: Vec<OHLCVBar> = Vec::new();
    for file in &dbn_files {
        match extract_ohlcv_bars_from_dbn(file) {
            Ok(bars) => all_bars.extend(bars),
            Err(e) => tracing::warn!("Failed to load {:?}: {e}", file.file_name()),
        }
    }
    all_bars.sort_by_key(|b| b.timestamp);
    info!("Loaded {} OHLCV bars in {:.1}s", all_bars.len(), t0.elapsed().as_secs_f64());

    // ── Step 2: Extract 42-dim feature vectors ──────────────────────────────
    let t1 = Instant::now();
    info!("Extracting features...");
    let feature_vectors = extract_features_from_bars(&all_bars)
        .context("Feature extraction failed")?;
    info!("Extracted {} feature vectors in {:.1}s", feature_vectors.len(), t1.elapsed().as_secs_f64());

    // ── Step 3: Build targets [preproc_close, preproc_next, raw_close, raw_next]
    const WARMUP: usize = 50;
    let n = feature_vectors.len().saturating_sub(1); // need next bar for target
    let features: Vec<[f64; 42]> = feature_vectors[..n].to_vec();
    let mut targets: Vec<[f64; 4]> = Vec::with_capacity(n);
    for i in 0..n {
        let raw_curr = all_bars[i + WARMUP].close;
        let raw_next = all_bars[i + WARMUP + 1].close;
        targets.push([raw_curr, raw_next, raw_curr, raw_next]);
    }

    // ── Step 3b: Build per-bar timestamps ─────────────────────────────────
    let timestamps: Vec<i64> = (0..n)
        .map(|i| all_bars[i + WARMUP].timestamp.timestamp_nanos_opt().unwrap_or(0))
        .collect();

    // ── Step 4: Compute OFI from MBP-10 + trades ────────────────────────────
    let t2 = Instant::now();
    let ofi: Vec<[f64; 8]> = if let Some(ref mbp10_path) = mbp10_dir {
        use ml::features::mbp10_loader::load_ofi_features_parallel;
        use ml::features::trades_loader::load_trades_sync;
        use ml::features::ofi_calculator::OFICalculator;

        info!("Loading MBP-10 snapshots from {}...", mbp10_path.display());
        let mut mbp10_files = collect_dbn_files_recursive(mbp10_path);
        mbp10_files.sort();

        if mbp10_files.is_empty() {
            info!("No MBP-10 files found, OFI will be zeros");
            vec![[0.0; 8]; n]
        } else {
            // Load MBP-10 snapshots (parallel)
            use data::providers::databento::dbn_parser::DbnParser;
            let per_file: Vec<Vec<_>> = {
                use rayon::prelude::*;
                mbp10_files.par_iter().filter_map(|file| {
                    let parser = match DbnParser::new() {
                        Ok(p) => p,
                        Err(e) => { tracing::warn!("Parser init failed: {e}"); return None; }
                    };
                    let mut snapshots = Vec::new();
                    match parser.parse_mbp10_streaming(file, 100, |snap| {
                        snapshots.push(snap.clone());
                    }) {
                        Ok(_) => {
                            info!("  {} -> {} snapshots", file.file_name().unwrap_or_default().to_string_lossy(), snapshots.len());
                            Some(snapshots)
                        }
                        Err(e) => {
                            tracing::warn!("  Failed {:?}: {e}", file.file_name());
                            None
                        }
                    }
                }).collect()
            };

            let mut all_snapshots = Vec::new();
            for snaps in per_file { all_snapshots.extend(snaps); }
            all_snapshots.sort_by_key(|s| s.timestamp);
            info!("Loaded {} MBP-10 snapshots in {:.1}s", all_snapshots.len(), t2.elapsed().as_secs_f64());

            // Load trades (parallel)
            let all_trades = if let Some(ref tdir) = trades_dir {
                let mut trade_files = collect_dbn_files_recursive(tdir);
                trade_files.sort();
                let per_file_trades: Vec<Vec<_>> = {
                    use rayon::prelude::*;
                    trade_files.par_iter().filter_map(|path| {
                        match load_trades_sync(path) {
                            Ok(t) => { info!("  {} trades from {:?}", t.len(), path.file_name()); Some(t) }
                            Err(e) => { tracing::warn!("  Failed {:?}: {e}", path.file_name()); None }
                        }
                    }).collect()
                };
                let mut trades = Vec::new();
                for t in per_file_trades { trades.extend(t); }
                trades.sort_by_key(|t| t.timestamp);
                if trades.is_empty() { None } else { Some(trades) }
            } else {
                None
            };

            // Compute per-bar OFI
            use ml::features::mbp10_loader::get_snapshots_for_timestamp;
            use ml::features::trades_loader::get_trades_for_bar;
            let mut calculator = OFICalculator::new();
            let mut ofi_per_bar = Vec::with_capacity(n);

            for i in 0..n {
                let bar = &all_bars[i + WARMUP];
                let bar_ts = bar.timestamp.timestamp_nanos_opt().unwrap_or(0) as u64;

                // Feed trades for this bar window
                if let Some(ref trades) = all_trades {
                    let bar_end_ts = all_bars.get(i + WARMUP + 1)
                        .map(|b| b.timestamp.timestamp_nanos_opt().unwrap_or(0) as u64)
                        .unwrap_or(bar_ts + 60_000_000_000);
                    for trade in get_trades_for_bar(trades, bar_ts, bar_end_ts) {
                        calculator.feed_trade(trade.price, trade.volume, trade.is_buy);
                    }
                }

                let window = get_snapshots_for_timestamp(&all_snapshots, bar_ts, 1);
                if let Some(snap) = window.first() {
                    match calculator.calculate(snap) {
                        Ok(f) if f.is_valid() => ofi_per_bar.push(f.to_array()),
                        _ => ofi_per_bar.push([0.0; 8]),
                    }
                } else {
                    ofi_per_bar.push([0.0; 8]);
                }
            }

            let non_zero = ofi_per_bar.iter().filter(|f| f.iter().any(|&v| v != 0.0)).count();
            info!("OFI computed: {} bars, {} non-zero ({:.1}%) in {:.1}s",
                ofi_per_bar.len(), non_zero,
                if ofi_per_bar.is_empty() { 0.0 } else { non_zero as f64 / ofi_per_bar.len() as f64 * 100.0 },
                t2.elapsed().as_secs_f64());
            ofi_per_bar
        }
    } else {
        info!("No MBP-10 directory, OFI will be zeros");
        vec![[0.0; 8]; n]
    };

    let total_len = n;

    // ── Compute cache key ────────────────────────────────────────────────────
    let hex_key = ml::feature_cache::calculate_dbn_cache_key_full(
        &data_dir,
        mbp10_dir.as_deref(),
        trades_dir.as_deref(),
    ).context("Failed to compute cache key")?;
    let cache_key: [u8; 32] = hex::decode(&hex_key)
        .context("Invalid hex key")?
        .try_into()
        .map_err(|_| anyhow::anyhow!("Cache key wrong length"))?;

    // ── Write .fxcache ───────────────────────────────────────────────────────
    std::fs::create_dir_all(&output_dir)
        .with_context(|| format!("Failed to create output directory: {}", output_dir.display()))?;

    let output_path = output_dir.join(format!("{hex_key}.fxcache"));

    let t2 = Instant::now();
    info!("Writing .fxcache to {}...", output_path.display());

    let bytes_written = ml::fxcache::write_fxcache(
        &output_path,
        &features,
        &targets,
        &ofi,
        &timestamps,
        cache_key,
        opts.bf16,
    )
    .context("Failed to write .fxcache file")?;

    let write_secs = t2.elapsed().as_secs_f64();
    let total_secs = t0.elapsed().as_secs_f64();

    // ── Summary ──────────────────────────────────────────────────────────────
    println!();
    println!("================================================================================");
    println!("PRECOMPUTE SUMMARY");
    println!("================================================================================");
    println!();
    let has_ofi = ofi.iter().any(|o| o.iter().any(|&v| v != 0.0));
    println!("Bars:          {}", total_len);
    println!("Features:      42-dim");
    println!("Targets:       4-dim");
    println!("OFI:           8-dim ({})", if has_ofi { "from MBP-10" } else { "zero-padded" });
    println!("Format:        {} (v{})", if opts.bf16 { "bf16" } else { "f64" }, if opts.bf16 { 2 } else { 1 });
    println!("Cache key:     {}", hex_key);
    println!("Output:        {}", output_path.display());
    println!("Size:          {:.2} MB", bytes_written as f64 / 1_048_576.0);
    println!("Write time:    {:.1}s", write_secs);
    println!("Total time:    {:.1}s", total_secs);
    println!();

    Ok(())
}