foxhunt/crates/ml/tests/smoke_test_real_data.rs

//! Smoke tests with real Databento market data.
//!
//! These tests validate the full pipeline from DBN parsing through GPU kernel
//! forward pass using actual ES.FUT data downloaded from MinIO.
//!
//! Data requirements (in `data/smoke-test/` relative to workspace root):
//!   - ohlcv/ES.FUT/ES.FUT_2024-Q1.dbn.zst   (~700 KB)
//!   - trades/ES.FUT/ES.FUT_2024-Q1.dbn.zst   (~120 MB)
//!   - trades/ES.FUT/ES.FUT_2025-Q2.dbn.zst   (~505 MB, matches MBP-10 quarter)
//!   - mbp10/ES.FUT/ES.FUT_2025-Q2.dbn.zst    (~125 MB)
//!
//! Run: `cargo test -p ml --test smoke_test_real_data -- --ignored --nocapture`

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

/// Resolve workspace root from CARGO_MANIFEST_DIR (crates/ml → ../..)
fn workspace_root() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .parent() // crates/
        .and_then(|p| p.parent()) // workspace root
        .expect("Failed to find workspace root")
        .to_path_buf()
}

fn smoke_data_dir() -> PathBuf {
    workspace_root().join("data").join("smoke-test")
}

/// Skip test if required data files are missing.
fn require_data(relative: &str) -> PathBuf {
    let path = smoke_data_dir().join(relative);
    if !path.exists() {
        eprintln!(
            "Smoke test data not found: {}. Download from MinIO first.",
            path.display()
        );
        std::process::exit(0);
    }
    path
}

// ==========================================================================
// Test 1: Parse OHLCV data and extract 42-dim features
// ==========================================================================

#[test]
#[ignore]
fn smoke_ohlcv_parse_and_extract_features() {
    let ohlcv_dir = require_data("ohlcv");

    // Load OHLCV bars from .dbn.zst
    let bars = ml::hyperopt::adapters::dbn_loader::load_bars_from_dbn_dir(&ohlcv_dir)
        .expect("Failed to load OHLCV bars from DBN");

    println!("Loaded {} OHLCV bars from {}", bars.len(), ohlcv_dir.display());
    assert!(bars.len() > 100, "Expected >100 bars, got {}", bars.len());

    // Verify bar structure
    let first = &bars[0];
    assert!(first.open > 0.0, "Open price should be positive");
    assert!(first.high >= first.low, "High >= Low invariant");
    assert!(first.volume >= 0.0, "Volume should be non-negative");

    // Extract 42-dim features (requires 50-bar warmup)
    let features = ml::features::extract_ml_features(&bars)
        .expect("Feature extraction failed");

    println!("Extracted {} feature vectors (42-dim each)", features.len());
    assert!(
        features.len() >= bars.len() - 50,
        "Expected at least {} features (bars - warmup), got {}",
        bars.len().saturating_sub(50),
        features.len()
    );

    // Verify features are finite and reasonable
    for (i, feat) in features.iter().enumerate().take(10) {
        for (j, &val) in feat.iter().enumerate() {
            assert!(
                val.is_finite(),
                "Feature [{i}][{j}] = {val} is not finite"
            );
        }
    }

    // Spot-check: close price feature (index 3) should match OHLCV bars
    // Features start at bar 50 (warmup period)
    let close_feat = features[0][3]; // Normalized close
    println!(
        "First feature vector close={close_feat:.4}, bar[50].close={:.4}",
        bars[50].close
    );
}

// ==========================================================================
// Test 2: Parse MBP-10 data and compute OFI features
// ==========================================================================

