foxhunt/crates/ml/src/benchmark/data_loader.rs

//! DBN Data Loader for ML Training Benchmarks
//!
//! Loads real market data from Databento DBN files and converts to ML training format.
//! Supports parallel loading for performance and comprehensive validation.

use anyhow::{Context, Result};
use dbn::decode::{DbnDecoder, DecodeRecordRef};
use dbn::{OhlcvMsg, VersionUpgradePolicy};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use tracing::{debug, info, warn};

/// Market data point representing a single OHLCV bar
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct MarketDataPoint {
    /// Timestamp in nanoseconds since UNIX epoch
    pub timestamp: i64,
    /// Trading symbol (e.g., "ES.FUT", "NQ.FUT")
    pub symbol: String,
    /// Open price
    pub open: f64,
    /// High price
    pub high: f64,
    /// Low price
    pub low: f64,
    /// Close price
    pub close: f64,
    /// Volume
    pub volume: f64,
}

impl MarketDataPoint {
    /// Validate OHLCV values
    pub fn is_valid(&self) -> bool {
        // Check for NaN or negative values
        if self.open.is_nan()
            || self.high.is_nan()
            || self.low.is_nan()
            || self.close.is_nan()
            || self.volume.is_nan()
        {
            return false;
        }

        if self.open < 0.0
            || self.high < 0.0
            || self.low < 0.0
            || self.close < 0.0
            || self.volume < 0.0
        {
            return false;
        }

        // OHLC consistency checks
        if self.high < self.low {
            return false;
        }

        if self.open < self.low || self.open > self.high {
            return false;
        }

        if self.close < self.low || self.close > self.high {
            return false;
        }

        true
    }
}

/// Statistics about loaded data
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DataStatistics {
    /// Total number of DBN files found
    pub total_files: usize,
    /// Total number of OHLCV bars loaded
    pub total_bars: usize,
    /// Date range of data (start, end)
    pub date_range: (String, String),
    /// Symbols in the dataset
    pub symbols: Vec<String>,
    /// Bars per symbol
    pub bars_per_symbol: HashMap<String, usize>,
    /// Invalid bars encountered
    pub invalid_bars: usize,
}

impl DataStatistics {
    pub fn new() -> Self {
        Self {
            total_files: 0,
            total_bars: 0,
            date_range: (String::new(), String::new()),
            symbols: Vec::new(),
            bars_per_symbol: HashMap::new(),
            invalid_bars: 0,
        }
    }
}

impl Default for DataStatistics {
    fn default() -> Self {
        Self::new()
    }
}

/// DBN data loader for ML training
#[derive(Debug)]
pub struct DbnDataLoader {
    /// Directory containing DBN files
    data_dir: PathBuf,
    /// Symbols to load (empty = load all)
    symbols: Vec<String>,
    /// Statistics about loaded data
    statistics: DataStatistics,
}

impl DbnDataLoader {
    /// Create new DBN data loader
    ///
    /// # Arguments
    /// * `data_dir` - Path to directory containing DBN files
    ///
    /// # Examples
    /// ```no_run
    /// use ml::benchmark::DbnDataLoader;
    /// use std::path::PathBuf;
    ///
    /// let loader = DbnDataLoader::new("test_data/real/databento/ml_training");
    /// ```
    pub fn new<P: AsRef<Path>>(data_dir: P) -> Self {
        Self {
            data_dir: data_dir.as_ref().to_path_buf(),
            symbols: Vec::new(),
            statistics: DataStatistics::new(),
        }
    }

    /// Set symbols to filter (empty = load all)
    pub fn with_symbols(mut self, symbols: Vec<String>) -> Self {
        self.symbols = symbols;
        self
    }

