foxhunt/crates/ml/tests/smoke_test_real_data.rs

#![allow(
    clippy::assertions_on_constants,
    clippy::assertions_on_result_states,
    clippy::clone_on_copy,
    clippy::decimal_literal_representation,
    clippy::doc_markdown,
    clippy::empty_line_after_doc_comments,
    clippy::field_reassign_with_default,
    clippy::get_unwrap,
    clippy::identity_op,
    clippy::inconsistent_digit_grouping,
    clippy::indexing_slicing,
    clippy::integer_division,
    clippy::len_zero,
    clippy::let_underscore_must_use,
    clippy::manual_div_ceil,
    clippy::manual_let_else,
    clippy::manual_range_contains,
    clippy::modulo_arithmetic,
    clippy::needless_range_loop,
    clippy::non_ascii_literal,
    clippy::redundant_clone,
    clippy::shadow_reuse,
    clippy::shadow_same,
    clippy::shadow_unrelated,
    clippy::single_match_else,
    clippy::str_to_string,
    clippy::string_slice,
    clippy::tests_outside_test_module,
    clippy::too_many_lines,
    clippy::unnecessary_wraps,
    clippy::unseparated_literal_suffix,
    clippy::use_debug,
    clippy::useless_vec,
    clippy::wildcard_enum_match_arm,
    clippy::else_if_without_else,
    clippy::expect_used,
    clippy::missing_const_for_fn,
    clippy::similar_names,
    clippy::type_complexity,
    clippy::collapsible_else_if,
    clippy::doc_lazy_continuation,
    clippy::items_after_test_module,
    clippy::map_clone,
    clippy::multiple_unsafe_ops_per_block,
    clippy::unwrap_or_default,
    clippy::assign_op_pattern,
    clippy::needless_borrow,
    clippy::println_empty_string,
    clippy::unnecessary_cast,
    clippy::used_underscore_binding,
    clippy::create_dir,
    clippy::implicit_saturating_sub,
    clippy::exit,
    clippy::expect_fun_call,
    clippy::too_many_arguments,
    clippy::unnecessary_map_or,
    clippy::unwrap_used,
    dead_code,
    unused_imports,
    unused_variables,
    clippy::cloned_ref_to_slice_refs,
    clippy::neg_multiply,
    clippy::while_let_loop,
    clippy::bool_assert_comparison,
    clippy::excessive_precision,
    clippy::trivially_copy_pass_by_ref,
    clippy::op_ref,
    clippy::redundant_closure,
    clippy::unnecessary_lazy_evaluations,
    clippy::if_then_some_else_none,
    clippy::unnecessary_to_owned,
    clippy::single_component_path_imports,
)]
//! 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 -- --nocapture`

use std::path::{Path, PathBuf};
use tracing::{info, warn};

/// 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 {
    // CI: TEST_DATA_DIR points to test-data-pvc on H100
    if let Ok(dir) = std::env::var("TEST_DATA_DIR") {
        return PathBuf::from(dir);
    }
    // Local dev: relative path from workspace root
    workspace_root().join("data").join("smoke-test")
}

/// Check if GPU has enough VRAM (MB). Returns false if no GPU or insufficient VRAM.
fn gpu_vram_mb() -> u64 {
    let output = std::process::Command::new("nvidia-smi")
        .args(["--query-gpu=memory.total", "--format=csv,noheader,nounits"])
        .output();
    match output {
        Ok(out) if out.status.success() => {
            String::from_utf8_lossy(&out.stdout)
                .trim()
                .parse()
                .unwrap_or(0)
        }
        _ => 0,
    }
}

/// Return data path if it exists, `None` otherwise (test should skip).
fn try_data(relative: &str) -> Option<PathBuf> {
    let path = smoke_data_dir().join(relative);
    if !path.exists() {
        warn!(path = %path.display(), "SKIP: smoke test data not found — download from MinIO first");
        return None;
    }
    Some(path)
}

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

#[test]
fn smoke_ohlcv_parse_and_extract_features() {
    let Some(ohlcv_dir) = try_data("ohlcv") else { return };

    // 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");

    info!(count = bars.len(), dir = %ohlcv_dir.display(), "Loaded OHLCV bars");
    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");

    info!(count = features.len(), dim = 42, "Extracted feature vectors");
    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
    info!(close_feat, bar_50_close = bars[50].close, "First feature vector spot-check");
}

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