#[test]
#[ignore]
fn smoke_mbp10_ofi_extraction() {
    let mbp10_dir = require_data("mbp10/ES.FUT");

    // Find the .dbn.zst file
    let files: Vec<_> = std::fs::read_dir(&mbp10_dir)
        .expect("Failed to read mbp10 dir")
        .filter_map(|e| e.ok())
        .map(|e| e.path())
        .filter(|p| {
            p.extension()
                .map_or(false, |ext| ext == "zst" || ext == "dbn")
        })
        .collect();

    assert!(!files.is_empty(), "No .dbn.zst files found in {}", mbp10_dir.display());

    let file = &files[0];
    println!("Computing OFI from: {}", file.display());

    let ofi_features = ml::features::mbp10_loader::compute_ofi_from_file(file)
        .expect("OFI computation failed");

    println!("Computed {} OFI feature vectors (8-dim each)", ofi_features.len());
    assert!(
        ofi_features.len() > 1000,
        "Expected >1000 OFI snapshots, got {}",
        ofi_features.len()
    );

    // Verify OFI features are finite
    let mut nan_count = 0;
    let mut zero_count = 0;
    for (i, feat) in ofi_features.iter().enumerate().take(100) {
        for (j, &val) in feat.iter().enumerate() {
            if !val.is_finite() {
                nan_count += 1;
            }
            if val == 0.0 {
                zero_count += 1;
            }
            // OFI values should not be astronomically large
            if val.is_finite() {
                assert!(
                    val.abs() < 1e12,
                    "OFI [{i}][{j}] = {val} unreasonably large"
                );
            }
        }
    }
    println!(
        "OFI quality (first 100): nan={nan_count}, zero={zero_count}/800"
    );
    // Some NaN is expected in early warmup, but not all
    assert!(
        nan_count < 50,
        "Too many NaN OFI features: {nan_count}/800"
    );
}

// ==========================================================================
// Test 2b: OFI with trade enrichment — validates VPIN/Kyle's Lambda are non-zero
// ==========================================================================

#[test]
#[ignore]
fn smoke_ofi_with_trade_enrichment() {
    let mbp10_file = require_data("mbp10/ES.FUT/ES.FUT_2025-Q2.dbn.zst");
    let trades_file = require_data("trades/ES.FUT/ES.FUT_2025-Q2.dbn.zst");

    // Compute OFI WITHOUT trades (baseline — VPIN/Kyle's Lambda will be zero)
    let ofi_without = ml::features::mbp10_loader::compute_ofi_from_file(&mbp10_file)
        .expect("OFI without trades failed");

    // Compute OFI WITH trades (enriched — VPIN/Kyle's Lambda should be non-zero)
    let ofi_with = ml::features::mbp10_loader::compute_ofi_with_trades(&mbp10_file, &trades_file)
        .expect("OFI with trades failed");

    println!(
        "OFI without trades: {} features, with trades: {} features",
        ofi_without.len(),
        ofi_with.len()
    );

    // Both should produce the same number of features (same MBP-10 snapshots)
    assert_eq!(
        ofi_without.len(),
        ofi_with.len(),
        "Feature count should match: without={} vs with={}",
        ofi_without.len(),
        ofi_with.len()
    );

    // OFI feature indices:
    // [0] ofi_level1  [1] ofi_level5  [2] depth_imbalance  [3] vpin
    // [4] kyle_lambda [5] bid_slope   [6] ask_slope         [7] trade_imbalance
    //
    // Without trades: indices 3, 4 should be ~all zero, 7 uses tick-rule fallback
    // With trades: indices 3, 4 should have real values from trade data

    // Count non-zero VPIN (index 3) in both
    let skip = 100; // skip warmup
    let sample_end = ofi_without.len().min(skip + 5000);
    let sample = skip..sample_end;

    let vpin_nonzero_without = ofi_without[sample.clone()]
        .iter()
        .filter(|f| f[3].abs() > 1e-10)
        .count();
    let vpin_nonzero_with = ofi_with[sample.clone()]
        .iter()
        .filter(|f| f[3].abs() > 1e-10)
        .count();

    let kyle_nonzero_without = ofi_without[sample.clone()]
        .iter()
        .filter(|f| f[4].abs() > 1e-10)
        .count();
    let kyle_nonzero_with = ofi_with[sample.clone()]
        .iter()
        .filter(|f| f[4].abs() > 1e-10)
        .count();

    let sample_size = sample.len();
    println!(
        "VPIN non-zero: without={vpin_nonzero_without}/{sample_size}, with={vpin_nonzero_with}/{sample_size}"
    );
    println!(
        "Kyle's Lambda non-zero: without={kyle_nonzero_without}/{sample_size}, with={kyle_nonzero_with}/{sample_size}"
    );

    // The WITHOUT-trades path should have mostly zero VPIN/Kyle's Lambda
    // (this is the bug we fixed)
    assert!(
        vpin_nonzero_without < sample_size / 10,
        "Without trades, VPIN should be mostly zero but got {vpin_nonzero_without}/{sample_size}"
    );

    // The WITH-trades path should have significantly more non-zero VPIN
    assert!(
        vpin_nonzero_with > sample_size / 2,
        "With trades, VPIN should be mostly non-zero but got {vpin_nonzero_with}/{sample_size}"
    );

    // Kyle's Lambda should also improve with trade data
    assert!(
        kyle_nonzero_with > kyle_nonzero_without,
        "Kyle's Lambda should improve with trades: without={kyle_nonzero_without} vs with={kyle_nonzero_with}"
    );

    // Spot-check some enriched values
    let sample_feat = &ofi_with[skip + 500];
    println!(
        "Sample enriched OFI: ofi_l1={:.4} ofi_l5={:.4} depth_imb={:.4} \
         vpin={:.4} kyle={:.6} bid_slope={:.4} ask_slope={:.4} trade_imb={:.4}",
        sample_feat[0],
        sample_feat[1],
        sample_feat[2],
        sample_feat[3],
        sample_feat[4],
        sample_feat[5],
        sample_feat[6],
        sample_feat[7]
    );
}