    /// Load all DBN data from directory
    ///
    /// # Returns
    /// Vector of market data points sorted by timestamp
    ///
    /// # Errors
    /// Returns error if directory doesn't exist or files can't be read
    pub fn load_all_data(&mut self) -> Result<Vec<MarketDataPoint>> {
        info!("Loading DBN data from: {}", self.data_dir.display());

        // Get all DBN files
        let dbn_files = self.find_dbn_files()?;

        if dbn_files.is_empty() {
            anyhow::bail!("No DBN files found in {}", self.data_dir.display());
        }

        self.statistics.total_files = dbn_files.len();
        info!("Found {} DBN files", dbn_files.len());

        // Load files sequentially (async not needed for I/O bound operations)
        let mut all_data = Vec::new();
        let mut progress_counter = 0;

        for file_path in dbn_files {
            match self.load_dbn_file(&file_path) {
                Ok(mut data) => {
                    all_data.append(&mut data);
                    progress_counter += 1;

                    // Progress reporting every 50 files
                    if progress_counter % 50 == 0 {
                        info!(
                            "Progress: {}/{} files loaded ({} bars)",
                            progress_counter,
                            self.statistics.total_files,
                            all_data.len()
                        );
                    }
                },
                Err(e) => {
                    warn!("Failed to load {}: {}", file_path.display(), e);
                },
            }
        }

        info!(
            "Loaded {} bars from {} files",
            all_data.len(),
            progress_counter
        );

        // Validate data
        self.validate_and_update_stats(&mut all_data);

        // Sort by timestamp
        all_data.sort_by_key(|point| point.timestamp);

        Ok(all_data)
    }

    /// Load data for a specific symbol
    ///
    /// # Arguments
    /// * `symbol` - Symbol to load (e.g., "ES.FUT")
    ///
    /// # Returns
    /// Vector of market data points for the symbol
    pub fn load_symbol_data(&mut self, symbol: &str) -> Result<Vec<MarketDataPoint>> {
        info!("Loading DBN data for symbol: {}", symbol);

        let dbn_files = self.find_symbol_files(symbol)?;

        if dbn_files.is_empty() {
            anyhow::bail!("No DBN files found for symbol {}", symbol);
        }

        let mut symbol_data = Vec::new();

        for file_path in dbn_files {
            match self.load_dbn_file(&file_path) {
                Ok(mut data) => {
                    symbol_data.append(&mut data);
                },
                Err(e) => {
                    warn!("Failed to load {}: {}", file_path.display(), e);
                },
            }
        }

        // Sort by timestamp
        symbol_data.sort_by_key(|point| point.timestamp);

        Ok(symbol_data)
    }

    /// Get data statistics
    pub fn data_statistics(&self) -> DataStatistics {
        self.statistics.clone()
    }

    /// Find all DBN files in directory
    fn find_dbn_files(&self) -> Result<Vec<PathBuf>> {
        let mut dbn_files = Vec::new();

        let entries = std::fs::read_dir(&self.data_dir).context("Failed to read data directory")?;

        for entry in entries {
            let entry = entry.context("Failed to read directory entry")?;
            let path = entry.path();

            if path.extension().and_then(|s| s.to_str()) == Some("dbn") {
                // Filter by symbol if specified
                if self.symbols.is_empty() || self.should_include_file(&path) {
                    dbn_files.push(path);
                }
            }
        }

        dbn_files.sort();
        Ok(dbn_files)
    }

    /// Find DBN files for a specific symbol
    fn find_symbol_files(&self, symbol: &str) -> Result<Vec<PathBuf>> {
        let mut symbol_files = Vec::new();

        let entries = std::fs::read_dir(&self.data_dir).context("Failed to read data directory")?;

        for entry in entries {
            let entry = entry.context("Failed to read directory entry")?;
            let path = entry.path();

            if let Some(file_name) = path.file_name().and_then(|s| s.to_str()) {
                if file_name.starts_with(symbol)
                    && path.extension().and_then(|s| s.to_str()) == Some("dbn")
                {
                    symbol_files.push(path);
                }
            }
        }

        symbol_files.sort();
        Ok(symbol_files)
    }

    /// Check if file should be included based on symbol filter
    fn should_include_file(&self, path: &Path) -> bool {
        if self.symbols.is_empty() {
            return true;
        }

        let file_name = match path.file_name().and_then(|s| s.to_str()) {
            Some(name) => name,
            None => return false,
        };

        self.symbols
            .iter()
            .any(|symbol| file_name.starts_with(symbol))
    }

