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`.

#[allow(dead_code)]
pub mod completion;

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.
///
/// **Calendar-spread filter:** when `stype_in=parent` is used to request
/// `ES.FUT`, Databento returns BOTH outright contracts (`ESH4`, `ESM4`, ...)
/// AND calendar-spread instruments (`ESH4-ESM4`, `ESM4-ESU4`, ...) in the
/// same DBN stream. Spreads trade at the price *difference* between adjacent
/// contracts (~$60-150 due to cost-of-carry) so their bars look like
/// ridiculously low ES quotes (e.g. close=$61 vs prev=$5141). The
/// `instrument_id → symbol` resolution comes from the DBN file's metadata
/// `symbol_map()`; any symbol containing `-` is a spread and is rejected at
/// parse time. Pre-fix this contamination was 23% of records on Q1 2024 ES
/// and ~8800 spread bars survived dedup-by-volume across the full dataset
/// because spreads occasionally carry higher minute-volume than the outright
/// during low-liquidity overnight periods.
///
/// **Same-timestamp dedup:** within outright records, multiple contracts
/// (front-month + back-month, e.g. `ESH4` + `ESM4` during rollover week)
/// can both report bars at the same `ts_event`. Keep the highest-volume bar
/// at each timestamp — that's the active front month.
fn decode_ohlcv_records<R: std::io::Read>(
    decoder: &mut dbn::decode::dbn::Decoder<R>,
    path: &Path,
) -> Result<Vec<OHLCVBar>> {
    use dbn::decode::DbnMetadata;
    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;
    let mut skipped_spread = 0_usize;

    // Build instrument_id → symbol resolver from the DBN metadata. This is the
    // authoritative source — the same map Databento's Python client uses for
    // its `symbol` column. Empty for files that lack symbology mappings (older
    // DBN versions or non-parent-symbol requests); in that case the spread
    // filter below silently no-ops, and the bar-level sanity gate in
    // `sanitize_bars` remains as a defense-in-depth backstop.
    let symbol_map = decoder
        .metadata()
        .symbol_map()
        .with_context(|| format!("Failed to build symbol map for {}", path.display()))?;

    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 timestamp = nanos_to_datetime(ts_nanos);

        // Resolve symbol via TsSymbolMap (date + instrument_id → symbol). Drop
        // any record whose symbol contains a hyphen — those are calendar
        // spread instruments and their prices are not on the underlying's
        // scale. If the symbol can't be resolved (no mapping for this date),
        // skip the record rather than guess.
        let secs_since_epoch = (ts_nanos / 1_000_000_000) as i64;
        let date = time::OffsetDateTime::from_unix_timestamp(secs_since_epoch)
            .map(|dt| dt.date())
            .ok();
        let symbol = match date.and_then(|d| symbol_map.get(d, record.hd.instrument_id)) {
            Some(s) => s.as_str(),
            None => {
                skipped_spread += 1;
                continue;
            }
        };
        if symbol.contains('-') {
            skipped_spread += 1;
            continue;
        }

        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 {
            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() - skipped_spread;
    if skipped_spread > 0 || deduped > 0 {
        info!(
            "  {}: {} records → {} bars (skipped {} spreads, deduped {} same-ts)",
            path.display(),
            total_records,
            bars.len(),
            skipped_spread,
            deduped,
        );
    }
    Ok(bars)
}

/// Convert nanosecond UNIX timestamp to chrono `DateTime`<Utc>.
fn nanos_to_datetime(nanos: u64) -> chrono::DateTime<Utc> {
    #[allow(clippy::integer_division)]
    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.
#[allow(clippy::cognitive_complexity)]
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);

    // Bar-level sanity gate. Drops bars with non-positive OHLC, high<low,
    // non-finite values, or close moves outside [0.5×, 2.0×] of the previous
    // bar (catches DBN parse glitches, broker tick errors, near-zero-open bars
    // that produce extreme log returns downstream). This is the FIRST defense
    // line; subsequent layers (`safe_log_return` clamp, `validate_features`
    // pre-norm bound, `NormStats::normalize` post-norm clamp,
    // `validate_normalized_features` final gate) are defense-in-depth backstops
    // for shapes this filter doesn't anticipate.
    let pre = all_bars.len();
    all_bars = sanitize_bars(all_bars);
    let dropped = pre.saturating_sub(all_bars.len());
    if dropped > 0 {
        warn!(
            "load_all_bars: dropped {} corrupt bars out of {} for symbol {} \
             (non-positive OHLC, high<low, or extreme close-to-prev ratio)",
            dropped, pre, symbol
        );
    }
    info!("Total bars loaded for {} (post-sanitize): {}", symbol, all_bars.len());
    Ok(all_bars)
}

/// Drop bars that fail OHLC sanity checks. Sequential — each bar is validated
/// against its sorted predecessor (`prev.close → bar.open` ratio).
fn sanitize_bars(bars: Vec<OHLCVBar>) -> Vec<OHLCVBar> {
    let mut out: Vec<OHLCVBar> = Vec::with_capacity(bars.len());
    let mut prev_close: Option<f64> = None;
    for bar in bars {
        // Per-bar invariants
        if !(bar.open.is_finite() && bar.high.is_finite() && bar.low.is_finite() && bar.close.is_finite()) {
            warn!("sanitize: drop non-finite OHLC at {}", bar.timestamp);
            continue;
        }
        if bar.open <= 0.0 || bar.high <= 0.0 || bar.low <= 0.0 || bar.close <= 0.0 {
            warn!("sanitize: drop non-positive OHLC at {} (o={} h={} l={} c={})",
                bar.timestamp, bar.open, bar.high, bar.low, bar.close);
            continue;
        }
        if bar.high < bar.low {
            warn!("sanitize: drop high<low at {} (h={} l={})",
                bar.timestamp, bar.high, bar.low);
            continue;
        }
        // Cross-bar invariant: close moves within ±100% of previous close.
        // Real markets stay well inside ±10% per bar; ±100% catches genuine
        // corruption while not rejecting any legitimate move (futures gap
        // limits are far below this).
        if let Some(prev) = prev_close {
            let ratio = bar.close / prev;
            if !ratio.is_finite() || !(0.5..=2.0).contains(&ratio) {
                warn!("sanitize: drop extreme close ratio at {} (close={} prev_close={} ratio={:.3e})",
                    bar.timestamp, bar.close, prev, ratio);
                continue;
            }
        }
        prev_close = Some(bar.close);
        out.push(bar);
    }
    out
}

/// 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
}