foxhunt/crates/ml-features/src/trades_loader.rs

//! Trade Data Loader for DBN `Schema::Trades`
//!
//! Loads Databento trade records from `.dbn` or `.dbn.zst` files and provides
//! time-windowed access for feeding into `OFICalculator::feed_trade()` (VPIN,
//! Kyle's Lambda).

use std::fs::File;
use std::io::BufReader;
use std::path::Path;

use dbn::decode::{DbnDecoder, DecodeRecordRef};
use dbn::RecordRefEnum;
use serde::{Deserialize, Serialize};

use crate::MLError;

/// A single parsed trade from a DBN file.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct DbnTrade {
    /// Event timestamp in nanoseconds since Unix epoch
    pub timestamp: u64,
    /// Trade price (f64, converted from i64 scaled by 1e-9)
    pub price: f64,
    /// Trade size in contracts
    pub volume: u64,
    /// True if buyer-initiated (side == 'B'), false if seller-initiated
    pub is_buy: bool,
    /// Databento instrument_id — identifies which contract month this trade belongs to
    pub instrument_id: u32,
}

/// Load all trades from a `.dbn` or `.dbn.zst` file (synchronous).
///
/// Auto-detects zstd compression by file extension. Returns trades sorted
/// by timestamp (as they appear in the file — DBN files are already ordered).
pub fn load_trades_sync(file_path: &Path) -> Result<Vec<DbnTrade>, MLError> {
    let file = File::open(file_path).map_err(|e| {
        MLError::InsufficientData(format!("Failed to open trade file {:?}: {}", file_path, e))
    })?;

    let is_zstd = file_path
        .to_string_lossy()
        .ends_with(".dbn.zst");

    let reader: Box<dyn std::io::Read> = if is_zstd {
        Box::new(
            zstd::Decoder::new(BufReader::new(file)).map_err(|e| {
                MLError::InsufficientData(format!("Failed to create zstd decoder: {}", e))
            })?,
        )
    } else {
        Box::new(BufReader::new(file))
    };

    let mut decoder = DbnDecoder::new(reader).map_err(|e| {
        MLError::InsufficientData(format!("Failed to create DBN decoder: {}", e))
    })?;

    let mut trades = Vec::new();

    loop {
        match decoder.decode_record_ref() {
            Ok(Some(record)) => {
                let record_enum = record.as_enum().map_err(|e| {
                    MLError::InsufficientData(format!("Failed to convert DBN record: {}", e))
                })?;

                if let RecordRefEnum::Trade(trade) = record_enum {
                    // Prices are i64 scaled by 1e-9 per DBN specification
                    let price_f64 = trade.price as f64 * 1e-9;
                    // Side: trade.side == b'B' as i8 => Buy, b'A' as i8 => Sell
                    let is_buy = trade.side == b'B' as i8;

                    trades.push(DbnTrade {
                        timestamp: trade.hd.ts_event,
                        price: price_f64,
                        volume: u64::from(trade.size),
                        is_buy,
                        instrument_id: trade.hd.instrument_id,
                    });
                }
            }
            Ok(None) => break,
            Err(e) => {
                return Err(MLError::InsufficientData(format!(
                    "DBN decode error: {}",
                    e
                )));
            }
        }
    }

    // Sort by timestamp (should already be ordered, but ensure correctness)
    trades.sort_by_key(|t| t.timestamp);

    Ok(trades)
}

/// Get the sub-slice of trades within `[bar_start_ns, bar_end_ns)`.
///
/// Uses binary search for O(log n) lookup on the pre-sorted trades slice.
pub fn get_trades_for_bar(trades: &[DbnTrade], bar_start_ns: u64, bar_end_ns: u64) -> &[DbnTrade] {
    if trades.is_empty() {
        return &[];
    }

    // First trade with timestamp >= bar_start_ns
    let start = trades.partition_point(|t| t.timestamp < bar_start_ns);
    // First trade with timestamp >= bar_end_ns
    let end = trades.partition_point(|t| t.timestamp < bar_end_ns);

    trades.get(start..end).unwrap_or(&[])
}

/// Default volume bar size: 100 contracts per bar for ES futures.
pub const DEFAULT_VOLUME_BAR_SIZE: u64 = 100;