    /// Load a single DBN file
    fn load_dbn_file(&self, path: &Path) -> Result<Vec<MarketDataPoint>> {
        debug!("Loading DBN file: {}", path.display());

        // Extract symbol from filename (e.g., "ES.FUT_ohlcv-1m_2024-01-02.dbn" -> "ES.FUT")
        let symbol = self.extract_symbol_from_filename(path)?;

        // Create DBN decoder using from_file (more efficient than BufReader)
        let mut decoder = DbnDecoder::from_file(path).context(format!(
            "Failed to create DBN decoder for file: {}",
            path.display()
        ))?;

        // Enable version upgrades for compatibility
        decoder.set_upgrade_policy(VersionUpgradePolicy::UpgradeToV3)?;

        let mut data_points = Vec::new();

        // Read records using decode_record_ref (zero-copy)
        while let Some(record_ref) = decoder
            .decode_record_ref()
            .context("Failed to decode DBN record")?
        {
            // Extract OHLCV message if present
            if let Some(ohlcv) = record_ref.get::<OhlcvMsg>() {
                // Convert DBN record to MarketDataPoint
                let point = self.convert_ohlcv_record(ohlcv, &symbol)?;
                data_points.push(point);
            }
        }

        debug!("Loaded {} bars from {}", data_points.len(), path.display());
        Ok(data_points)
    }

    /// Extract symbol from DBN filename
    fn extract_symbol_from_filename(&self, path: &Path) -> Result<String> {
        let file_name = path
            .file_name()
            .and_then(|s| s.to_str())
            .context("Invalid filename")?;

        // Format: "SYMBOL_ohlcv-1m_DATE.dbn"
        let symbol = file_name
            .split('_')
            .next()
            .context("Failed to extract symbol from filename")?
            .to_string();

        Ok(symbol)
    }

    /// Convert DBN OHLCV record to MarketDataPoint
    fn convert_ohlcv_record(&self, record: &OhlcvMsg, symbol: &str) -> Result<MarketDataPoint> {
        // DBN 0.42 uses hd.ts_event for timestamp (nanoseconds since UNIX epoch)
        // Prices are in fixed-point (9 decimal places)
        let price_scale = 1_000_000_000.0;

        let point = MarketDataPoint {
            timestamp: record.hd.ts_event as i64,
            symbol: symbol.to_string(),
            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: record.volume as f64,
        };

        Ok(point)
    }

    /// Validate data and update statistics
    fn validate_and_update_stats(&mut self, data: &mut Vec<MarketDataPoint>) {
        self.statistics.total_bars = data.len();

        if data.is_empty() {
            return;
        }

        // Count bars per symbol
        let mut symbol_counts: HashMap<String, usize> = HashMap::new();
        for point in data.iter() {
            *symbol_counts.entry(point.symbol.clone()).or_insert(0) += 1;
        }

        // Validate and filter invalid bars
        let mut invalid_count = 0;
        data.retain(|point| {
            if !point.is_valid() {
                invalid_count += 1;
                warn!(
                    "Invalid bar: symbol={}, timestamp={}, open={}, high={}, low={}, close={}, volume={}",
                    point.symbol, point.timestamp, point.open, point.high, point.low, point.close, point.volume
                );
                false
            } else {
                true
            }
        });

        self.statistics.invalid_bars = invalid_count;
        self.statistics.bars_per_symbol = symbol_counts.clone();
        self.statistics.symbols = symbol_counts.keys().cloned().collect();
        self.statistics.symbols.sort();

        // Calculate date range
        if !data.is_empty() {
            let start_ts = data.first().map(|p| p.timestamp).unwrap_or(0);
            let end_ts = data.last().map(|p| p.timestamp).unwrap_or(0);

            self.statistics.date_range = (
                Self::timestamp_to_date_string(start_ts),
                Self::timestamp_to_date_string(end_ts),
            );
        }

        // Check for missing data gaps
        self.check_for_gaps(data);
    }