#[test]
fn smoke_mbp10_ofi_extraction() {
    let Some(mbp10_dir) = try_data("mbp10/ES.FUT") else { return };

    // 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];
    info!(file = %file.display(), "Computing OFI");

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

    info!(count = ofi_features.len(), dim = 8, "Computed OFI feature vectors");
    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"
                );
            }
        }
    }
    info!(nan_count, zero_count, sample = 800, "OFI quality (first 100 features)");
    // 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]
fn smoke_ofi_with_trade_enrichment() {
    let Some(mbp10_file) = try_data("mbp10/ES.FUT/ES.FUT_2025-Q2.dbn.zst") else { return };
    let Some(trades_file) = try_data("trades/ES.FUT/ES.FUT_2025-Q2.dbn.zst") else { return };

    // 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");

    info!(without = ofi_without.len(), with = ofi_with.len(), "OFI feature counts");

    // 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();
    info!(vpin_nonzero_without, vpin_nonzero_with, sample_size, "VPIN non-zero counts");
    info!(kyle_nonzero_without, kyle_nonzero_with, sample_size, "Kyle's Lambda non-zero counts");

    // 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];
    info!(
        ofi_l1 = sample_feat[0],
        ofi_l5 = sample_feat[1],
        depth_imb = sample_feat[2],
        vpin = sample_feat[3],
        kyle = sample_feat[4],
        bid_slope = sample_feat[5],
        ask_slope = sample_feat[6],
        trade_imb = sample_feat[7],
        "Sample enriched OFI"
    );
}

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

#[test]
fn smoke_trades_parse() {
    let Some(trades_dir) = try_data("trades/ES.FUT") else { return };

    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];
    info!(file = %file.display(), "Loading trades");

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

    info!(count = trades.len(), "Loaded trades");
    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
// ==========================================================================

mod gpu_smoke {
    use super::*;
    use std::sync::Arc;

    // candle eliminated — use MlDevice for stream extraction
    use ml_core::device::MlDevice;
    use ml::cuda_pipeline::gpu_experience_collector::{
        ExperienceCollectorConfig, GpuExperienceCollector,
    };
    use ml::cuda_pipeline::shared_cublas_handle::PerStreamCublasHandles;
    use ml::cuda_pipeline::mapped_pinned::MappedF32Buffer;

    type CudaStream = cudarc::driver::CudaStream;

    fn try_cuda() -> Option<(MlDevice, Arc<CudaStream>)> {
        match MlDevice::cuda(0) {
            Ok(dev) => {
                let stream = dev.cuda_stream().expect("CUDA stream required").clone();
                Some((dev, stream))
            }
            Err(_) => {
                warn!("SKIP: CUDA not available");
                None
            }
        }
    }

    fn has_enough_vram(min_mb: u64) -> bool {
        let vram = gpu_vram_mb();
        if vram < min_mb {
            warn!(vram, min_mb, "SKIP: insufficient GPU VRAM");
            false
        } else {
            true
        }
    }

