foxhunt/crates/ml/tests/sp21_per_trade_predicted_q_test.rs

//! SP21 T2.2 Phase 1.5 (2026-05-11) — `predicted_q` per-trade tape GPU oracle test.
//!
//! Validates the `entry_q` tracking added to `backtest_env_step_batch`:
//!
//!   1. **predicted_q populated on close** — open a Long position at step 0
//!      with a controlled `q_values_per_window[Long_Full] = 2.5`, hold for two
//!      bars, close with Flat at step 3. Read the per-trade tape and assert
//!      the single emitted trade carries `predicted_q ≈ 2.5` (the entry_q
//!      captured at open and snapshotted before the close-detection block).
//!   2. **NULL q_values_per_window keeps predicted_q at 0** — same scenario
//!      with `q_values_per_window = NULL` (single-step `evaluate()` semantics).
//!      Kernel's open-time write is gated on the NULL guard, so
//!      `entry_q_state[w]` stays 0.0 and the tape's `predicted_q` is 0.0.
//!   3. **entry_q persists across hold bars** — the persistent `entry_q_state`
//!      buffer is per-window, written ONCE at open, read back at close. Hold
//!      bars don't touch it. The test indirectly proves persistence by the
//!      single-trade success of test (1).
//!
//! These cover the entry_q capture path end-to-end at the kernel level. The
//! eval pipeline (cuBLAS forward → action select → env step) integration test
//! lives in `gpu_backtest_validation.rs`; this file focuses on the per-trade
//! tape's new `predicted_q` field as a kernel-direct oracle.
//!
//! Run on a GPU host:
//!
//!     SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 \
//!       cargo test -p ml --test sp21_per_trade_predicted_q_test --features cuda \
//!         -- --ignored --nocapture

#![allow(clippy::tests_outside_test_module)]
#![allow(clippy::too_many_lines)]

#[cfg(feature = "cuda")]
#[allow(unsafe_code)] // CUDA kernel launch.
mod gpu {
    use std::sync::Arc;

    use cudarc::driver::{
        CudaContext, CudaFunction, CudaStream, DevicePtr, LaunchConfig, PushKernelArg,
    };

    /// Same cubin loaded by `gpu_backtest_evaluator.rs`. Test-side
    /// `include_bytes!` reads from the per-crate `OUT_DIR` set by
    /// `crates/ml/build.rs` — the same OUT_DIR used by the lib build.
    const ENV_CUBIN: &[u8] = include_bytes!(concat!(
        env!("OUT_DIR"),
        "/backtest_env_kernel.cubin"
    ));

    /// Branch sizes — match `DqnBacktestConfig::from_network_dims` defaults.
    const B0: i32 = 4; // Short / Hold / Long / Flat
    const B1: i32 = 3; // Quarter / Half / Full
    const B2: i32 = 3; // Order type
    const B3: i32 = 3; // Patient / Normal / Aggressive
    /// Total factored action count. Matches the kernel's
    /// `num_actions` arg.
    const NUM_ACTIONS: i32 = B0 * B1 * B2 * B3; // 108

    /// Direction constants (mirrors `state_layout.cuh:123-126`).
    const DIR_LONG: i32 = 2;
    const DIR_HOLD: i32 = 1;
    const DIR_FLAT: i32 = 3;

    /// Magnitude constants (mirrors `state_layout.cuh:133-135`).
    const MAG_FULL: i32 = 2;
    const MAG_QUARTER: i32 = 0;

    /// Encode a factored action: `dir * b1*b2*b3 + mag * b2*b3 + order * b3 + urgency`.
    /// Mirrors `trade_physics.cuh::decode_direction_4b` etc. exactly so the
    /// kernel decoder maps our test actions back to the canonical DIR/MAG.
    const fn encode_action(dir: i32, mag: i32, order: i32, urgency: i32) -> i32 {
        dir * (B1 * B2 * B3) + mag * (B2 * B3) + order * B3 + urgency
    }

    /// `MAX_TRADES_PER_WINDOW` matches the kernel guard at the per-trade
    /// emission site. We size at 16 here — far smaller than production's
    /// 200_000 — because the test only emits a handful of trades and the
    /// per-window buffer alloc is `n_windows × max_trades`.
    const MAX_TRADES_PER_WINDOW: i32 = 16;

    /// Tolerance for f32 comparisons against expected entry_q values.
    /// The kernel writes the Q-value verbatim from `q_values_per_window`
    /// (no math), so the round-trip is exact under IEEE-754 — but a small
    /// epsilon defends against any future float-rewrite of the open path.
    const Q_TOL: f32 = 1e-5;

