foxhunt/crates/ml/examples/baseline_common/mod.rs

//! Shared utilities for the baseline training/evaluation pipeline.
//!
//! Contains DBN file loading, OHLCV bar decoding with timestamp dedup,
//! and cost helpers used by `train_baseline`, `evaluate_baseline`, and
//! `hyperopt_baseline`.

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

use anyhow::{Context, Result};
use chrono::{TimeZone, Utc};
use tracing::{info, warn};

use dbn::decode::DecodeRecord;
use ml::types::OHLCVBar;

/// Recursively discover .dbn.zst files under `dir`.
pub fn find_dbn_files(dir: &Path) -> Result<Vec<PathBuf>> {
    let mut files = Vec::new();
    if !dir.exists() {
        anyhow::bail!("Data directory does not exist: {}", dir.display());
    }
    collect_dbn_files_recursive(dir, &mut files)?;
    files.sort();
    Ok(files)
}

/// Recursive helper that walks the directory tree without `walkdir`.
fn collect_dbn_files_recursive(dir: &Path, out: &mut Vec<PathBuf>) -> Result<()> {
    let entries = std::fs::read_dir(dir)
        .with_context(|| format!("Cannot read directory: {}", dir.display()))?;
    for entry in entries {
        let dir_entry = entry.with_context(|| "Failed to read dir entry")?;
        let entry_path = dir_entry.path();
        if entry_path.is_dir() {
            collect_dbn_files_recursive(&entry_path, out)?;
        } else if entry_path
            .file_name()
            .and_then(|n| n.to_str())
            .is_some_and(|n| n.ends_with(".dbn.zst"))
        {
            out.push(entry_path);
        } else {
            // Skip non-.dbn.zst files
        }
    }
    Ok(())
}

/// Load OHLCV bars from a single .dbn or .dbn.zst file using the `dbn` crate decoder.
///
/// Reads OhlcvMsg records and converts them to [`OHLCVBar`].
fn load_bars_from_dbn(path: &Path) -> Result<Vec<OHLCVBar>> {
    if path.to_string_lossy().ends_with(".dbn.zst") {
        let mut decoder = dbn::decode::dbn::Decoder::from_zstd_file(path)
            .with_context(|| format!("Failed to create zstd DBN decoder for {}", path.display()))?;
        decode_ohlcv_records(&mut decoder, path)
    } else {
        let file = std::fs::File::open(path)
            .with_context(|| format!("Cannot open DBN file: {}", path.display()))?;
        let buf = std::io::BufReader::new(file);
        let mut decoder = dbn::decode::dbn::Decoder::new(buf)
            .with_context(|| format!("Failed to create DBN decoder for {}", path.display()))?;
        decode_ohlcv_records(&mut decoder, path)
    }
}

/// Decode OHLCV records from an already-opened DBN decoder.
///
/// When a `.FUT` parent symbol is used, the DBN file may contain duplicate
/// bars at the same timestamp from overlapping contracts (e.g. ESH5, ESM5 at
/// roll dates). We deduplicate by timestamp, keeping the bar with the highest
/// volume (front-month contract).
fn decode_ohlcv_records<R: std::io::Read>(
    decoder: &mut dbn::decode::dbn::Decoder<R>,
    path: &Path,
) -> Result<Vec<OHLCVBar>> {
    use dbn::OhlcvMsg;
    use std::collections::BTreeMap;

    let price_scale = 1e-9_f64;
    let mut by_ts: BTreeMap<u64, OHLCVBar> = BTreeMap::new();
    let mut total_records = 0_usize;

    while let Some(record) = decoder
        .decode_record::<OhlcvMsg>()
        .with_context(|| format!("Error decoding records from {}", path.display()))?
    {
        total_records += 1;
        let ts_nanos = record.hd.ts_event;
        let volume = record.volume as f64;

        let existing_vol = by_ts.get(&ts_nanos).map(|b| b.volume).unwrap_or(0.0);
        if volume >= existing_vol {
            let timestamp = nanos_to_datetime(ts_nanos);
            by_ts.insert(
                ts_nanos,
                OHLCVBar {
                    timestamp,
                    open: record.open as f64 * price_scale,
                    high: record.high as f64 * price_scale,
                    low: record.low as f64 * price_scale,
                    close: record.close as f64 * price_scale,
                    volume,
                },
            );
        }
    }

    let bars: Vec<OHLCVBar> = by_ts.into_values().collect();
    let deduped = total_records - bars.len();
    if deduped > 0 {
        info!(
            "  Deduplicated {}/{} bars by timestamp from {}",
            deduped,
            total_records,
            path.display()
        );
    }
    Ok(bars)
}

/// Convert nanosecond UNIX timestamp to chrono DateTime<Utc>.
fn nanos_to_datetime(nanos: u64) -> chrono::DateTime<Utc> {
    let secs = (nanos / 1_000_000_000) as i64;
    let subsec_nanos = (nanos % 1_000_000_000) as u32;
    Utc.timestamp_opt(secs, subsec_nanos)
        .single()
        .unwrap_or_else(Utc::now)
}

/// Load all OHLCV bars for a single symbol from .dbn.zst files, sorted chronologically.
///
/// Only loads files from `data_dir/symbol/` to avoid mixing different futures
/// contracts (e.g. ES, NQ, 6E, ZN) into a single price series.
pub fn load_all_bars(data_dir: &Path, symbol: &str) -> Result<Vec<OHLCVBar>> {
    let symbol_dir = data_dir.join(symbol);
    if !symbol_dir.exists() {
        anyhow::bail!(
            "Symbol directory does not exist: {}. Available: {:?}",
            symbol_dir.display(),
            list_subdirs(data_dir),
        );
    }
    let dbn_files = find_dbn_files(&symbol_dir)?;
    if dbn_files.is_empty() {
        anyhow::bail!(
            "No .dbn.zst files found in {}. Run download_baseline first.",
            symbol_dir.display()
        );
    }
    info!("Found {} .dbn.zst files for symbol {} in {}", dbn_files.len(), symbol, symbol_dir.display());

    let mut all_bars = Vec::new();
    for path in &dbn_files {
        match load_bars_from_dbn(path) {
            Ok(bars) => {
                info!("  {} -> {} bars", path.display(), bars.len());
                all_bars.extend(bars);
            }
            Err(e) => {
                warn!("  Skipping {} -- {}", path.display(), e);
            }
        }
    }

    // Sort chronologically
    all_bars.sort_by_key(|b| b.timestamp);
    info!("Total bars loaded for {}: {}", symbol, all_bars.len());
    Ok(all_bars)
}

/// List subdirectory names for error messages.
pub fn list_subdirs(dir: &Path) -> Vec<String> {
    std::fs::read_dir(dir)
        .ok()
        .map(|entries| {
            entries
                .filter_map(|e| e.ok())
                .filter(|e| e.path().is_dir())
                .filter_map(|e| e.file_name().into_string().ok())
                .collect()
        })
        .unwrap_or_default()
}

/// Compute round-trip spread slippage in basis points for a given price.
///
/// Uses the bid-ask spread model from `backtesting/src/slippage.rs`:
/// half-spread = tick_size * spread_ticks / 2.  Round-trip pays the full spread.
/// Converted to bps: spread_price / price * 10_000.
pub fn spread_cost_bps(price: f64, tick_size: f64, spread_ticks: f64) -> f64 {
    if price.abs() < 1e-10 {
        return 0.0;
    }
    let spread_price = tick_size * spread_ticks; // full spread in price units
    spread_price / price * 10_000.0
}