// ==========================================================================
// Test 3: Parse trades data
// ==========================================================================

#[test]
#[ignore]
fn smoke_trades_parse() {
    let trades_dir = require_data("trades/ES.FUT");

    let files: Vec<_> = std::fs::read_dir(&trades_dir)
        .expect("Failed to read trades dir")
        .filter_map(|e| e.ok())
        .map(|e| e.path())
        .filter(|p| {
            p.extension()
                .map_or(false, |ext| ext == "zst" || ext == "dbn")
        })
        .collect();

    assert!(!files.is_empty(), "No .dbn.zst files in {}", trades_dir.display());

    let file = &files[0];
    println!("Loading trades from: {}", file.display());

    let trades = ml::features::trades_loader::load_trades_sync(file)
        .expect("Trade loading failed");

    println!("Loaded {} trades", trades.len());
    assert!(
        trades.len() > 10_000,
        "Expected >10K trades for a full quarter, got {}",
        trades.len()
    );

    // Verify trade structure
    let first = &trades[0];
    assert!(first.price > 0.0, "Trade price should be positive");
    assert!(first.volume > 0, "Trade volume should be positive");
    assert!(first.timestamp > 0, "Trade timestamp should be set");
}

// ==========================================================================
// Test 4: Full GPU pipeline — real data → kernel forward pass → Q-values
// ==========================================================================

#[cfg(feature = "cuda")]
mod gpu_smoke {
    use super::*;
    use std::sync::Arc;

    use candle_core::Device;
    use ml::cuda_pipeline::gpu_experience_collector::{
        ExperienceCollectorConfig, GpuExperienceCollector,
    };
    use ml::dqn::dueling::{DuelingConfig, DuelingQNetwork};

    type CudaStream = candle_core::cuda_backend::cudarc::driver::CudaStream;
    type CudaSlice<T> = candle_core::cuda_backend::cudarc::driver::CudaSlice<T>;

    fn require_cuda() -> Device {
        match Device::cuda_if_available(0) {
            Ok(dev) if dev.is_cuda() => dev,
            _ => {
                eprintln!("CUDA not available, skipping GPU smoke test");
                std::process::exit(0);
            }
        }
    }

    fn cuda_stream(device: &Device) -> Arc<CudaStream> {
        match device {
            Device::Cuda(cuda_dev) => cuda_dev.cuda_stream(),
            _ => panic!("Expected CUDA device"),
        }
    }

