//! Data Loading for DQN Training
//!
//! Methods for loading market data from Parquet files and DBN files,
//! including feature extraction, preprocessing, and train/val splitting.

use std::path::Path;
use std::sync::Arc;

use anyhow::{Context, Result};
use tracing::{debug, info, warn};

use crate::features::extraction::OHLCVBar;
use crate::preprocessing::{clip_outliers_with_bounds, compute_clip_bounds, compute_log_returns, windowed_normalize, PreprocessConfig};
use crate::training_pipeline::FinancialFeatures;
use crate::TrainingMetrics;

use crate::features::extraction::FeatureVector;
use super::trainer::DQNTrainer;

/// Detect zstd-compressed files by magic bytes (0x28B52FFD).
/// File extension is unreliable — Databento downloads often use `.dbn` for zstd-compressed data.
fn is_zstd_file(path: &Path) -> Result<bool> {
    let mut file = std::fs::File::open(path)?;
    let mut magic = [0_u8; 4];
    use std::io::Read;
    if file.read_exact(&mut magic).is_ok() {
        Ok(magic == [0x28, 0xB5, 0x2F, 0xFD])
    } else {
        Ok(false)
    }
}

/// Recursively collect all .dbn and .dbn.zst files from a directory and its subdirectories.
fn collect_dbn_files_recursive(dir: &Path) -> Vec<std::path::PathBuf> {
    let mut files = Vec::new();
    if let Ok(entries) = std::fs::read_dir(dir) {
        for entry in entries.flatten() {
            let path = entry.path();
            if path.is_dir() {
                files.extend(collect_dbn_files_recursive(&path));
            } else if path.extension().and_then(|s| s.to_str()) == Some("dbn")
                || path.to_string_lossy().ends_with(".dbn.zst")
            {
                files.push(path);
            } else {
                // Skip non-DBN files
            }
        }
    }
    files
}

impl DQNTrainer {
    /// Train DQN on market data from Parquet file (Wave 12 Group 3)
    ///
    /// # Arguments
    ///
    /// * `parquet_path` - Path to Parquet file containing OHLCV bars
    /// * `checkpoint_callback` - Callback for saving checkpoints (epoch, model_data) -> `Result<String>`
    ///
    /// # Returns
    ///
    /// Training metrics (loss, accuracy, gradient norms, Q-values)
    pub async fn train_from_parquet<F>(
        &mut self,
        parquet_path: &str,
        checkpoint_callback: F,
    ) -> Result<TrainingMetrics>
    where
        F: FnMut(usize, Vec<u8>, bool) -> Result<String> + Send,
    {
        info!("Starting DQN training from Parquet file: {}", parquet_path);

        // Load market data from Parquet file (returns train/val split)
        let (mut training_data, mut validation_data) =
            self.load_training_data_from_parquet(parquet_path).await?;

        info!(
            "Loaded {} training samples, {} validation samples",
            training_data.len(),
            validation_data.len()
        );


        // Calculate feature statistics from all training samples
        info!("📊 Calculating feature statistics from {} training samples...", training_data.len());
        let feature_stats = self.calculate_feature_statistics(&training_data)?;
        self.feature_stats = Some(feature_stats.clone());
        info!("✅ Feature statistics calculated: {} features normalized", feature_stats.mean.len());

        // Normalize all training and validation samples BEFORE training starts
        info!("📊 Normalizing all samples with z-score normalization...");
        self.normalize_dataset(&mut training_data)?;
        self.normalize_dataset(&mut validation_data)?;
        info!("✅ Dataset normalization complete");

        // Store normalized validation data for validation loss computation
        self.val_data = validation_data;

        // Use the same training loop as DBN-based training
        self.train_with_data_full_loop(&training_data, checkpoint_callback)
            .await
    }