    /// 4-bar single-window scenario: open Long Full at step 0, hold for two
    /// bars, Flat-close at step 3. Returns (n_windows=1, max_len=4, chunk_len=4).
    const N_WINDOWS: usize = 1;
    const MAX_LEN: usize = 4;
    const CHUNK_LEN: usize = 4;
    const FEATURE_DIM: i32 = 4; // < 41 → kernel uses fixed-width trail (no features deref)

    fn try_cuda_stream() -> Option<Arc<CudaStream>> {
        match CudaContext::new(0) {
            Ok(ctx) => Some(ctx.default_stream()),
            Err(_) => {
                eprintln!("CUDA not available, skipping per-trade-tape predicted_q test");
                None
            }
        }
    }

    fn load_env_batch_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
        let module = stream
            .context()
            .load_cubin(ENV_CUBIN.to_vec())
            .expect("load backtest_env_kernel.cubin");
        module
            .load_function("backtest_env_step_batch")
            .expect("load backtest_env_step_batch function")
    }

    /// One window of OHLC. All bars at price 100.0 (flat market) → trade
    /// closes at break-even, isolating the entry_q signal from any P&L
    /// noise. The close emission predicate `pre_entry_price > 0.0f &&
    /// value > 1.0f` only requires entry_price was non-zero (set on open
    /// at step 0) and capital is positive — both satisfied with flat
    /// prices and `initial_capital = 100_000`.
    fn flat_market_prices() -> Vec<f32> {
        let mut prices = Vec::with_capacity(N_WINDOWS * MAX_LEN * 4);
        for _ in 0..(N_WINDOWS * MAX_LEN) {
            // OHLC = [100, 100, 100, 100]
            prices.push(100.0_f32);
            prices.push(100.0_f32);
            prices.push(100.0_f32);
            prices.push(100.0_f32);
        }
        prices
    }

    /// Initial portfolio state matching `GpuBacktestEvaluator::init_portfolio_state`.
    fn init_portfolio_state(initial_capital: f32) -> Vec<f32> {
        let mut state = vec![0.0_f32; N_WINDOWS * 8];
        state[0] = initial_capital; // value
        state[2] = initial_capital; // cash
        state[4] = initial_capital; // max_equity
        state
    }

    /// Action sequence: open Long Full at step 0, hold, hold, close Flat.
    /// Layout `[chunk_len, n_windows]` flattened — the kernel reads via
    /// `chunked_actions[s * n_windows + w]`.
    fn open_hold_close_actions() -> Vec<i32> {
        // step 0: Long Full Market Normal (open)
        let a_open = encode_action(DIR_LONG, MAG_FULL, 0, 1);
        // step 1, 2: Hold Quarter Market Normal (carry the position; Hold = no-op,
        // see `project_hold_action_design.md`)
        let a_hold = encode_action(DIR_HOLD, MAG_QUARTER, 0, 1);
        // step 3: Flat Quarter Market Normal (close all positions to zero)
        let a_close = encode_action(DIR_FLAT, MAG_QUARTER, 0, 1);
        vec![a_open, a_hold, a_hold, a_close]
    }

    /// Q-values buffer layout `[chunk_len, n_windows, num_actions]` flattened.
    /// The kernel indexes via `q[s * n_windows * num_actions + w * num_actions
    /// + action_taken]` — at step 0 we set the slot for the open action only;
    /// other slots stay at 0.0 (the kernel never reads them since it always
    /// indexes by the chosen action).
    fn q_values_with_open_at(open_action: i32, open_q: f32) -> Vec<f32> {
        let n = N_WINDOWS;
        let mut q = vec![0.0_f32; CHUNK_LEN * n * NUM_ACTIONS as usize];
        let s = 0_usize;
        let w = 0_usize;
        let a = open_action as usize;
        q[s * n * NUM_ACTIONS as usize + w * NUM_ACTIONS as usize + a] = open_q;
        q
    }

    /// Configuration constants for the kernel — match `GpuBacktestConfig::default`
    /// where it matters and zero out the optional shaping params.
    struct KernelConfig {
        max_position: f32,
        tx_cost_bps: f32,
        spread_cost: f32,
        contract_multiplier: f32,
        margin_pct: f32,
        initial_capital: f32,
    }

    impl KernelConfig {
        fn for_test() -> Self {
            Self {
                max_position: 1.0,
                // Zero costs so any P&L deviation comes from the kernel logic,
                // not transaction friction. The test asserts on `predicted_q`
                // (independent of P&L) but the close emission predicate's
                // `value > 1.0f` guard is more comfortably satisfied with
                // no cost drain.
                tx_cost_bps: 0.0,
                spread_cost: 0.0,
                contract_multiplier: 1.0,
                margin_pct: 0.0,
                initial_capital: 100_000.0,
            }
        }
    }

    /// Run the batch kernel one chunk over the 4-bar window with the given
    /// q_values_per_window pointer (real or NULL=0u64) and `num_actions`.
    /// Returns (per_trade_count[0], per_trade_predicted_q[0..count]).
    #[allow(clippy::too_many_arguments)]
    fn run_one_chunk(
        stream: &Arc<CudaStream>,
        kernel: &CudaFunction,
        cfg: &KernelConfig,
        actions: &[i32],
        q_values_dev_ptr: u64, // 0 = NULL
        num_actions: i32,
    ) -> (u32, Vec<f32>) {
        // ── Allocate input buffers ──
        let mut prices_dev = stream
            .alloc_zeros::<f32>(N_WINDOWS * MAX_LEN * 4)
            .expect("alloc prices");
        let prices_host = flat_market_prices();
        stream.memcpy_htod(&prices_host, &mut prices_dev).expect("HtoD prices");

        let mut window_lens_dev = stream
            .alloc_zeros::<i32>(N_WINDOWS)
            .expect("alloc window_lens");
        let window_lens_host = vec![MAX_LEN as i32; N_WINDOWS];
        stream
            .memcpy_htod(&window_lens_host, &mut window_lens_dev)
            .expect("HtoD window_lens");

        // Features buffer — sized `[n * max_len * feature_dim]`. With
        // FEATURE_DIM=4 (< 41), the kernel's `compute_regime_trail_scales`
        // takes the fallback (no feature deref), so zeros are safe.
        let features_dev = stream
            .alloc_zeros::<f32>(N_WINDOWS * MAX_LEN * FEATURE_DIM as usize)
            .expect("alloc features");

        let mut actions_dev = stream
            .alloc_zeros::<i32>(CHUNK_LEN * N_WINDOWS)
            .expect("alloc chunked_actions");
        stream
            .memcpy_htod(actions, &mut actions_dev)
            .expect("HtoD chunked_actions");

        let mut portfolio_dev = stream
            .alloc_zeros::<f32>(N_WINDOWS * 8)
            .expect("alloc portfolio");
        let portfolio_host = init_portfolio_state(cfg.initial_capital);
        stream
            .memcpy_htod(&portfolio_host, &mut portfolio_dev)
            .expect("HtoD portfolio");

        let step_rewards_dev = stream
            .alloc_zeros::<f32>(N_WINDOWS)
            .expect("alloc step_rewards");
        let step_returns_dev = stream
            .alloc_zeros::<f32>(N_WINDOWS * MAX_LEN)
            .expect("alloc step_returns");
        let done_dev = stream
            .alloc_zeros::<i32>(N_WINDOWS)
            .expect("alloc done_flags");
        let actions_history_dev = stream
            .alloc_zeros::<i32>(N_WINDOWS * MAX_LEN)
            .expect("alloc actions_history");
        let kelly_stats_dev = stream
            .alloc_zeros::<f32>(N_WINDOWS * 4)
            .expect("alloc kelly_stats");

        // Per-trade tape buffers (5 SoA + count + 1 new SoA for predicted_q).
        let per_trade_pnl_dev = stream
            .alloc_zeros::<f32>(N_WINDOWS * MAX_TRADES_PER_WINDOW as usize)
            .expect("alloc per_trade_pnl");
        let per_trade_holding_bars_dev = stream
            .alloc_zeros::<u32>(N_WINDOWS * MAX_TRADES_PER_WINDOW as usize)
            .expect("alloc per_trade_holding_bars");
        let per_trade_bar_index_dev = stream
            .alloc_zeros::<u32>(N_WINDOWS * MAX_TRADES_PER_WINDOW as usize)
            .expect("alloc per_trade_bar_index");
        let per_trade_dir_mag_dev = stream
            .alloc_zeros::<u32>(N_WINDOWS * MAX_TRADES_PER_WINDOW as usize)
            .expect("alloc per_trade_dir_mag");
        let per_trade_count_dev = stream
            .alloc_zeros::<u32>(N_WINDOWS)
            .expect("alloc per_trade_count");
        let per_trade_predicted_q_dev = stream
            .alloc_zeros::<f32>(N_WINDOWS * MAX_TRADES_PER_WINDOW as usize)
            .expect("alloc per_trade_predicted_q");
        // Phase 1.5 entry_q persistent state — one float per window.
        // Initialized to 0.0; the kernel writes it on open / reverse.
        let entry_q_state_dev = stream
            .alloc_zeros::<f32>(N_WINDOWS)
            .expect("alloc entry_q_state");

        // Scalar args.
        let n_i32 = N_WINDOWS as i32;
        let max_len_i32 = MAX_LEN as i32;
        let start_step_i32 = 0_i32;
        let chunk_len_i32 = CHUNK_LEN as i32;
        let max_trades_i32 = MAX_TRADES_PER_WINDOW;
        let feature_dim_i32 = FEATURE_DIM;

        // ── Launch ──
        let env_blocks = ((N_WINDOWS as u32) + 255) / 256;
        let launch_cfg = LaunchConfig {
            grid_dim: (env_blocks.max(1), 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };

        unsafe {
            stream
                .launch_builder(kernel)
                .arg(&prices_dev)
                .arg(&window_lens_dev)
                .arg(&features_dev)
                .arg(&feature_dim_i32)
                .arg(&actions_dev)
                .arg(&portfolio_dev)
                .arg(&step_rewards_dev)
                .arg(&step_returns_dev)
                .arg(&done_dev)
                .arg(&actions_history_dev)
                .arg(&n_i32)
                .arg(&max_len_i32)
                .arg(&cfg.max_position)
                .arg(&cfg.tx_cost_bps)
                .arg(&cfg.spread_cost)
                .arg(&start_step_i32)
                .arg(&chunk_len_i32)
                .arg(&B0)
                .arg(&B1)
                .arg(&B2)
                .arg(&B3)
                .arg(&0.0_f32) // holding_cost_rate (unused per Phase 3)
                .arg(&0.0_f32) // churn_threshold (unused)
                .arg(&0.0_f32) // churn_penalty_scale (unused)
                .arg(&cfg.contract_multiplier)
                .arg(&cfg.margin_pct)
                .arg(&0.0_f32) // opp_cost_scale (unused)
                .arg(&kelly_stats_dev)
                .arg(&0u64) // isv_signals NULL → health=0.5 fallback
                .arg(&0u64) // exploration_scale_ptr NULL
                .arg(&0u64) // shaping_scale_ptr NULL
                .arg(&0u64) // conviction_ptr NULL → 1.0 fallback
                .arg(&0u64) // magnitude_conviction_ptr NULL → 0.0 fallback
                .arg(&per_trade_pnl_dev)
                .arg(&per_trade_holding_bars_dev)
                .arg(&per_trade_bar_index_dev)
                .arg(&per_trade_dir_mag_dev)
                .arg(&per_trade_count_dev)
                .arg(&max_trades_i32)
                // SP21 T2.2 Phase 1.5 args
                .arg(&entry_q_state_dev)
                .arg(&q_values_dev_ptr) // 0 = NULL for the NULL-tolerance test
                .arg(&num_actions)
                .arg(&per_trade_predicted_q_dev)
                .launch(launch_cfg)
                .expect("backtest_env_step_batch launch");
        }

        stream.synchronize().expect("post-launch sync");

        // ── Read tape ──
        let mut count_host = vec![0u32; N_WINDOWS];
        stream
            .memcpy_dtoh(&per_trade_count_dev, &mut count_host)
            .expect("DtoH per_trade_count");
        let mut predicted_q_host =
            vec![0.0_f32; N_WINDOWS * MAX_TRADES_PER_WINDOW as usize];
        stream
            .memcpy_dtoh(&per_trade_predicted_q_dev, &mut predicted_q_host)
            .expect("DtoH per_trade_predicted_q");

        // Drop device buffers explicitly to free GPU memory before the
        // next test runs (each #[test] gets its own isolated alloc set).
        drop(prices_dev);
        drop(window_lens_dev);
        drop(features_dev);
        drop(actions_dev);
        drop(portfolio_dev);
        drop(step_rewards_dev);
        drop(step_returns_dev);
        drop(done_dev);
        drop(actions_history_dev);
        drop(kelly_stats_dev);
        drop(per_trade_pnl_dev);
        drop(per_trade_holding_bars_dev);
        drop(per_trade_bar_index_dev);
        drop(per_trade_dir_mag_dev);
        drop(per_trade_count_dev);
        drop(per_trade_predicted_q_dev);
        drop(entry_q_state_dev);

        let count_w0 = count_host[0];
        let predicted_q_w0: Vec<f32> = predicted_q_host[0..(count_w0 as usize)].to_vec();
        (count_w0, predicted_q_w0)
    }

    /// Test 1: production path (real q_values) populates `predicted_q` from
    /// the entry_q snapshot at trade open.
    #[test]
    #[ignore = "requires GPU"]
    fn predicted_q_populated_on_close_with_real_q_values() {
        let stream = match try_cuda_stream() {
            Some(s) => s,
            None => return,
        };
        let kernel = load_env_batch_kernel(&stream);

        let cfg = KernelConfig::for_test();
        let actions = open_hold_close_actions();
        let open_action = actions[0];

        // The Q-value we expect predicted_q to carry across the trade.
        let expected_entry_q: f32 = 2.5;
        let q_host = q_values_with_open_at(open_action, expected_entry_q);
        let mut q_dev = stream
            .alloc_zeros::<f32>(q_host.len())
            .expect("alloc q_values");
        stream.memcpy_htod(&q_host, &mut q_dev).expect("HtoD q_values");
        let q_dev_ptr = {
            let (ptr, _g) = q_dev.device_ptr(&stream);
            ptr
        };

        let (count, predicted_q) = run_one_chunk(
            &stream, &kernel, &cfg, &actions, q_dev_ptr, NUM_ACTIONS,
        );

        // Drop the q_values buffer ourselves — it's owned by this test only.
        drop(q_dev);

        assert_eq!(
            count, 1,
            "expected exactly 1 trade closed (open Long@step 0, Flat@step 3); got count={count}"
        );
        assert!(
            !predicted_q.is_empty(),
            "predicted_q tape should have ≥1 entry when count={count}"
        );
        let actual_q = predicted_q[0];
        assert!(
            (actual_q - expected_entry_q).abs() < Q_TOL,
            "predicted_q[0] = {actual_q}, expected {expected_entry_q} (entry_q captured from \
             q_values_per_window[w=0, a={open_action}] at step-0 open and snapshotted before the \
             step-3 close emission)"
        );
    }

    /// Test 2: NULL `q_values_per_window` (single-step `evaluate()` semantics)
    /// keeps `predicted_q` at the buffer-init 0.0 value. The kernel's
    /// open-time write is gated on the NULL guard at the bottom of the
    /// per-step loop, so `entry_q_state[w]` stays 0.0 and the close emit
    /// snapshots 0.0 into the tape.
    #[test]
    #[ignore = "requires GPU"]
    fn predicted_q_stays_zero_when_q_values_is_null() {
        let stream = match try_cuda_stream() {
            Some(s) => s,
            None => return,
        };
        let kernel = load_env_batch_kernel(&stream);

        let cfg = KernelConfig::for_test();
        let actions = open_hold_close_actions();

        // q_values_per_window = NULL (0u64). num_actions = 0 so the kernel's
        // open-time guard `num_actions > 0` short-circuits independently.
        let (count, predicted_q) = run_one_chunk(
            &stream, &kernel, &cfg, &actions, 0u64, 0,
        );

        assert_eq!(
            count, 1,
            "trade still closes regardless of q_values availability; got count={count}"
        );
        let actual_q = predicted_q[0];
        assert!(
            actual_q.abs() < Q_TOL,
            "predicted_q[0] = {actual_q}, expected ~0.0 (NULL q_values_per_window means the \
             kernel skips the entry_q write at open; the persistent entry_q_state[w] stays at \
             buffer-init 0.0 and the close emission snapshots 0.0 into the tape)"
        );
    }

    /// Test 3: explicit verification that the entry_q write is GATED on
    /// `is_open || is_reverse_open` — a sequence with NO open (e.g.,
    /// flat-flat-flat-flat) should produce zero trades and never touch
    /// `entry_q_state`. Provides an additional path-coverage check beyond
    /// tests (1) and (2).
    #[test]
    #[ignore = "requires GPU"]
    fn no_trades_no_predicted_q_emission() {
        let stream = match try_cuda_stream() {
            Some(s) => s,
            None => return,
        };
        let kernel = load_env_batch_kernel(&stream);

        let cfg = KernelConfig::for_test();
        // All Flat actions: no position ever opens, no close emission.
        let a_flat = encode_action(DIR_FLAT, MAG_QUARTER, 0, 1);
        let actions = vec![a_flat; CHUNK_LEN];

        let q_host = q_values_with_open_at(0, 99.9); // won't be read
        let mut q_dev = stream
            .alloc_zeros::<f32>(q_host.len())
            .expect("alloc q_values");
        stream.memcpy_htod(&q_host, &mut q_dev).expect("HtoD q_values");
        let q_dev_ptr = {
            let (ptr, _g) = q_dev.device_ptr(&stream);
            ptr
        };

        let (count, _predicted_q) = run_one_chunk(
            &stream, &kernel, &cfg, &actions, q_dev_ptr, NUM_ACTIONS,
        );

        drop(q_dev);

        assert_eq!(
            count, 0,
            "all-Flat action sequence should yield zero trades; got count={count}"
        );
    }
}