    /// Build flat GPU buffers from real OHLCV features + OFI.
    ///
    /// Returns (market_buf [N*51], target_buf [N*4], num_bars).
    fn real_market_data(
        ohlcv_dir: &Path,
        mbp10_dir: Option<&Path>,
        device: &Device,
    ) -> (CudaSlice<f32>, CudaSlice<f32>, usize) {
        let stream = match device {
            Device::Cuda(cuda_dev) => cuda_dev.cuda_stream(),
            _ => panic!("Expected CUDA device"),
        };

        // 1. Load OHLCV bars
        let bars = ml::hyperopt::adapters::dbn_loader::load_bars_from_dbn_dir(ohlcv_dir)
            .expect("Failed to load OHLCV bars");
        println!("Loaded {} OHLCV bars", bars.len());

        // 2. Extract 42-dim features
        let features = ml::features::extract_ml_features(&bars)
            .expect("Feature extraction failed");
        let n = features.len();
        println!("Extracted {} feature vectors", n);
        assert!(n > 200, "Need at least 200 bars for GPU test, got {n}");

        // 3. Optionally load OFI features from MBP-10
        let ofi_features: Option<Vec<[f64; 8]>> = mbp10_dir.and_then(|dir| {
            let files: Vec<_> = std::fs::read_dir(dir)
                .ok()?
                .filter_map(|e| e.ok())
                .map(|e| e.path())
                .filter(|p| {
                    p.extension()
                        .map_or(false, |ext| ext == "zst" || ext == "dbn")
                })
                .collect();
            let file = files.first()?;
            println!("Loading OFI from: {}", file.display());
            ml::features::mbp10_loader::compute_ofi_from_file(file).ok()
        });

        let ofi_count = ofi_features.as_ref().map_or(0, |v| v.len());
        println!("OFI features: {ofi_count}");

        // 4. Build flat market buffer [n * MARKET_DIM(51)]
        const MARKET_DIM: usize = 51;
        let mut market_data = vec![0.0_f32; n * MARKET_DIM];
        for (i, feat) in features.iter().enumerate() {
            // Indices 0-39: market features
            for (j, &val) in feat.iter().enumerate() {
                market_data[i * MARKET_DIM + j] = val as f32;
            }
            // Indices 40-42: zeros (dead regime features)
            // Indices 43-50: OFI features
            if let Some(ref ofi) = ofi_features {
                if let Some(ofi_row) = ofi.get(i) {
                    for (j, &val) in ofi_row.iter().enumerate() {
                        if val.is_finite() {
                            market_data[i * MARKET_DIM + 43 + j] = val as f32;
                        }
                    }
                }
            }
        }

        // 5. Build flat target buffer [n * 4]
        // Layout: [current_close, next_close, current_close_raw, next_close_raw]
        // Features are offset by warmup (50 bars), so features[i] ≈ bars[i + 50]
        let warmup = bars.len() - n; // typically 50
        let mut target_data = vec![0.0_f32; n * 4];
        for i in 0..n {
            let bar_idx = i + warmup;
            let close = bars[bar_idx].close as f32;
            let next_close = if bar_idx + 1 < bars.len() {
                bars[bar_idx + 1].close as f32
            } else {
                close
            };
            target_data[i * 4] = close;
            target_data[i * 4 + 1] = next_close;
            target_data[i * 4 + 2] = close;
            target_data[i * 4 + 3] = next_close;
        }

        // 6. Upload to GPU
        let mut market_buf = stream.alloc_zeros::<f32>(n * MARKET_DIM).unwrap();
        stream.memcpy_htod(&market_data, &mut market_buf).unwrap();

        let mut target_buf = stream.alloc_zeros::<f32>(n * 4).unwrap();
        stream.memcpy_htod(&target_data, &mut target_buf).unwrap();

        (market_buf, target_buf, n)
    }

    // ------------------------------------------------------------------
    // Test 4a: Real OHLCV data → GPU kernel → valid Q-values
    // ------------------------------------------------------------------