    /// Load training data from Parquet file (Wave 12 Group 3)
    /// Returns (train_data, val_data) with 80/20 split
    pub async fn load_training_data_from_parquet(
        &mut self,
        parquet_path: &str,
    ) -> Result<(
        Vec<(FeatureVector, Vec<f64>)>,
        Vec<(FeatureVector, Vec<f64>)>,
    )> {
        use arrow::array::{Array, Float64Array, PrimitiveArray, UInt64Array};
        use arrow::datatypes::TimestampNanosecondType;
        use arrow::record_batch::RecordBatch;
        use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
        use std::fs::File;

        info!("Loading Parquet file: {}", parquet_path);

        // TRY CACHE FIRST
        if let Some(cache_dir) = &self.feature_cache_dir {
            let parquet_path_obj = Path::new(parquet_path);
            let mbp10 = Path::new("test_data/mbp10");

            info!("🔍 Checking feature cache...");

            // Calculate cache key
            match crate::feature_cache::calculate_cache_key(
                parquet_path_obj,
                mbp10,
                50, // warmup period
            ) {
                Ok(cache_key) => {
                    // Try to load from cache
                    match crate::feature_cache::load_features_from_cache(
                        cache_dir,
                        &cache_key,
                    ).await {
                        Ok(Some(features)) => {
                            info!("🚀 Loaded {} feature vectors from cache", features.len());
                            info!("   ⚡ Savings vs compute: ~2m 25s");

                            // Split into train/val (80/20)
                            let split_idx = (features.len() as f64 * 0.8) as usize;
                            let train_data: Vec<(FeatureVector, Vec<f64>)> = features[..split_idx]
                                .iter()
                                .map(|f| (*f, vec![]))
                                .collect();
                            let val_data: Vec<(FeatureVector, Vec<f64>)> = features[split_idx..]
                                .iter()
                                .map(|f| (*f, vec![]))
                                .collect();

                            return Ok((train_data, val_data));
                        }
                        Ok(None) => {
                            debug!("Cache miss, computing features from scratch...");
                        }
                        Err(e) => {
                            debug!("Cache load unavailable: {}, computing features from scratch", e);
                        }
                    }
                }
                Err(e) => {
                    debug!("Cache key calculation not possible: {}, skipping cache", e);
                }
            }
        }

        // FALLBACK: Original feature extraction
        info!("📊 Computing features from scratch...");

        // Open Parquet file
        let file = File::open(parquet_path)
            .with_context(|| format!("Failed to open Parquet file: {}", parquet_path))?;

        // Create Parquet reader
        let builder = ParquetRecordBatchReaderBuilder::try_new(file)
            .with_context(|| "Failed to create Parquet reader")?;

        let reader = builder
            .build()
            .with_context(|| "Failed to build Parquet reader")?;

        // Read all batches
        let mut all_ohlcv_bars = Vec::new();

        for batch_result in reader {
            let batch: RecordBatch = batch_result.with_context(|| "Failed to read record batch")?;

            // Extract columns by name (schema-agnostic approach)
            // Required columns: timestamp_ns (or ts_event), open, high, low, close, volume

            // Try timestamp_ns first (our schema), fallback to ts_event (Databento schema)
            let timestamp_col = batch
                .column_by_name("timestamp_ns")
                .or_else(|| batch.column_by_name("ts_event"))
                .ok_or_else(|| {
                    anyhow::anyhow!(
                        "Missing timestamp column. Expected 'timestamp_ns' or 'ts_event'"
                    )
                })?;

            let timestamps = timestamp_col
                .as_any()
                .downcast_ref::<PrimitiveArray<TimestampNanosecondType>>()
                .ok_or_else(|| {
                    anyhow::anyhow!(
                        "Failed to downcast timestamp column. Expected Timestamp(Nanosecond), got: {:?}",
                        timestamp_col.data_type()
                    )
                })?;

            // Extract OHLCV columns by name
            let opens = batch
                .column_by_name("open")
                .ok_or_else(|| anyhow::anyhow!("Missing 'open' column in Parquet schema"))?
                .as_any()
                .downcast_ref::<Float64Array>()
                .ok_or_else(|| anyhow::anyhow!("Invalid 'open' column type. Expected Float64"))?;

            let highs = batch
                .column_by_name("high")
                .ok_or_else(|| anyhow::anyhow!("Missing 'high' column in Parquet schema"))?
                .as_any()
                .downcast_ref::<Float64Array>()
                .ok_or_else(|| anyhow::anyhow!("Invalid 'high' column type. Expected Float64"))?;

            let lows = batch
                .column_by_name("low")
                .ok_or_else(|| anyhow::anyhow!("Missing 'low' column in Parquet schema"))?
                .as_any()
                .downcast_ref::<Float64Array>()
                .ok_or_else(|| anyhow::anyhow!("Invalid 'low' column type. Expected Float64"))?;

            let closes = batch
                .column_by_name("close")
                .ok_or_else(|| anyhow::anyhow!("Missing 'close' column in Parquet schema"))?
                .as_any()
                .downcast_ref::<Float64Array>()
                .ok_or_else(|| anyhow::anyhow!("Invalid 'close' column type. Expected Float64"))?;

            let volumes = batch
                .column_by_name("volume")
                .ok_or_else(|| anyhow::anyhow!("Missing 'volume' column in Parquet schema"))?
                .as_any()
                .downcast_ref::<UInt64Array>()
                .ok_or_else(|| anyhow::anyhow!("Invalid 'volume' column type. Expected UInt64"))?;

            // Convert to OHLCVBar structs
            for i in 0..batch.num_rows() {
                let timestamp_ns = timestamps.value(i);
                let timestamp = chrono::DateTime::from_timestamp_nanos(timestamp_ns);

                let bar = OHLCVBar {
                    timestamp,
                    open: opens.value(i),
                    high: highs.value(i),
                    low: lows.value(i),
                    close: closes.value(i),
                    volume: volumes.value(i) as f64, // Convert u64 to f64
                };
                all_ohlcv_bars.push(bar);
            }
        }

        info!(
            "Successfully loaded {} OHLCV bars from Parquet file",
            all_ohlcv_bars.len()
        );

        // Sort bars by timestamp (critical for rolling window feature extraction)
        debug!("Sorting bars chronologically by timestamp...");
        all_ohlcv_bars.sort_by_key(|bar| bar.timestamp);
        debug!("Bars sorted successfully");

        // Wave 14 Agent 32: Preprocess close prices for stationarity
        let preprocessed_closes = if self.hyperparams.enable_preprocessing {
            info!("🔬 Preprocessing enabled: Applying log returns + windowed normalization + outlier clipping");

            // Extract close prices (CPU-side preprocessing)
            let close_prices_f64: Vec<f64> = all_ohlcv_bars.iter().map(|b| b.close).collect();
            let close_prices_f32: Vec<f32> = close_prices_f64.iter().map(|&x| x as f32).collect();

            // Configure preprocessing
            let preprocess_config = PreprocessConfig {
                window_size: self.hyperparams.preprocessing_window,
                clip_sigma: self.hyperparams.preprocessing_clip_sigma,
                use_log_returns: true,
            };

            info!("  Window size: {}", preprocess_config.window_size);
            info!("  Clip sigma: +/-{:.1}sigma", preprocess_config.clip_sigma);

            // Compute returns on CPU (preprocessing functions work on &[f32])
            let returns = if preprocess_config.use_log_returns {
                compute_log_returns(&close_prices_f32)
                    .context("Failed to compute log returns")?
            } else {
                // Simple returns: (curr - prev) / prev, with leading zero
                let n = close_prices_f32.len();
                let mut simple_ret = vec![0.0_f32];
                for i in 1..n {
                    let prev = close_prices_f32.get(i - 1).copied().unwrap_or(1.0);
                    let curr = close_prices_f32.get(i).copied().unwrap_or(prev);
                    simple_ret.push(if prev.abs() > 1e-10 { (curr - prev) / prev } else { 0.0 });
                }
                simple_ret
            };

            let normalized = windowed_normalize(&returns, preprocess_config.window_size)
                .context("Failed windowed normalization")?;

            // Compute clip bounds from training portion only (80/20 split)
            let n_norm = normalized.len();
            let train_end = (n_norm * 80) / 100;
            let train_normalized = &normalized[..train_end];
            let clip_bounds = compute_clip_bounds(train_normalized, preprocess_config.clip_sigma)
                .context("Failed to compute clip bounds from training data")?;

            info!("  Clip bounds (train-only): [{:.4}, {:.4}]", clip_bounds.0, clip_bounds.1);

            // Apply training-derived clip bounds to the full normalized series
            let preprocessed_vec = clip_outliers_with_bounds(
                &normalized, clip_bounds.0, clip_bounds.1,
            ).context("Failed to clip outliers with training bounds")?;

            // Compute validation stats on CPU
            let warmup = preprocess_config.window_size as usize;
            let n_total = preprocessed_vec.len();
            let post_warmup = &preprocessed_vec[warmup.min(n_total)..];
            let n_pw = post_warmup.len() as f64;
            let mean = if n_pw > 0.0 {
                post_warmup.iter().map(|&v| v as f64).sum::<f64>() / n_pw
            } else { 0.0 };
            let variance = if n_pw > 1.0 {
                post_warmup.iter().map(|&v| { let d = v as f64 - mean; d * d }).sum::<f64>() / n_pw
            } else { 0.0 };
            let std = variance.sqrt();
            let max_abs = post_warmup.iter().map(|&v| (v as f64).abs()).fold(0.0_f64, f64::max);

            let preprocessed_f64: Vec<f64> = preprocessed_vec.iter().map(|&x| x as f64).collect();

            debug!("✅ Preprocessing complete:");
            debug!("  • Mean: {:.6} (expected ~0 for normalized data)", mean);
            debug!("  • Std: {:.4} (expected ~1 for normalized data)", std);
            debug!(
                "  • Max absolute value: {:.4} (clipped at ±{:.1}σ)",
                max_abs, preprocess_config.clip_sigma
            );

            Some(preprocessed_f64)
        } else {
            info!("⚠️  Preprocessing disabled: Using raw close prices (NON-STATIONARY)");
            None
        };

        // WAVE 2-A2: Load MBP-10 snapshots for OFI calculation (parallel)
        use data::providers::databento::dbn_parser::DbnParser;
        let mbp10_dir_path = self.hyperparams.mbp10_data_dir.as_ref().map(Path::new);
        let mbp10_snapshots = if let Some(mbp10_dir) = mbp10_dir_path {
            if mbp10_dir.exists() {
                let mut dbn_files: Vec<_> = collect_dbn_files_recursive(mbp10_dir);
                dbn_files.sort();

                if !dbn_files.is_empty() {
                    let start_mbp10 = std::time::Instant::now();
                    let n_files = dbn_files.len();
                    info!("📊 Loading MBP-10 snapshots from {} files (parallel)...", n_files);

                    let per_file: Vec<Vec<_>> = {
                        use rayon::prelude::*;
                        dbn_files.par_iter().filter_map(|file| {
                            let parser = match DbnParser::new() {
                                Ok(p) => p,
                                Err(e) => { warn!("  DBN 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) => {
                                    warn!("  ⚠️ Failed to load {:?}: {e}", file.file_name());
                                    None
                                }
                            }
                        }).collect()
                    };

                    let total_snaps: usize = per_file.iter().map(|v| v.len()).sum();
                    let mut all_snapshots = Vec::with_capacity(total_snaps);
                    for snaps in per_file {
                        all_snapshots.extend(snaps);
                    }
                    all_snapshots.sort_by_key(|s| s.timestamp);
                    let mbp10_secs = start_mbp10.elapsed().as_secs_f64();
                    info!("✅ Loaded {} MBP-10 snapshots from {} files in {:.1}s", total_snaps, n_files, mbp10_secs);

                    if all_snapshots.is_empty() { None } else { Some(all_snapshots) }
                } else {
                    info!("No MBP-10 .dbn files found. OFI features will be zeros.");
                    None
                }
            } else {
                info!("MBP-10 directory not found at {:?}. OFI features disabled.", mbp10_dir);
                None
            }
        } else {
            info!("ℹ️  No MBP-10 data directory configured. OFI features disabled.");
            None
        };

        // Extract 40-feature vectors (technical indicators, time, statistical features)
        info!("Extracting 42-feature vectors from OHLCV bars...");
        let feature_vectors = self.extract_full_features(&all_ohlcv_bars, mbp10_snapshots.as_deref())?;

        info!(
            "Extracted {} feature vectors (42 dimensions: technical indicators, time, statistical, regime features)",
            feature_vectors.len()
        );

        // Load trade data for real VPIN/Kyle's Lambda (parallel)
        let trades = if let Some(ref trades_dir_str) = self.hyperparams.trades_data_dir {
            use crate::features::trades_loader::load_trades_sync;
            let trades_dir = Path::new(trades_dir_str);
            if trades_dir.exists() {
                let mut trade_files: Vec<_> = collect_dbn_files_recursive(trades_dir);
                trade_files.sort();
                if !trade_files.is_empty() {
                    info!("📊 Loading trade data from {} files (parallel)...", trade_files.len());
                    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) => {
                                    warn!("  ⚠️ Failed to load {:?}: {e}", path.file_name());
                                    None
                                }
                            }
                        }).collect()
                    };
                    let mut all_trades = Vec::new();
                    for t in per_file_trades {
                        all_trades.extend(t);
                    }
                    if all_trades.is_empty() {
                        info!("No trades loaded. VPIN/Kyle's Lambda will use tick-rule proxy.");
                        None
                    } else {
                        all_trades.sort_by_key(|t| t.timestamp);
                        info!("✅ Total trades loaded: {} (sorted by timestamp)", all_trades.len());
                        Some(all_trades)
                    }
                } else {
                    None
                }
            } else {
                info!("Trades directory not found at {:?}. Using tick-rule proxy.", trades_dir);
                None
            }
        } else {
            None
        };

        // Compute OFI features from MBP-10 snapshots (8 features per bar)
        if let Some(ref snapshots) = mbp10_snapshots {
            use crate::features::ofi_calculator::OFICalculator;
            use crate::features::mbp10_loader::get_snapshots_for_timestamp;
            use crate::features::trades_loader::get_trades_for_bar;

            info!("📊 Computing OFI features from {} MBP-10 snapshots...", snapshots.len());
            let mut ofi_calculator = OFICalculator::new();
            let mut ofi_per_bar = Vec::with_capacity(feature_vectors.len());

            const WARMUP: usize = 50;
            for i in 0..feature_vectors.len() {
                let bar = &all_ohlcv_bars[i + WARMUP];
                let bar_ts = bar.timestamp.timestamp_nanos_opt().unwrap_or(0) as u64;

                // Feed real trade data for this bar's time window (VPIN + Kyle's Lambda)
                if let Some(ref all_trades) = trades {
                    // Bar duration: use next bar's timestamp or add 60s (1-min bars)
                    let bar_end_ts = all_ohlcv_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); // 60s fallback
                    let bar_trades = get_trades_for_bar(all_trades, bar_ts, bar_end_ts);
                    for trade in bar_trades {
                        ofi_calculator.feed_trade(trade.price, trade.volume, trade.is_buy);
                    }
                }

                // Find the nearest MBP-10 snapshot for this bar's timestamp
                let window = get_snapshots_for_timestamp(snapshots, bar_ts, 1);
                if let Some(snap) = window.first() {
                    match ofi_calculator.calculate(snap) {
                        Ok(features) => {
                            let arr = features.to_array();
                            if features.is_valid() {
                                ofi_per_bar.push(arr);
                            } else {
                                ofi_per_bar.push([0.0; 8]);
                            }
                        }
                        Err(_) => 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();
            let has_trades = trades.is_some();
            info!("✅ OFI features computed: {} total, {} non-zero ({:.1}%), trade_data={}",
                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 },
                has_trades);
            self.ofi_features = Some(Arc::from(ofi_per_bar));
        }

        // Create training data pairs (features, target)
        // Target: [preprocessed_current, preprocessed_next, raw_current, raw_next]
        // WAVE 3 BUG FIX: Include raw prices for triple barrier tracker (needs actual market prices in cents)
        // Wave 14 Agent 32: Use preprocessed closes if enabled
        let mut training_data = Vec::new();
        for i in 0..feature_vectors.len().saturating_sub(1) {
            let (preprocessed_current, preprocessed_next, raw_current, raw_next) = if let Some(ref preprocessed) = preprocessed_closes {
                // Use preprocessed values for reward calculation + raw for barrier tracker
                (
                    preprocessed[i + 50],
                    preprocessed[i + 1 + 50],
                    all_ohlcv_bars[i + 50].close,
                    all_ohlcv_bars[i + 1 + 50].close,
                )
            } else {
                // Use raw prices for both (original behavior)
                let raw_curr = all_ohlcv_bars[i + 50].close;
                let raw_next = all_ohlcv_bars[i + 1 + 50].close;
                (raw_curr, raw_next, raw_curr, raw_next)
            };
            training_data.push((feature_vectors[i], vec![preprocessed_current, preprocessed_next, raw_current, raw_next]));
        }
        // Last sample targets itself
        if !feature_vectors.is_empty() {
            let idx = all_ohlcv_bars.len() - 1;
            let (preprocessed_close, raw_close) = if let Some(ref preprocessed) = preprocessed_closes {
                (preprocessed[idx], all_ohlcv_bars[idx].close)
            } else {
                let raw = all_ohlcv_bars[idx].close;
                (raw, raw)
            };
            training_data.push((
                feature_vectors[feature_vectors.len() - 1],
                vec![preprocessed_close, preprocessed_close, raw_close, raw_close],
            ));
        }

        info!(
            "Created {} total samples with 42-dim features",
            training_data.len()
        );

        // Split training data 80/20 for train/validation
        let split_idx = (training_data.len() * 80) / 100;
        let train_data = training_data[..split_idx].to_vec(); // cpu-side split
        let val_data = training_data[split_idx..].to_vec(); // cpu-side split

        info!(
            "Split data - Training samples: {}, Validation samples: {}",
            train_data.len(),
            val_data.len()
        );

        Ok((train_data, val_data))
    }

    /// Load training data from DBN files using official dbn crate decoder
    /// Returns (train_data, val_data) with 80/20 split
    pub(crate) async fn load_training_data(
        &mut self,
        dbn_data_dir: &str,
    ) -> Result<(
        Vec<(FeatureVector, Vec<f64>)>,
        Vec<(FeatureVector, Vec<f64>)>,
    )> {
        // ── DBN feature cache (dev/test speedup) ──────────────────────────────
        // Check for a previously computed and serialised training dataset.
        // The cache is keyed on all .dbn/.dbn.zst filenames + sizes + mtimes so
        // any data change automatically invalidates it.
        // Disabled by FOXHUNT_FEATURE_CACHE=0.
        let dir_path_for_cache = Path::new(dbn_data_dir);
        let mbp10_owned = self.hyperparams.mbp10_data_dir.clone();
        let trades_owned = self.hyperparams.trades_data_dir.clone();
        let mbp10_path = mbp10_owned.as_ref().map(|s| Path::new(s.as_str()));
        let trades_path = trades_owned.as_ref().map(|s| Path::new(s.as_str()));
        if dir_path_for_cache.exists() {
            match crate::feature_cache::load_dbn_training_cache_full(
                dir_path_for_cache,
                mbp10_path,
                trades_path,
            ) {
                Ok(Some((train, val))) => {
                    info!(
                        "DBN feature cache hit: {} train + {} val samples (skipped ~60s load)",
                        train.len(),
                        val.len()
                    );
                    return Ok((train, val));
                }
                Ok(None) => {
                    debug!("DBN feature cache miss — loading from disk");
                }
                Err(e) => {
                    debug!("DBN feature cache unavailable (non-fatal): {e}");
                }
            }
        }

        // Find all DBN files in directory
        let dir_path = Path::new(dbn_data_dir);
        if !dir_path.exists() {
            return Err(anyhow::anyhow!(
                "Data directory not found: {}",
                dbn_data_dir
            ));
        }

        let mut dbn_files = collect_dbn_files_recursive(dir_path);
        dbn_files.sort();

        if dbn_files.is_empty() {
            return Err(anyhow::anyhow!("No DBN files found in: {}", dbn_data_dir));
        }

        info!("Found {} DBN files to load", dbn_files.len());

        let mut all_ohlcv_bars = Vec::new();

        // Load and decode each DBN file to collect OHLCV bars
        for (file_idx, file_path) in dbn_files.iter().enumerate() {
            debug!(
                "Loading DBN file {}/{}: {}",
                file_idx + 1,
                dbn_files.len(),
                file_path.display()
            );

            // Extract raw OHLCV bars from file
            let file_bars = self.extract_ohlcv_bars_from_dbn(file_path)?;

            debug!(
                "Extracted {} OHLCV bars from {}",
                file_bars.len(),
                file_path.file_name().unwrap_or_default().to_string_lossy()
            );

            all_ohlcv_bars.extend(file_bars);
        }

        if all_ohlcv_bars.is_empty() {
            return Err(anyhow::anyhow!(
                "No OHLCV bars extracted from DBN files. Check if files contain OHLCV messages."
            ));
        }

        info!(
            "Successfully loaded {} OHLCV bars from {} DBN files",
            all_ohlcv_bars.len(),
            dbn_files.len()
        );

        // Sort bars by timestamp (critical for rolling window feature extraction)
        debug!("Sorting bars chronologically by timestamp...");
        all_ohlcv_bars.sort_by_key(|bar| bar.timestamp);
        debug!("Bars sorted successfully");

        // Extract 40-feature vectors (technical indicators, time, statistical features)
        info!("Extracting 42-feature vectors from OHLCV bars...");
        let feature_vectors = self.extract_full_features(&all_ohlcv_bars, None)?;

        info!(
            "Extracted {} feature vectors (42 dimensions: technical indicators, time, statistical, regime features)",
            feature_vectors.len()
        );

        // Compute per-bar OFI features from MBP-10 snapshots (8 features per bar)
        // Uses rayon par_iter to load files in parallel (~7-9x speedup over sequential)
        if let Some(ref mbp10_dir_str) = self.hyperparams.mbp10_data_dir {
            use crate::features::ofi_calculator::OFICalculator;
            use crate::features::mbp10_loader::get_snapshots_for_timestamp;
            use crate::features::trades_loader::get_trades_for_bar;
            use data::providers::databento::dbn_parser::DbnParser;

            let mbp10_dir = Path::new(mbp10_dir_str);
            if mbp10_dir.exists() {
                let mut dbn_files: Vec<_> = collect_dbn_files_recursive(mbp10_dir);
                dbn_files.sort();

                if !dbn_files.is_empty() {
                    let start_mbp10 = std::time::Instant::now();
                    let n_files = dbn_files.len();
                    info!("Loading MBP-10 snapshots from {n_files} files (parallel)...");

                    // Parse all MBP-10 files in parallel with rayon.
                    // Each file produces a sorted Vec of snapshots; we merge after.
                    let per_file: Vec<Vec<_>> = {
                        use rayon::prelude::*;
                        dbn_files.par_iter().filter_map(|file| {
                            let parser = match DbnParser::new() {
                                Ok(p) => p,
                                Err(e) => { warn!("  DBN 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) => {
                                    warn!("  Failed to load {:?}: {e}", file.file_name());
                                    None
                                }
                            }
                        }).collect()
                    };

                    let total_snaps: usize = per_file.iter().map(|v| v.len()).sum();
                    let mut all_snapshots = Vec::with_capacity(total_snaps);
                    for snaps in per_file {
                        all_snapshots.extend(snaps);
                    }
                    all_snapshots.sort_by_key(|s| s.timestamp);
                    let mbp10_secs = start_mbp10.elapsed().as_secs_f64();
                    info!("Loaded {total_snaps} MBP-10 snapshots from {n_files} files in {mbp10_secs:.1}s");

                    if !all_snapshots.is_empty() {
                        // Load trade data in parallel for VPIN/Kyle's Lambda enrichment
                        let trades = if let Some(ref trades_dir_str) = self.hyperparams.trades_data_dir {
                            use crate::features::trades_loader::load_trades_sync;
                            let trades_dir = Path::new(trades_dir_str);
                            if trades_dir.exists() {
                                let mut trade_files: Vec<_> = collect_dbn_files_recursive(trades_dir);
                                trade_files.sort();
                                let per_file_trades: Vec<Vec<_>> = {
                                    use rayon::prelude::*;
                                    trade_files.par_iter().filter_map(|path| {
                                        load_trades_sync(path).ok()
                                    }).collect()
                                };
                                let mut all_trades = Vec::new();
                                for t in per_file_trades {
                                    all_trades.extend(t);
                                }
                                (!all_trades.is_empty()).then(|| {
                                    all_trades.sort_by_key(|t| t.timestamp);
                                    all_trades
                                })
                            } else {
                                None
                            }
                        } else {
                            None
                        };

                        info!("Computing per-bar OFI from {} MBP-10 snapshots...", all_snapshots.len());
                        let mut ofi_calculator = OFICalculator::new();
                        let mut ofi_per_bar = Vec::with_capacity(feature_vectors.len());
                        const WARMUP: usize = 50;

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

                            if let Some(ref all_trades) = trades {
                                let bar_end_ts = all_ohlcv_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);
                                let bar_trades = get_trades_for_bar(all_trades, bar_ts, bar_end_ts);
                                for trade in bar_trades {
                                    ofi_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 ofi_calculator.calculate(snap) {
                                    Ok(features) => {
                                        if features.is_valid() {
                                            ofi_per_bar.push(features.to_array());
                                        } else {
                                            ofi_per_bar.push([0.0; 8]);
                                        }
                                    }
                                    Err(_) => 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!("Per-bar OFI computed: {} total, {} non-zero ({:.1}%)",
                            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 });
                        self.ofi_features = Some(Arc::from(ofi_per_bar));
                    }
                }
            }
        }

        // Create training data pairs (features, target)
        // Target layout: [preproc_close, preproc_next, raw_close, raw_next]
        //   [0:1] = network input (raw prices here; preprocessing applied in feature vector)
        //   [2:3] = raw dollar prices for portfolio simulation P&L + tx costs
        let mut training_data = Vec::new();
        for i in 0..feature_vectors.len().saturating_sub(1) {
            let raw_current = all_ohlcv_bars[i + 50].close;
            let raw_next = all_ohlcv_bars[i + 1 + 50].close;
            training_data.push((
                feature_vectors[i],
                vec![raw_current, raw_next, raw_current, raw_next],
            ));
        }
        // Last sample targets itself
        if !feature_vectors.is_empty() {
            let idx = all_ohlcv_bars.len() - 1;
            let raw_close = all_ohlcv_bars[idx].close;
            training_data.push((
                feature_vectors[feature_vectors.len() - 1],
                vec![raw_close, raw_close, raw_close, raw_close],
            ));
        }

        info!(
            "Created {} total samples with 42-dim features",
            training_data.len()
        );

        // Split training data 80/20 for train/validation
        let split_idx = (training_data.len() * 80) / 100;
        let train_data = training_data[..split_idx].to_vec(); // cpu-side split
        let val_data = training_data[split_idx..].to_vec(); // cpu-side split

        info!(
            "Split data - Training samples: {}, Validation samples: {}",
            train_data.len(),
            val_data.len()
        );

        // ── Persist to cache for future runs ─────────────────────────────────
        crate::feature_cache::save_dbn_training_cache_full(
            Path::new(dbn_data_dir),
            mbp10_path,
            trades_path,
            &train_data,
            &val_data,
        );

        Ok((train_data, val_data))
    }

    /// Extract raw OHLCV bars from DBN file using official dbn crate decoder
    ///
    /// This replaces the custom parser that only extracted 2 messages (header metadata).
    /// Now extracts all OHLCV bars (400-500+ records per file).
    ///
    /// Public for testing purposes.
    pub fn extract_ohlcv_bars_from_dbn(&self, file_path: &Path) -> Result<Vec<OHLCVBar>> {
        // Detect zstd by magic bytes (0x28B52FFD) — file extension is unreliable
        if is_zstd_file(file_path)? {
            let mut decoder = dbn::decode::DbnDecoder::from_zstd_file(file_path)
                .map_err(|e| anyhow::anyhow!("Failed to create zstd DBN decoder for {}: {}", file_path.display(), e))?;
            Self::decode_ohlcv_records(&mut decoder, file_path)
        } else {
            let mut decoder = dbn::decode::DbnDecoder::from_file(file_path)
                .map_err(|e| anyhow::anyhow!("Failed to create DBN decoder for {}: {}", file_path.display(), e))?;
            Self::decode_ohlcv_records(&mut decoder, file_path)
        }
    }

    fn decode_ohlcv_records<R: std::io::Read>(
        decoder: &mut dbn::decode::dbn::Decoder<R>,
        file_path: &Path,
    ) -> Result<Vec<OHLCVBar>> {
        use dbn::decode::{DbnMetadata, DecodeRecordRef};

        let mut ohlcv_bars = Vec::new();

        // Read metadata — reject non-OHLCV schemas at file level
        let metadata = decoder.metadata();
        debug!(
            "DBN file metadata: dataset={:?}, schema={:?}, symbols={:?}",
            metadata.dataset, metadata.schema, metadata.symbols
        );

        // Reject files with non-OHLCV schema upfront (MBP-10, trades, etc.)
        if let Some(schema) = metadata.schema {
            let schema_str = format!("{schema:?}");
            if !schema_str.contains("Ohlcv") {
                tracing::warn!(
                    schema = %schema_str,
                    file = %file_path.display(),
                    "Skipping non-OHLCV DBN file"
                );
                return Ok(Vec::new());
            }
        }

        // Decode all OHLCV records
        let mut ohlcv_count = 0;
        let mut other_count = 0;
        let mut idx = 0;

        loop {
            match decoder.decode_record_ref() {
                Ok(Some(record)) => {
                    idx += 1;

                    // Convert RecordRef to RecordRefEnum for pattern matching
                    let record_enum = record
                        .as_enum()
                        .map_err(|e| anyhow::anyhow!(
                            "Corrupt record at idx {} in {}: {}",
                            idx, file_path.display(), e
                        ))?;

                    match record_enum {
                        dbn::RecordRefEnum::Ohlcv(ohlcv) => {
                            ohlcv_count += 1;

                            // Extract OHLCV values (prices are i64 scaled by 1e-9 per DBN spec, volume is u64)
                            let open_f64 = ohlcv.open as f64 * 1e-9;
                            let high_f64 = ohlcv.high as f64 * 1e-9;
                            let low_f64 = ohlcv.low as f64 * 1e-9;
                            let close_f64 = ohlcv.close as f64 * 1e-9;
                            let volume_u64 = ohlcv.volume;

                            // WAVE 8 AGENT 36: Validate all price values are finite (not NaN/Inf)
                            // Skip bars with invalid data to prevent NaN propagation
                            if !open_f64.is_finite()
                                || !high_f64.is_finite()
                                || !low_f64.is_finite()
                                || !close_f64.is_finite()
                            {
                                debug!(
                                    "Skipping OHLCV bar {} with non-finite values: open={}, high={}, low={}, close={}",
                                    ohlcv_count, open_f64, high_f64, low_f64, close_f64
                                );
                                continue;
                            }

                            // Log first few records for validation
                            if ohlcv_count <= 5 {
                                debug!(
                                    "Raw OHLCV #{}: open={}, high={}, low={}, close={}",
                                    ohlcv_count, ohlcv.open, ohlcv.high, ohlcv.low, ohlcv.close
                                );
                                debug!(
                                    "Scaled OHLCV #{}: open={:.6}, high={:.6}, low={:.6}, close={:.6}",
                                    ohlcv_count, open_f64, high_f64, low_f64, close_f64
                                );
                            }

                            // Convert timestamp from nanoseconds since epoch to DateTime
                            let timestamp_nanos = ohlcv.hd.ts_event as i64;
                            let timestamp_secs = timestamp_nanos / 1_000_000_000;
                            let timestamp_nanos_remainder =
                                (timestamp_nanos.rem_euclid(1_000_000_000)) as u32;
                            let timestamp = chrono::DateTime::<chrono::Utc>::from_timestamp(
                                timestamp_secs,
                                timestamp_nanos_remainder,
                            )
                            .unwrap_or_else(|| chrono::Utc::now());

                            // Create OHLCVBar for feature extraction pipeline
                            let bar = OHLCVBar {
                                timestamp,
                                open: open_f64,
                                high: high_f64,
                                low: low_f64,
                                close: close_f64,
                                volume: volume_u64 as f64,
                            };

                            ohlcv_bars.push(bar);
                        },
                        dbn::RecordRefEnum::Mbo(_)
                        | dbn::RecordRefEnum::Trade(_)
                        | dbn::RecordRefEnum::Mbp1(_)
                        | dbn::RecordRefEnum::Mbp10(_)
                        | dbn::RecordRefEnum::Status(_)
                        | dbn::RecordRefEnum::InstrumentDef(_)
                        | dbn::RecordRefEnum::Imbalance(_)
                        | dbn::RecordRefEnum::Stat(_)
                        | dbn::RecordRefEnum::Error(_)
                        | dbn::RecordRefEnum::SymbolMapping(_)
                        | dbn::RecordRefEnum::System(_)
                        | dbn::RecordRefEnum::Cmbp1(_)
                        | dbn::RecordRefEnum::Bbo(_)
                        | dbn::RecordRefEnum::Cbbo(_) => {
                            other_count += 1;
                            if other_count <= 5 {
                                debug!("Skipping non-OHLCV record at index {}", idx);
                            }
                        },
                    }
                },
                Ok(None) => {
                    // End of stream
                    break;
                },
                Err(e) => {
                    return Err(anyhow::anyhow!(
                        "Failed to decode record {} in {}: {}",
                        idx, file_path.display(), e
                    ));
                },
            }
        }

        info!(
            "Extracted {} OHLCV bars from {:?} ({} other records skipped)",
            ohlcv_count,
            file_path.file_name().unwrap_or_default(),
            other_count
        );

        Ok(ohlcv_bars)
    }

    /// Create features from OHLCV data
    pub(crate) fn create_ohlcv_features(
        &self,
        open: f64,
        high: f64,
        low: f64,
        close: f64,
        volume: u64,
    ) -> Result<FinancialFeatures> {
        use std::collections::HashMap;

        // Use absolute values for Price type (futures data can have negative values)
        // For ML training, the absolute magnitude is what matters for feature extraction
        let close_price =
            common::Price::from_f64(close.abs()).unwrap_or_else(|_| common::Price::ZERO);
        let open_price =
            common::Price::from_f64(open.abs()).unwrap_or_else(|_| common::Price::ZERO);
        let high_price =
            common::Price::from_f64(high.abs()).unwrap_or_else(|_| common::Price::ZERO);
        let low_price = common::Price::from_f64(low.abs()).unwrap_or_else(|_| common::Price::ZERO);

        // Calculate technical indicators
        let mut indicators = HashMap::new();

        // Price-based features
        let price_range = high - low;
        let body_size = (close - open).abs();
        let upper_shadow = high - close.max(open);
        let lower_shadow = close.min(open) - low;

        indicators.insert("price_range".to_owned(), price_range);
        indicators.insert("body_size".to_owned(), body_size);
        indicators.insert("upper_shadow".to_owned(), upper_shadow);
        indicators.insert("lower_shadow".to_owned(), lower_shadow);
        indicators.insert("close_to_high".to_owned(), (close - high).abs());
        indicators.insert("close_to_low".to_owned(), (close - low).abs());

        // Microstructure features
        let spread_bps = ((high - low) / close * 10000.0) as i32;
        let trade_intensity = volume as f64;

        Ok(FinancialFeatures {
            prices: vec![open_price, high_price, low_price, close_price],
            volumes: vec![volume as i64],
            technical_indicators: indicators,
            microstructure: crate::training_pipeline::MicrostructureFeatures {
                spread_bps,
                imbalance: 0.0, // Not available from OHLCV
                trade_intensity,
                vwap: close_price, // Approximate VWAP as close
            },
            risk_metrics: crate::training_pipeline::RiskFeatures {
                var_5pct: -0.02, // Placeholder
                expected_shortfall: -0.03,
                max_drawdown: -0.05,
                sharpe_ratio: 1.0,
            },
            timestamp: chrono::Utc::now(),
        })
    }
}