/// Build volume bars from a sorted slice of trades.
///
/// Aggregates trades until cumulative volume reaches `bar_size` contracts.
/// Each bar records OHLCV from the trades within that volume bucket.
/// Trades MUST be pre-sorted by timestamp and pre-filtered to a single instrument_id.
///
/// Returns bars sorted by timestamp (bar timestamp = last trade in bar).
pub fn build_volume_bars(trades: &[DbnTrade], bar_size: u64) -> Vec<ml_core::types::OHLCVBar> {
    if trades.is_empty() || bar_size == 0 {
        return Vec::new();
    }

    let mut bars = Vec::new();
    let mut bar_open = 0.0_f64;
    let mut bar_high = 0.0_f64;
    let mut bar_low = 0.0_f64;
    #[allow(unused_assignments)]
    let mut bar_close = 0.0_f64;
    let mut bar_volume: u64 = 0;
    #[allow(unused_assignments)]
    let mut bar_ts: u64 = 0;

    for trade in trades {
        if trade.price <= 0.0 || trade.volume == 0 {
            continue;
        }

        if bar_volume == 0 {
            bar_open = trade.price;
            bar_high = trade.price;
            bar_low = trade.price;
        }

        bar_high = bar_high.max(trade.price);
        bar_low = bar_low.min(trade.price);
        bar_close = trade.price;
        bar_ts = trade.timestamp;
        bar_volume += trade.volume;

        if bar_volume >= bar_size {
            let ts_secs = (bar_ts / 1_000_000_000) as i64;
            let ts_nanos = (bar_ts % 1_000_000_000) as u32;
            let timestamp = chrono::DateTime::<chrono::Utc>::from_timestamp(ts_secs, ts_nanos)
                .unwrap_or_else(|| chrono::Utc::now());

            bars.push(ml_core::types::OHLCVBar {
                timestamp,
                open: bar_open,
                high: bar_high,
                low: bar_low,
                close: bar_close,
                volume: bar_volume as f64,
            });

            bar_volume = 0;
        }
    }

    bars
}