    #[test]
    #[ignore]
    fn smoke_gpu_real_ohlcv_forward() {
        let ohlcv_dir = require_data("ohlcv");
        let device = require_cuda();
        let stream = cuda_stream(&device);

        let (market_buf, target_buf, num_bars) =
            real_market_data(&ohlcv_dir, None, &device);

        println!("GPU buffer: {num_bars} bars × 51 features");

        // Create dueling networks (state_dim=54: 51 market + 3 portfolio)
        let config = DuelingConfig::new(54, 5, vec![256, 256], 128, 128);
        let online = DuelingQNetwork::new(config.clone(), device.clone()).unwrap();
        let target_net = DuelingQNetwork::new(config, device.clone()).unwrap();

        let mut collector = GpuExperienceCollector::new(
            stream.clone(),
            online.vars(),
            target_net.vars(),
            None,
            100_000.0,
            0.01, 0.05,
            (256, 256, 128, 128),
            (54, 51, 51),
            4, 50,
        ).unwrap();

        // Run 4 episodes of 50 timesteps each
        let n_episodes = 4_usize;
        let timesteps = 50;
        let episode_starts: Vec<i32> = (0..n_episodes)
            .map(|i| (i * (num_bars / n_episodes)) as i32)
            .collect();

        let exp_config = ExperienceCollectorConfig {
            n_episodes: n_episodes as i32,
            timesteps_per_episode: timesteps,
            total_bars: num_bars as i32,
            episode_length: timesteps,
            epsilon: 0.1,
            gamma: 0.99,
            use_noisy_nets: false,
            use_distributional: false,
            ..Default::default()
        };

        let batch = collector
            .collect_experiences(&market_buf, &target_buf, &episode_starts, &exp_config)
            .expect("GPU experience collection failed on real data");

        let expected_total = n_episodes * timesteps as usize;
        assert_eq!(batch.actions.len(), expected_total);
        assert_eq!(batch.rewards.len(), expected_total);
        assert_eq!(batch.target_q_values.len(), expected_total);

        // Q-values should be finite
        for (i, &q) in batch.target_q_values.iter().enumerate() {
            assert!(q.is_finite(), "target_q[{i}] = {q} is not finite");
        }

        // Actions in [0, 4]
        for (i, &a) in batch.actions.iter().enumerate() {
            assert!((0..5).contains(&a), "action[{i}] = {a} out of range");
        }

        // Rewards should reflect real price movements (not all zero)
        let nonzero_rewards = batch.rewards.iter().filter(|&&r| r.abs() > 1e-8).count();
        println!(
            "Real OHLCV pipeline: {expected_total} experiences, \
             nonzero_rewards={nonzero_rewards}/{expected_total}"
        );

        let q_min = batch.target_q_values.iter().cloned().fold(f32::INFINITY, f32::min);
        let q_max = batch.target_q_values.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        let r_min = batch.rewards.iter().cloned().fold(f32::INFINITY, f32::min);
        let r_max = batch.rewards.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        println!("Q range: [{q_min:.4}, {q_max:.4}], reward range: [{r_min:.6}, {r_max:.6}]");
    }

    // ------------------------------------------------------------------
    // Test 4b: Real OHLCV + MBP-10 OFI data → GPU kernel → valid Q-values
    // ------------------------------------------------------------------

    #[test]
    #[ignore]
    fn smoke_gpu_real_ohlcv_plus_ofi_forward() {
        let ohlcv_dir = require_data("ohlcv");
        let mbp10_dir = require_data("mbp10/ES.FUT");
        let device = require_cuda();
        let stream = cuda_stream(&device);

        let (market_buf, target_buf, num_bars) =
            real_market_data(&ohlcv_dir, Some(&mbp10_dir), &device);

        println!("GPU buffer: {num_bars} bars × 51 features (with OFI)");

        let config = DuelingConfig::new(54, 5, vec![256, 256], 128, 128);
        let online = DuelingQNetwork::new(config.clone(), device.clone()).unwrap();
        let target_net = DuelingQNetwork::new(config, device.clone()).unwrap();

        let mut collector = GpuExperienceCollector::new(
            stream.clone(),
            online.vars(),
            target_net.vars(),
            None,
            100_000.0,
            0.01, 0.05,
            (256, 256, 128, 128),
            (54, 51, 51),
            4, 50,
        ).unwrap();

        let n_episodes = 4_usize;
        let timesteps = 50;
        let episode_starts: Vec<i32> = (0..n_episodes)
            .map(|i| (i * (num_bars / n_episodes)) as i32)
            .collect();

        let exp_config = ExperienceCollectorConfig {
            n_episodes: n_episodes as i32,
            timesteps_per_episode: timesteps,
            total_bars: num_bars as i32,
            episode_length: timesteps,
            epsilon: 0.1,
            gamma: 0.99,
            use_noisy_nets: false,
            use_distributional: false,
            ..Default::default()
        };

        let batch = collector
            .collect_experiences(&market_buf, &target_buf, &episode_starts, &exp_config)
            .expect("GPU experience collection with OFI failed");

        let expected_total = n_episodes * timesteps as usize;
        assert_eq!(batch.actions.len(), expected_total);

        for (i, &q) in batch.target_q_values.iter().enumerate() {
            assert!(q.is_finite(), "OFI target_q[{i}] = {q} is not finite");
        }

        let q_min = batch.target_q_values.iter().cloned().fold(f32::INFINITY, f32::min);
        let q_max = batch.target_q_values.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        println!(
            "OHLCV+OFI pipeline: {expected_total} experiences, Q range: [{q_min:.4}, {q_max:.4}]"
        );
    }