    /// 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>,
        _stream: &Arc<CudaStream>,
    ) -> Result<(MappedF32Buffer, MappedF32Buffer, usize), anyhow::Error> {
        // 1. Load OHLCV bars
        let bars = ml::hyperopt::adapters::dbn_loader::load_bars_from_dbn_dir(ohlcv_dir)
            .expect("Failed to load OHLCV bars");
        info!(count = bars.len(), "Loaded OHLCV bars");

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

        // 3. Optionally load OFI features from MBP-10
        const OFI_DIM: usize = ml_core::state_layout::OFI_DIM;
        let ofi_features: Option<Vec<[f64; OFI_DIM]>> = 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()?;
            info!(file = %file.display(), "Loading OFI");
            ml::features::mbp10_loader::compute_ofi_from_file(file).ok().map(|ofi8| {
                ofi8.into_iter().map(|row| {
                    let mut wide = [0.0_f64; OFI_DIM];
                    wide[..8].copy_from_slice(&row);
                    wide
                }).collect::<Vec<_>>()
            })
        });

        let ofi_count = ofi_features.as_ref().map_or(0, |v| v.len());
        info!(ofi_count, "OFI features loaded");

        // 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 via mapped pinned (cuMemHostAlloc DEVICEMAP).
        // Safety: caller has confirmed a CUDA context is active (try_cuda() succeeded).
        let market_buf = unsafe { MappedF32Buffer::new(n * MARKET_DIM) }
            .expect("MappedF32Buffer alloc for market_buf");
        market_buf.write_from_slice(&market_data);

        let target_buf = unsafe { MappedF32Buffer::new(n * 4) }
            .expect("MappedF32Buffer alloc for target_buf");
        target_buf.write_from_slice(&target_data);

        Ok((market_buf, target_buf, n))
    }

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

    #[test]
    fn smoke_gpu_real_ohlcv_forward() -> Result<(), anyhow::Error> {
        let Some(ohlcv_dir) = try_data("ohlcv") else { return Ok(()) };
        let Some((_device, stream)) = try_cuda() else { return Ok(()) };

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

        info!(num_bars, market_dim = 51, "GPU buffer allocated");

        // Create collector (zero-copy: no weight sets needed at construction)
        let shared_cublas = Arc::new(PerStreamCublasHandles::new(&stream).unwrap());
        let mut collector = GpuExperienceCollector::new(
            stream.clone(),
            shared_cublas,
            0,     // bottleneck_dim (disabled for test)
            51,    // market_dim_cfg
            100_000.0,
            0.01, 0.05,
            (64, 64, 32, 32),
            (51, 51),
            4, 50,
            0.0,   // curiosity_weight (disabled for test)
        ).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,
            ..Default::default()
        };

        let batch = collector
            .collect_experiences_gpu(&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.n_episodes * batch.timesteps, expected_total);

        // Download actions and rewards from GPU to host for validation
        let mut actions_host = vec![0i32; batch.actions.len()];
        stream.memcpy_dtoh(&batch.actions, &mut actions_host).unwrap();
        let mut rewards_host = vec![0.0f32; batch.rewards.len()];
        stream.memcpy_dtoh(&batch.rewards, &mut rewards_host).unwrap();

        // Branching DQN: 9×3×3 = 81 factored actions (0..81)
        for (i, &a) in actions_host.iter().enumerate() {
            assert!((0..81).contains(&a), "action[{i}] = {a} out of range (expected 0..81 for 9-action branching DQN)");
        }

        // Rewards should reflect real price movements (not all zero)
        let nonzero_rewards = rewards_host.iter().filter(|&&r| r.abs() > 1e-8).count();
        let r_min = rewards_host.iter().cloned().fold(f32::INFINITY, f32::min);
        let r_max = rewards_host.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        info!(
            expected_total, nonzero_rewards,
            r_min, r_max,
            "Real OHLCV pipeline results"
        );
        Ok(())
    }

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

    #[test]
    fn smoke_gpu_real_ohlcv_plus_ofi_forward() -> Result<(), anyhow::Error> {
        if !has_enough_vram(2048) { return Ok(()) }
        let Some(ohlcv_dir) = try_data("ohlcv") else { return Ok(()) };
        let Some(mbp10_dir) = try_data("mbp10/ES.FUT") else { return Ok(()) };
        let Some((_device, stream)) = try_cuda() else { return Ok(()) };

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

        info!(num_bars, market_dim = 51, "GPU buffer allocated (with OFI)");

        let shared_cublas = Arc::new(PerStreamCublasHandles::new(&stream).unwrap());
        let mut collector = GpuExperienceCollector::new(
            stream.clone(),
            shared_cublas,
            0,     // bottleneck_dim (disabled for test)
            51,    // market_dim_cfg
            100_000.0,
            0.01, 0.05,
            (64, 64, 32, 32),
            (51, 51),
            4, 50,
            0.0,   // curiosity_weight (disabled for test)
        ).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,
            ..Default::default()
        };

        let batch = collector
            .collect_experiences_gpu(&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.n_episodes * batch.timesteps, expected_total);

        // Download rewards from GPU to host for validation
        let mut rewards_host = vec![0.0f32; batch.rewards.len()];
        stream.memcpy_dtoh(&batch.rewards, &mut rewards_host).unwrap();

        for (i, &r) in rewards_host.iter().enumerate() {
            assert!(r.is_finite(), "OFI reward[{i}] = {r} is not finite");
        }

        let r_min = rewards_host.iter().cloned().fold(f32::INFINITY, f32::min);
        let r_max = rewards_host.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        info!(expected_total, r_min, r_max, "OHLCV+OFI pipeline results");
        Ok(())
    }

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

    #[test]
    fn smoke_gpu_real_data_noisy_distributional() -> Result<(), anyhow::Error> {
        let Some(ohlcv_dir) = try_data("ohlcv") else { return Ok(()) };
        let Some((_device, stream)) = try_cuda() else { return Ok(()) };

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

        let shared_cublas = Arc::new(PerStreamCublasHandles::new(&stream).unwrap());
        let mut collector = GpuExperienceCollector::new(
            stream.clone(),
            shared_cublas,
            0,     // bottleneck_dim (disabled for test)
            51,    // market_dim_cfg
            100_000.0,
            0.01, 0.05,
            (64, 64, 32, 32),
            (51, 51),
            4, 50,
            0.0,   // curiosity_weight (disabled for test)
        ).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, // 10% random ensures ≥2 distinct actions in 200 samples
            gamma: 0.99,
            // sigma must be large enough that NoisyNet noise dominates Q-value gaps
            // even on wide networks (256-dim). With random Xavier init on a 256-wide
            // network, Q-values are O(1/sqrt(256)) ≈ 0.06 apart. sigma=2.0 ensures
            // noise O(sigma/sqrt(fan_in)) ≈ 0.12 exceeds this gap.
            noisy_sigma_init: 2.0,
            num_atoms: 51,
            v_min: -10.0,
            v_max: 10.0,
            ..Default::default()
        };

        let batch = collector
            .collect_experiences_gpu(&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.n_episodes * batch.timesteps, expected_total);

        // Download actions and rewards from GPU to host for validation
        let mut actions_host = vec![0i32; batch.actions.len()];
        stream.memcpy_dtoh(&batch.actions, &mut actions_host).unwrap();
        let mut rewards_host = vec![0.0f32; batch.rewards.len()];
        stream.memcpy_dtoh(&batch.rewards, &mut rewards_host).unwrap();

        // Rewards should be finite
        for (i, &r) in rewards_host.iter().enumerate() {
            assert!(r.is_finite(), "C51 reward[{i}] = {r} is not finite");
        }

        // NoisyNet should produce diverse actions (not all the same)
        // Branching DQN: 9×3×3 = 81 factored actions
        let mut action_counts = [0_usize; 45];
        for &a in &actions_host {
            if let Some(count) = action_counts.get_mut(a as usize) {
                *count += 1;
            }
        }
        let unique_actions = action_counts.iter().filter(|&&c| c > 0).count();
        info!(unique_actions, "C51+NoisyNet action distribution");
        assert!(
            unique_actions >= 2,
            "NoisyNet should produce at least 2 distinct actions, got {unique_actions}"
        );

        let r_min = rewards_host.iter().cloned().fold(f32::INFINITY, f32::min);
        let r_max = rewards_host.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        info!(r_min, r_max, "C51 reward range");
        Ok(())
    }
}