    /// Check for significant data gaps
    fn check_for_gaps(&self, data: &[MarketDataPoint]) {
        if data.len() < 2 {
            return;
        }

        // Expected interval: 1 minute = 60 seconds = 60_000_000_000 nanoseconds
        let expected_interval_ns = 60_000_000_000_i64;
        let gap_threshold = expected_interval_ns * 10; // 10 minutes

        for i in 1..data.len() {
            let gap = data[i].timestamp - data[i - 1].timestamp;
            if gap > gap_threshold {
                let gap_minutes = gap / 60_000_000_000;
                warn!(
                    "Data gap detected: {} minutes between {} and {}",
                    gap_minutes,
                    Self::timestamp_to_date_string(data[i - 1].timestamp),
                    Self::timestamp_to_date_string(data[i].timestamp)
                );
            }
        }
    }

    /// Convert timestamp to date string
    fn timestamp_to_date_string(timestamp_ns: i64) -> String {
        use chrono::{DateTime, Utc};

        // Convert nanoseconds to seconds
        let timestamp_secs = timestamp_ns / 1_000_000_000;

        match DateTime::from_timestamp(timestamp_secs, 0) {
            Some(dt) => {
                let utc: DateTime<Utc> = dt;
                utc.format("%Y-%m-%d").to_string()
            },
            None => "Invalid timestamp".to_owned(),
        }
    }
}

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

    #[test]
    fn test_market_data_point_validation() {
        // Valid point
        let valid_point = MarketDataPoint {
            timestamp: 1704067200000000000,
            symbol: "ES.FUT".to_owned(),
            open: 4750.0,
            high: 4755.0,
            low: 4745.0,
            close: 4752.0,
            volume: 1000.0,
        };
        assert!(valid_point.is_valid());

        // Invalid: NaN values
        let nan_point = MarketDataPoint {
            timestamp: 1704067200000000000,
            symbol: "ES.FUT".to_owned(),
            open: f64::NAN,
            high: 4755.0,
            low: 4745.0,
            close: 4752.0,
            volume: 1000.0,
        };
        assert!(!nan_point.is_valid());

        // Invalid: negative values
        let negative_point = MarketDataPoint {
            timestamp: 1704067200000000000,
            symbol: "ES.FUT".to_owned(),
            open: -4750.0,
            high: 4755.0,
            low: 4745.0,
            close: 4752.0,
            volume: 1000.0,
        };
        assert!(!negative_point.is_valid());

        // Invalid: high < low
        let invalid_range = MarketDataPoint {
            timestamp: 1704067200000000000,
            symbol: "ES.FUT".to_owned(),
            open: 4750.0,
            high: 4740.0, // High < Low
            low: 4745.0,
            close: 4752.0,
            volume: 1000.0,
        };
        assert!(!invalid_range.is_valid());
    }

    #[test]
    fn test_dbn_data_loader_creation() {
        let loader = DbnDataLoader::new("test_data/real/databento/ml_training");
        assert_eq!(loader.symbols.len(), 0);
        assert_eq!(loader.statistics.total_files, 0);
    }

    #[test]
    fn test_dbn_data_loader_with_symbols() {
        let symbols = vec!["ES.FUT".to_owned(), "NQ.FUT".to_owned()];
        let loader = DbnDataLoader::new("test_data/real/databento/ml_training")
            .with_symbols(symbols.clone());
        assert_eq!(loader.symbols, symbols);
    }

    #[test]
    fn test_extract_symbol_from_filename() {
        let loader = DbnDataLoader::new("test_data");

        let path = Path::new("ES.FUT_ohlcv-1m_2024-01-02.dbn");
        let symbol = loader.extract_symbol_from_filename(path).unwrap();
        assert_eq!(symbol, "ES.FUT");

        let path2 = Path::new("NQ.FUT_ohlcv-1m_2024-01-03.dbn");
        let symbol2 = loader.extract_symbol_from_filename(path2).unwrap();
        assert_eq!(symbol2, "NQ.FUT");
    }

    #[test]
    fn test_data_statistics_default() {
        let stats = DataStatistics::default();
        assert_eq!(stats.total_files, 0);
        assert_eq!(stats.total_bars, 0);
        assert_eq!(stats.symbols.len(), 0);
    }
}