/// Filter trades to keep only the front-month contract (highest cumulative volume).
///
/// Databento DBN files contain trades from ALL contract months for a symbol.
/// This function groups trades by `instrument_id` and keeps only the instrument
/// with the highest total volume — the front-month continuous contract.
pub fn filter_front_month(trades: &[DbnTrade]) -> Vec<DbnTrade> {
    if trades.is_empty() {
        return Vec::new();
    }

    let mut volume_by_id: std::collections::HashMap<u32, u64> = std::collections::HashMap::new();
    for trade in trades {
        *volume_by_id.entry(trade.instrument_id).or_insert(0) += trade.volume;
    }

    let front_month_id = volume_by_id
        .into_iter()
        .max_by_key(|&(_, vol)| vol)
        .map(|(id, _)| id)
        .unwrap_or(0);

    tracing::info!(
        instrument_id = front_month_id,
        "Front-month contract selected (highest volume)"
    );

    trades.iter()
        .filter(|t| t.instrument_id == front_month_id)
        .cloned()
        .collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_trade(timestamp: u64, price: f64, volume: u64, is_buy: bool) -> DbnTrade {
        DbnTrade {
            timestamp,
            price,
            volume,
            is_buy,
            instrument_id: 0,
        }
    }

    #[test]
    fn test_get_trades_for_bar_empty() {
        let trades: &[DbnTrade] = &[];
        let result = get_trades_for_bar(trades, 1000, 2000);
        assert!(result.is_empty());
    }

    #[test]
    fn test_get_trades_for_bar_binary_search() {
        let trades = vec![
            make_trade(1000, 100.0, 10, true),
            make_trade(2000, 101.0, 20, false),
            make_trade(3000, 102.0, 30, true),
            make_trade(4000, 103.0, 40, false),
            make_trade(5000, 104.0, 50, true),
        ];

        // Window [2000, 4000) should return trades at timestamps 2000 and 3000
        let result = get_trades_for_bar(&trades, 2000, 4000);
        assert_eq!(result.len(), 2);
        assert_eq!(result.first().map(|t| t.timestamp), Some(2000));
        assert_eq!(result.last().map(|t| t.timestamp), Some(3000));

        // Window [1000, 6000) should return all 5 trades
        let all = get_trades_for_bar(&trades, 1000, 6000);
        assert_eq!(all.len(), 5);

        // Window [5001, 6000) should return nothing (past all trades)
        let none = get_trades_for_bar(&trades, 5001, 6000);
        assert!(none.is_empty());

        // Window [0, 1000) should return nothing (before all trades)
        let before = get_trades_for_bar(&trades, 0, 1000);
        assert!(before.is_empty());
    }

    #[test]
    fn test_build_volume_bars_basic() {
        let trades = vec![
            DbnTrade { timestamp: 1000, price: 5100.0, volume: 30, is_buy: true, instrument_id: 1 },
            DbnTrade { timestamp: 2000, price: 5101.0, volume: 40, is_buy: false, instrument_id: 1 },
            DbnTrade { timestamp: 3000, price: 5099.0, volume: 30, is_buy: true, instrument_id: 1 },
            DbnTrade { timestamp: 4000, price: 5102.0, volume: 50, is_buy: true, instrument_id: 1 },
            DbnTrade { timestamp: 5000, price: 5103.0, volume: 50, is_buy: false, instrument_id: 1 },
        ];
        let bars = super::build_volume_bars(&trades, 100);
        assert_eq!(bars.len(), 2, "Should produce 2 bars from 200 contracts");
        assert!((bars[0].open - 5100.0).abs() < 0.01);
        assert!((bars[0].high - 5101.0).abs() < 0.01);
        assert!((bars[0].low - 5099.0).abs() < 0.01);
        assert!((bars[0].close - 5099.0).abs() < 0.01);
        assert!((bars[0].volume - 100.0).abs() < 0.01);
        assert!((bars[1].open - 5102.0).abs() < 0.01);
        assert!((bars[1].close - 5103.0).abs() < 0.01);
    }

    #[test]
    fn test_build_volume_bars_large_trade() {
        let trades = vec![
            DbnTrade { timestamp: 1000, price: 5100.0, volume: 500, is_buy: true, instrument_id: 1 },
        ];
        let bars = super::build_volume_bars(&trades, 100);
        assert_eq!(bars.len(), 1);
        assert!((bars[0].volume - 500.0).abs() < 0.01);
    }

    #[test]
    fn test_build_volume_bars_empty() {
        let bars = super::build_volume_bars(&[], 100);
        assert!(bars.is_empty());
    }

    #[test]
    fn test_build_volume_bars_partial_dropped() {
        let trades = vec![
            DbnTrade { timestamp: 1000, price: 5100.0, volume: 50, is_buy: true, instrument_id: 1 },
            DbnTrade { timestamp: 2000, price: 5101.0, volume: 25, is_buy: false, instrument_id: 1 },
        ];
        let bars = super::build_volume_bars(&trades, 100);
        assert!(bars.is_empty(), "Partial bar should be dropped");
    }

    #[test]
    fn test_filter_front_month() {
        let trades = vec![
            DbnTrade { timestamp: 1000, price: 5100.0, volume: 100, is_buy: true, instrument_id: 1 },
            DbnTrade { timestamp: 2000, price: 5101.0, volume: 200, is_buy: false, instrument_id: 1 },
            DbnTrade { timestamp: 1500, price: 5160.0, volume: 30, is_buy: true, instrument_id: 2 },
            DbnTrade { timestamp: 2500, price: 5161.0, volume: 20, is_buy: false, instrument_id: 2 },
            DbnTrade { timestamp: 3000, price: 61.25, volume: 10, is_buy: true, instrument_id: 3 },
        ];
        let filtered = super::filter_front_month(&trades);
        assert_eq!(filtered.len(), 2, "Should keep only front month (id=1)");
        assert!(filtered.iter().all(|t| t.instrument_id == 1));
    }

    #[test]
    fn test_filter_front_month_empty() {
        let filtered = super::filter_front_month(&[]);
        assert!(filtered.is_empty());
    }

    #[test]
    fn test_filter_front_month_per_file_pattern() {
        // Simulates quarterly roll: Q1 has front-month id=100, Q2 has id=200.
        // If you filter globally, only one quarter survives (the one with more volume).
        // Correct behavior: filter PER FILE, then concatenate both quarters.
        let q1_trades = vec![
            DbnTrade { timestamp: 1000, price: 5100.0, volume: 200, is_buy: true, instrument_id: 100 },
            DbnTrade { timestamp: 2000, price: 5101.0, volume: 100, is_buy: false, instrument_id: 100 },
            DbnTrade { timestamp: 1500, price: 5160.0, volume: 30, is_buy: true, instrument_id: 101 }, // back month
        ];
        let q2_trades = vec![
            DbnTrade { timestamp: 5000, price: 5200.0, volume: 300, is_buy: true, instrument_id: 200 },
            DbnTrade { timestamp: 6000, price: 5201.0, volume: 150, is_buy: false, instrument_id: 200 },
            DbnTrade { timestamp: 5500, price: 5260.0, volume: 20, is_buy: true, instrument_id: 201 }, // back month
        ];

        // Per-file filtering (CORRECT): each file keeps its own front month
        let q1_filtered = super::filter_front_month(&q1_trades);
        let q2_filtered = super::filter_front_month(&q2_trades);
        assert_eq!(q1_filtered.len(), 2, "Q1 should keep id=100 (2 trades)");
        assert_eq!(q2_filtered.len(), 2, "Q2 should keep id=200 (2 trades)");
        assert!(q1_filtered.iter().all(|t| t.instrument_id == 100));
        assert!(q2_filtered.iter().all(|t| t.instrument_id == 200));
        let combined_len = q1_filtered.len() + q2_filtered.len();
        assert_eq!(combined_len, 4, "Per-file: 4 trades from both quarters");

        // Global filtering (WRONG): merges all, picks single highest-volume id
        let mut all_trades = q1_trades.clone();
        all_trades.extend(q2_trades.clone());
        let global_filtered = super::filter_front_month(&all_trades);
        // id=200 has 450 volume vs id=100's 300 → global keeps only Q2
        assert!(global_filtered.iter().all(|t| t.instrument_id == 200),
            "Global filter incorrectly drops Q1 entirely");
        assert_eq!(global_filtered.len(), 2, "Global filter loses Q1 data");

        // This test documents WHY per-file filtering is required:
        // global loses 50% of the data (one quarter).
        assert!(combined_len > global_filtered.len(),
            "Per-file ({combined_len}) must produce more data than global ({})",
            global_filtered.len());
    }
}