// ==========================================================================
// Test 6: E2E DQN Training Loop — validates the full training pipeline
// using local smoke-test OHLCV data with branching DQN on CUDA GPU path.
//
// Validates:
//   1. All epochs complete (no crash, no premature early stopping)
//   2. Loss decreases (gradient flow works through branching network)
//   3. All losses finite (no NaN/Inf from gradient clipping)
//   4. Q-value divergence (model develops action preferences)
//   5. Action diversity uses 81-action factored space
//   6. Gradient norms are non-zero (clipping is functional)
// ==========================================================================

#[tokio::test]
async fn smoke_e2e_dqn_training_loop() {
    use ml::trainers::dqn::{DQNHyperparameters, DQNTrainer};

    if gpu_vram_mb() > 0 && gpu_vram_mb() < 2048 {
        return;
    }
    let Some(ohlcv_dir) = try_data("ohlcv") else {
        return;
    };
    let data_dir = ohlcv_dir.to_string_lossy().to_string();

    // Configure for a fast smoke run: few epochs, small batch, low warmup
    let mut hyperparams = DQNHyperparameters::conservative();
    hyperparams.epochs = 2;
    hyperparams.batch_size = 32;
    hyperparams.learning_rate = 0.0003;
    hyperparams.epsilon_start = 1.0;
    hyperparams.epsilon_end = 0.1;
    hyperparams.epsilon_decay = 0.7;
    hyperparams.early_stopping_enabled = false;
    hyperparams.checkpoint_frequency = 100;
    hyperparams.warmup_steps = 0;
    hyperparams.min_replay_size = 50;
    // CI smoke: 16 episodes x 50 timesteps = 800 experiences/epoch.
    // Exercises full fused CUDA kernel (branching+C51+NoisyNets+DSR).
    hyperparams.gpu_timesteps_per_episode = 10;
    // CI: cap training steps to avoid full 204K-bar dataset sweep (6375->64 steps).
    // 64 steps x batch_size 32 = 2048 gradient updates — sufficient to validate
    // finite loss, gradient flow, and action diversity without 400s/epoch overhead.
    // Curiosity disabled: kernel crashes on RTX 3050 (needs investigation)
    // C51 atom count from GPU profile (replaces hardcoded VRAM if/else)
    let gpu_profile = ml_core::gpu::profile::GpuProfile::load();
    hyperparams.num_atoms = gpu_profile.training.num_atoms;

    let checkpoint_dir = tempfile::tempdir().expect("Failed to create temp dir");
    let ckpt_path = checkpoint_dir.path().to_path_buf();
    let mut trainer = DQNTrainer::new(hyperparams).expect("Failed to create DQN trainer");

    info!(data_dir, "Starting E2E DQN training smoke test");

    let metrics = trainer
        .train(&data_dir, "ES.FUT", move |epoch, checkpoint_data, is_best| {
            if is_best {
                let path = ckpt_path.join("smoke_e2e_best.safetensors");
                std::fs::write(&path, &checkpoint_data)?;
                Ok(path.to_string_lossy().to_string())
            } else {
                Ok(format!("epoch_{epoch}"))
            }
        })
        .await
        .expect("DQN training failed");

    // === ASSERTIONS ===

    // 1. All epochs completed
    assert_eq!(
        metrics.epochs_trained, 2,
        "Expected 2 epochs, got {}",
        metrics.epochs_trained
    );

    // 2. Loss decreases (gradient flow works)
    let loss_history = trainer.loss_history();
    assert!(
        loss_history.len() >= 2,
        "Need at least 2 epochs of loss history, got {}",
        loss_history.len()
    );
    let initial_loss = loss_history.first().copied().unwrap_or(f64::NAN);
    let final_loss = loss_history.last().copied().unwrap_or(f64::NAN);
    info!(initial_loss, final_loss, "Loss trajectory");

    // 3. All losses must be finite (NaN would indicate broken gradient clipping)
    for (i, &loss) in loss_history.iter().enumerate() {
        assert!(
            loss.is_finite(),
            "Loss at epoch {i} is not finite: {loss}"
        );
    }

    // 4. Q-value divergence (model has preferences)
    let avg_q = metrics
        .additional_metrics
        .get("avg_q_value")
        .copied()
        .unwrap_or(0.0);
    info!(avg_q, "Average Q-value");

    // 5. Action diversity uses 81-action factored space
    let action_space = metrics
        .additional_metrics
        .get("action_space_size")
        .copied()
        .unwrap_or(5.0);
    let action_diversity = metrics
        .additional_metrics
        .get("action_diversity")
        .copied()
        .unwrap_or(0.0);
    let total_actions = metrics
        .additional_metrics
        .get("total_actions")
        .copied()
        .unwrap_or(0.0);
    info!(action_space, action_diversity, total_actions, "Action metrics");
    // Branching DQN: action_space=81 (9×3×3) when factored counts populated,
    // or 9 (exposure-only fallback) for very short runs.
    assert!(
        action_space >= 9.0 || total_actions == 0.0,
        "Branching DQN should report >=9 action space, got {action_space}"
    );

    // 6. Gradient norms should be non-zero (clipping is functional)
    let avg_grad_norm = metrics
        .additional_metrics
        .get("avg_gradient_norm")
        .copied()
        .unwrap_or(0.0);
    info!(avg_grad_norm, "Average gradient norm");
    // With the TensorId fallback fix, grad norms should be reported correctly
    assert!(
        avg_grad_norm.is_finite(),
        "Gradient norm should be finite, got {avg_grad_norm}"
    );

    // === REPORT ===
    let loss_change = if initial_loss > 0.0 && initial_loss.is_finite() {
        (1.0 - final_loss / initial_loss) * 100.0
    } else {
        f64::NAN
    };
    info!(
        epochs = metrics.epochs_trained,
        initial_loss, final_loss, loss_change,
        avg_q, action_space, action_diversity,
        avg_grad_norm,
        training_time_s = metrics.training_time_seconds,
        "E2E DQN training smoke test report"
    );
}