    // ------------------------------------------------------------------
    // Test 4c: NoisyNet + C51 distributional on real data
    // ------------------------------------------------------------------

    #[test]
    #[ignore]
    fn smoke_gpu_real_data_noisy_distributional() {
        use ml::dqn::distributional_dueling::{
            DistributionalDuelingConfig, DistributionalDuelingQNetwork,
        };

        let ohlcv_dir = require_data("ohlcv");
        let device = require_cuda();
        let stream = cuda_stream(&device);

        let (market_buf, target_buf, num_bars) =
            real_market_data(&ohlcv_dir, None, &device);

        let config = DistributionalDuelingConfig::new(54, 5, 51, vec![256, 256], 128, 128);
        let online = DistributionalDuelingQNetwork::new(config.clone(), device.clone()).unwrap();
        let target_net = DistributionalDuelingQNetwork::new(config, device.clone()).unwrap();

        let mut collector = GpuExperienceCollector::new(
            stream.clone(),
            online.vars(),
            target_net.vars(),
            None,
            100_000.0,
            0.01, 0.05,
            (256, 256, 128, 128),
            (54, 51, 51),
            4, 50,
        ).unwrap();

        // Sync RMSNorm weights (C51 distributional)
        collector.sync_online_rmsnorm(online.vars()).unwrap();
        collector.sync_target_rmsnorm(target_net.vars()).unwrap();

        let n_episodes = 4_usize;
        let timesteps = 50;
        let episode_starts: Vec<i32> = (0..n_episodes)
            .map(|i| (i * (num_bars / n_episodes)) as i32)
            .collect();

        let exp_config = ExperienceCollectorConfig {
            n_episodes: n_episodes as i32,
            timesteps_per_episode: timesteps,
            total_bars: num_bars as i32,
            episode_length: timesteps,
            epsilon: 0.0, // NoisyNets handle exploration
            gamma: 0.99,
            use_noisy_nets: true,
            noisy_sigma_init: 0.5,
            use_distributional: true,
            num_atoms: 51,
            v_min: -10.0,
            v_max: 10.0,
            ..Default::default()
        };

        let batch = collector
            .collect_experiences(&market_buf, &target_buf, &episode_starts, &exp_config)
            .expect("Distributional+NoisyNet failed on real data");

        let expected_total = n_episodes * timesteps as usize;
        assert_eq!(batch.actions.len(), expected_total);

        for (i, &q) in batch.target_q_values.iter().enumerate() {
            assert!(q.is_finite(), "C51 target_q[{i}] = {q} is not finite");
        }

        // NoisyNet should produce diverse actions (not all the same)
        let mut action_counts = [0_usize; 5];
        for &a in &batch.actions {
            if let Some(count) = action_counts.get_mut(a as usize) {
                *count += 1;
            }
        }
        let unique_actions = action_counts.iter().filter(|&&c| c > 0).count();
        println!(
            "C51+NoisyNet: actions={action_counts:?}, unique={unique_actions}"
        );
        assert!(
            unique_actions >= 2,
            "NoisyNet should produce at least 2 distinct actions, got {unique_actions}"
        );

        let q_min = batch.target_q_values.iter().cloned().fold(f32::INFINITY, f32::min);
        let q_max = batch.target_q_values.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        println!("C51 Q range: [{q_min:.4}, {q_max:.4}]");
    }
}