#![allow(unsafe_code)] // Required for CUDA kernel launch + mapped pinned memory
#![allow(dead_code)] // Populated by Tasks 3-8 — wrappers are exercised by tests added in subsequent plan tasks.

//! Phase 0 GPU-direct verification tests for the Thompson sampling
//! design (spec: docs/superpowers/specs/2026-04-26-distributional-rl-aggregation-design.md).
//!
//! All tests in this module are `#[ignore]` because they require a CUDA
//! GPU. Run them with:
//!     cargo test -p ml --lib distributional_q_tests -- --ignored --nocapture
//!
//! ## No memcpy_dtoh on the test path
//!
//! Earlier revisions of this module wrapped the Thompson kernel with a
//! single-pick launcher that copied one `i32` back via `clone_dtoh`,
//! and called it inside a 100 000-seed `for` loop. That violates the
//! `feedback_gpu_cpu_roundtrip.md` invariant — even on a test path,
//! per-launch device-to-host copies are not how this codebase ships.
//!
//! The current design matches `cuda_pipeline/gpu_training_guard.rs`
//! `MappedBuffer`: kernels write `int` outputs to a device pointer
//! obtained via `cuMemHostGetDevicePointer_v2` for a host allocation
//! made with `cuMemHostAlloc(DEVICEMAP|PORTABLE)`. The kernel finishes
//! each write with `__threadfence_system()` so the host can read
//! through the same allocation via `read_volatile`. No `memcpy_dtoh`
//! anywhere — and one batched launch replaces the 10 000- and
//! 100 000-iteration loops.

use std::sync::{Arc, OnceLock};
use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
use crate::cuda_pipeline::mapped_pinned::{MappedF32Buffer, MappedI32Buffer};

const TEST_KERNEL_CUBIN: &[u8] =
    include_bytes!(concat!(env!("OUT_DIR"), "/thompson_test_kernel.cubin"));

const N_IQN_QUANTILES: usize = 5;
const B0_SIZE: usize = 4;  // 4-direction action space

// ── Cached kernel set ─────────────────────────────────────────────────────

/// Cached kernel handle set; loaded once per process.
struct KernelSet {
    thompson_direction_batched: CudaFunction,
    argmax_eq: CudaFunction,
    sigma_c51: CudaFunction,
}

/// Process-wide cached `KernelSet`. Loading the cubin and resolving function
/// handles on every test invocation would dominate the kernel time. Pattern
/// matches `cuda_pipeline/signal_adapter.rs` and `cuda_pipeline/gpu_her.rs`.
static KERNELS: OnceLock<KernelSet> = OnceLock::new();

fn kernels(stream: &Arc<CudaStream>) -> &'static KernelSet {
    KERNELS.get_or_init(|| {
        let module = stream
            .context()
            .load_cubin(TEST_KERNEL_CUBIN.to_vec())
            .expect("load thompson_test_kernel cubin");
        let thompson_direction_batched = module
            .load_function("thompson_direction_test_batched")
            .expect("load thompson_direction_test_batched");
        let argmax_eq = module
            .load_function("argmax_eq_test")
            .expect("load argmax_eq_test");
        let sigma_c51 = module
            .load_function("compute_sigma_c51_test")
            .expect("load compute_sigma_c51_test");
        KernelSet { thompson_direction_batched, argmax_eq, sigma_c51 }
    })
}

// ── Kernel launchers (batched, pinned-mapped output) ──────────────────────

/// Launch the batched Thompson direction kernel and return one direction
/// pick per seed. Output is read directly from mapped pinned memory — the
/// host never issues a `memcpy_dtoh`. A single launch replaces the 10 000-
/// or 100 000-iteration single-pick loops the earlier wrappers used.
fn launch_thompson_direction_batched(
    stream: &Arc<CudaStream>,
    c51_probs: &CudaSlice<f32>,
    atom_values: &CudaSlice<f32>,
    iqn_quantiles: &CudaSlice<f32>,
    n_atoms: i32,
    base_seed: u32,
    n_seeds: usize,
) -> Vec<u8> {
    assert!(n_seeds > 0, "launch_thompson_direction_batched: n_seeds must be > 0");
    let kernel = &kernels(stream).thompson_direction_batched;

    // Safety: CUDA context is active on this thread (we just resolved the
    // kernel through the stream's context above).
    let buffer = unsafe { MappedI32Buffer::new(n_seeds) }
        .expect("alloc MappedI32Buffer for thompson seeds");

    let n_seeds_i32 = i32::try_from(n_seeds).expect("n_seeds fits in i32");
    let block_dim_x: u32 = 256;
    let grid_dim_x: u32 = (n_seeds as u32).div_ceil(block_dim_x);

    let dev_ptr = buffer.dev_ptr;

    unsafe {
        stream.launch_builder(kernel)
            .arg(c51_probs)
            .arg(atom_values)
            .arg(iqn_quantiles)
            .arg(&(B0_SIZE as i32))
            .arg(&n_atoms)
            .arg(&base_seed)
            .arg(&n_seeds_i32)
            .arg(&dev_ptr)
            .launch(LaunchConfig {
                grid_dim: (grid_dim_x, 1, 1),
                block_dim: (block_dim_x, 1, 1),
                shared_mem_bytes: 0,
            })
            .expect("launch thompson_direction_test_batched");
    }

    stream.synchronize().expect("synchronize after thompson batched launch");

    buffer
        .read_all()
        .into_iter()
        .enumerate()
        .map(|(idx, raw)| {
            assert!(
                (0..(B0_SIZE as i32)).contains(&raw),
                "thompson kernel wrote direction index {} outside [0, {}) at seed {}",
                raw,
                B0_SIZE,
                idx,
            );
            raw as u8
        })
        .collect()
}

/// Launch the eval-mode argmax kernel. Output is single-i32 mapped pinned
/// memory — same mechanism as the batched Thompson path so the test code
/// has one read pattern to reason about.
fn launch_argmax_eq_pinned(
    stream: &Arc<CudaStream>,
    c51_probs: &CudaSlice<f32>,
    atom_values: &CudaSlice<f32>,
    iqn_quantiles: &CudaSlice<f32>,
    n_atoms: i32,
) -> u8 {
    let kernel = &kernels(stream).argmax_eq;

    // Safety: CUDA context is active on this thread.
    let buffer = unsafe { MappedI32Buffer::new(1) }
        .expect("alloc MappedI32Buffer for argmax");
    let dev_ptr = buffer.dev_ptr;

    unsafe {
        stream.launch_builder(kernel)
            .arg(c51_probs)
            .arg(atom_values)
            .arg(iqn_quantiles)
            .arg(&(B0_SIZE as i32))
            .arg(&n_atoms)
            .arg(&dev_ptr)
            .launch(LaunchConfig {
                grid_dim: (1, 1, 1),
                block_dim: (1, 1, 1),
                shared_mem_bytes: 0,
            })
            .expect("launch argmax_eq_test");
    }

    stream.synchronize().expect("synchronize after argmax_eq launch");

    let host = buffer.read_all();
    let raw = host.first().copied().expect("argmax mapped buf empty");
    assert!(
        (0..(B0_SIZE as i32)).contains(&raw),
        "argmax kernel wrote direction index {} outside [0, {})",
        raw,
        B0_SIZE,
    );
    raw as u8
}

/// Helper: build CUDA stream + context.
fn make_test_stream() -> Arc<CudaStream> {
    let ctx = CudaContext::new(0).expect("CUDA context — is GPU available?");
    ctx.default_stream()
}

/// Helper: upload a flat Vec<f32> to GPU.
fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> CudaSlice<f32> {
    stream.clone_htod(host).expect("htod upload")
}

// Test cases are added by Tasks 3-8 of plan-A-phase-0.md.

/// Phase 0 Test 0.A — Bias mechanism reproduces under controlled C51 atoms.
///
/// Construct synthetic C51 atom probabilities representing the val-Flat-collapse:
///   Flat dist:  δ(v=0)              — single atom mass at v=0
///   Long/Short dist: Gaussian-shaped mass centered at -0.001 (post-tx_cost)
///   Hold dist:  δ(v=0)              — same as Flat (no-op semantics)
///
/// Verify:
///   - argmax(E[Q]) picks Flat (or Hold) ~100% of trials  → BIAS exists
///   - Thompson sampling picks Long+Short ≥ 30% of trials → fix would work
#[test]
#[ignore]
fn test_0a_bias_reproduces_argmax_vs_thompson() {
    let stream = make_test_stream();
    let n_atoms = 21;
    let atoms_host: Vec<f32> = (0..n_atoms)
        .map(|i| -1.0 + (i as f32) * (2.0 / (n_atoms - 1) as f32))
        .collect();
    let atom_values = upload(&stream, &atoms_host);

    // Build per-direction C51 probs:
    //   index 0=Short, 1=Hold, 2=Long, 3=Flat
    let mut c51_probs_host = vec![0.0_f32; B0_SIZE * n_atoms as usize];
    let center_atom = (n_atoms / 2) as usize;  // atom at v=0

    // Flat (d=3): δ(v=0)
    c51_probs_host[3 * n_atoms as usize + center_atom] = 1.0;
    // Hold (d=1): δ(v=0) too
    c51_probs_host[n_atoms as usize + center_atom] = 1.0;
    // Long (d=2): Gaussian-like spread around v=-0.001 (one atom left of center)
    for a in 0..n_atoms as usize {
        let v = atoms_host[a];
        let mean = -0.001;
        let sigma = 0.3;  // wide
        let p = (-((v - mean).powi(2)) / (2.0 * sigma * sigma)).exp();
        c51_probs_host[2 * n_atoms as usize + a] = p;
    }
    let long_sum: f32 = c51_probs_host[2 * n_atoms as usize..3 * n_atoms as usize].iter().sum();
    for a in 2 * n_atoms as usize..3 * n_atoms as usize {
        c51_probs_host[a] /= long_sum;
    }
    // Short (d=0): same shape as Long, mirrored
    for a in 0..n_atoms as usize {
        c51_probs_host[a] = c51_probs_host[2 * n_atoms as usize + a];
    }
    let c51_probs = upload(&stream, &c51_probs_host);

    // IQN quantiles: tight for Flat/Hold, wide for Long/Short
    let mut iqn_host = vec![0.0_f32; B0_SIZE * N_IQN_QUANTILES];
    // Flat/Hold: all quantiles = 0
    // Long/Short: quantiles spread (Gaussian-like)
    let long_quantiles = [-0.45, -0.20, -0.001, 0.20, 0.45];
    for q_idx in 0..N_IQN_QUANTILES {
        iqn_host[q_idx] = long_quantiles[q_idx];  // Short
        iqn_host[2 * N_IQN_QUANTILES + q_idx] = long_quantiles[q_idx];  // Long
    }
    let iqn_quantiles = upload(&stream, &iqn_host);

    // 1. argmax should pick Flat (or Hold) deterministically
    let argmax_pick = launch_argmax_eq_pinned(&stream, &c51_probs, &atom_values, &iqn_quantiles, n_atoms);
    assert!(argmax_pick == 1 || argmax_pick == 3,
        "argmax expected Flat (3) or Hold (1), got {}", argmax_pick);

    // 2. Thompson over 10 000 seeds in ONE batched launch — host counts
    //    on the returned Vec<u8>. Per-sample-arrays + host-reduce; no
    //    atomicAdd, no per-seed dtoh.
    let n_seeds = 10_000_usize;
    let picks = launch_thompson_direction_batched(
        &stream, &c51_probs, &atom_values, &iqn_quantiles, n_atoms, 0, n_seeds,
    );
    let mut counts = [0_u32; B0_SIZE];
    for d in &picks {
        counts[*d as usize] += 1;
    }
    let active = counts[0] + counts[2];  // Short + Long
    assert!(active >= 3000,
        "Thompson active_frac too low: counts={:?}, active={}", counts, active);
}

/// Phase 0 Test 0.B — Thompson C51 inverse-CDF sampling matches input distribution.
///
/// Build a single-direction (d=0) C51 distribution with known probabilities
/// p = [0.1, 0.7, 0.2] over atoms [-1, 0, +1]. Other directions get a
/// δ(v=0) distribution so their C51 sample is always exactly 0. IQN
/// quantiles are all 0 for every direction, so the combined sample is
/// just the C51 atom value.
///
/// What this test measures: P(d=0 is chosen) over many seeds. The kernel's
/// tie-break is strict-> with `best_d=0` initialised, so d=0 wins whenever
/// its C51 sample is ≥ 0 (i.e. atom v=0 OR atom v=+1 was drawn). Other
/// directions always sample exactly 0, so d=0 wins iff `s_c51_d0 >= 0`.
///
///     P(d=0) = P(atom = v=0) + P(atom = v=+1) = p[1] + p[2] = 0.7 + 0.2 = 0.9
///
/// This still verifies inverse-CDF correctness end-to-end: a wrong CDF
/// would shift the v=-1 mass into the [0, +1] region (or vice versa) and
/// move P(d=0) measurably away from 0.9.
#[test]
#[ignore]
fn test_0b_thompson_c51_distribution_matches_input() {
    let stream = make_test_stream();
    let n_atoms = 3;
    let atoms_host = vec![-1.0_f32, 0.0, 1.0];
    let atom_values = upload(&stream, &atoms_host);

    // d=0 carries the distribution under test
    let mut c51_probs_host = vec![0.0_f32; B0_SIZE * n_atoms as usize];
    c51_probs_host[0] = 0.1;
    c51_probs_host[1] = 0.7;
    c51_probs_host[2] = 0.2;
    // Other directions: δ(v=0) so they always sample exactly 0
    for d in 1..B0_SIZE {
        for a in 0..n_atoms as usize {
            c51_probs_host[d * n_atoms as usize + a] = if a == 1 { 1.0 } else { 0.0 };
        }
    }
    let c51_probs = upload(&stream, &c51_probs_host);

    // IQN: all quantiles 0 → IQN sample is always 0 for every direction
    let iqn_host = vec![0.0_f32; B0_SIZE * N_IQN_QUANTILES];
    let iqn_quantiles = upload(&stream, &iqn_host);

    // Single batched launch: 100 000 seeds in one kernel, output read
    // via mapped pinned memory. d=0 wins iff its C51 sample ≥ 0.
    let n_seeds = 100_000_usize;
    let picks = launch_thompson_direction_batched(
        &stream, &c51_probs, &atom_values, &iqn_quantiles, n_atoms as i32, 0, n_seeds,
    );
    let total_d0 = picks.iter().filter(|&&d| d == 0).count() as u32;

    // P(d=0) = p[v=0] + p[v=+1] = 0.7 + 0.2 = 0.9
    let p_d0 = total_d0 as f32 / n_seeds as f32;
    assert!((p_d0 - 0.9).abs() < 0.01,
        "P(d=0) expected ~0.9 (= P(s>=0)), got {} (total_d0={})", p_d0, total_d0);
}

/// Phase 0 Test 0.C — Thompson IQN quantile interpolation symmetry check.
///
/// Setup: d=0 has standard-normal-shaped IQN quantiles
/// `[-1.645, -0.674, 0, 0.674, 1.645]` at the fixed τ-anchors
/// {0.05, 0.25, 0.50, 0.75, 0.95}; other directions have all-zero
/// quantiles. C51 is δ(v=0) for ALL directions, so every direction's
/// C51 sample is exactly 0.
///
/// Combined sample per direction:
///   d=0 → 0 + s_iqn = s_iqn         (s_iqn drawn via uniform-τ interp)
///   d=1,2,3 → 0 + 0 = 0
///
/// Kernel `best_d=0` is initialised with strict-`>` tie-break, so:
///   - d=0 wins whenever `s_iqn ≥ 0` (since 0 > 0 is false, d=1 doesn't override)
///   - By symmetry of the standard-normal IQN, P(s_iqn ≥ 0) = 0.5
///   - The exact-zero tie (u = 0.5 hitting τ=0.5 anchor) is ~6e-8 with
///     the LCG RNG, so it doesn't move the mean measurably.
///
/// A wrong interpolation (e.g. dropping a τ-anchor, swapping bounds, or
/// non-linear blend) would break the symmetry and shift P(d=0) off 0.5.
#[test]
#[ignore]
fn test_0c_thompson_iqn_quantile_interp() {
    let stream = make_test_stream();
    let n_atoms: i32 = 3;
    let atoms_host = vec![-1.0_f32, 0.0, 1.0];
    let atom_values = upload(&stream, &atoms_host);

    // C51: all directions deterministic at v=0
    let mut c51_probs_host = vec![0.0_f32; B0_SIZE * n_atoms as usize];
    for d in 0..B0_SIZE {
        c51_probs_host[d * n_atoms as usize + 1] = 1.0;
    }
    let c51_probs = upload(&stream, &c51_probs_host);

    // IQN: d=0 has standard-normal quantiles; others all-zero
    let mut iqn_host = vec![0.0_f32; B0_SIZE * N_IQN_QUANTILES];
    let normal_quantiles = [-1.645_f32, -0.674, 0.0, 0.674, 1.645];
    for i in 0..N_IQN_QUANTILES {
        iqn_host[i] = normal_quantiles[i];
    }
    let iqn_quantiles = upload(&stream, &iqn_host);

    // Single batched launch: 100 000 seeds in one kernel.
    let n_seeds = 100_000_usize;
    let picks = launch_thompson_direction_batched(
        &stream, &c51_probs, &atom_values, &iqn_quantiles, n_atoms, 0, n_seeds,
    );
    let win_d0 = picks.iter().filter(|&&d| d == 0).count() as u32;

    let p_d0 = win_d0 as f32 / n_seeds as f32;
    // Should be near 0.5 (symmetric standard-normal IQN, strict-> tie-break).
    assert!((p_d0 - 0.5).abs() < 0.02,
        "P(d=0) expected ~0.5, got {} (deviation > 2% suggests IQN interp wrong)", p_d0);
}

/// Phase 0 Test 0.D — Thompson reverses bias on synthetic failure mode.
///
/// Construct directions with: E_flat=0, σ_flat=0, E_long=-0.001, σ_long=0.05.
/// This is the realistic synthetic failure mode — narrower σ than Test 0.A
/// (σ=0.3) so it stresses whether Thompson exploration is sufficient at
/// the σ scale we actually observe in trained checkpoints.
///
/// Confirm:
///   - argmax(E[Q]) deterministically picks Flat (or Hold)
///   - Thompson sampling picks Long with probability ≥ 0.20 over 10k seeds
#[test]
#[ignore]
fn test_0d_thompson_reverses_bias() {
    let stream = make_test_stream();
    let n_atoms: i32 = 21;
    let atoms_host: Vec<f32> = (0..n_atoms)
        .map(|i| -0.5 + (i as f32) * (1.0 / (n_atoms - 1) as f32))
        .collect();
    let atom_values = upload(&stream, &atoms_host);

    // C51: Flat = δ(0); Long = Gaussian centered at -0.001 with σ=0.05
    let mut c51_probs_host = vec![0.0_f32; B0_SIZE * n_atoms as usize];
    let center_atom = (n_atoms / 2) as usize;
    // Flat (d=3) and Hold (d=1): δ(0)
    c51_probs_host[3 * n_atoms as usize + center_atom] = 1.0;
    c51_probs_host[n_atoms as usize + center_atom] = 1.0;
    // Long (d=2): Gaussian centered at -0.001 with σ=0.05 (tight)
    for a in 0..n_atoms as usize {
        let v = atoms_host[a];
        let mean = -0.001;
        let sigma = 0.05;
        let p = (-((v - mean).powi(2)) / (2.0 * sigma * sigma)).exp();
        c51_probs_host[2 * n_atoms as usize + a] = p;
    }
    let long_sum: f32 = c51_probs_host[2 * n_atoms as usize..3 * n_atoms as usize].iter().sum();
    for a in 2 * n_atoms as usize..3 * n_atoms as usize {
        c51_probs_host[a] /= long_sum;
    }
    // Short (d=0): same as Long
    for a in 0..n_atoms as usize {
        c51_probs_host[a] = c51_probs_host[2 * n_atoms as usize + a];
    }
    let c51_probs = upload(&stream, &c51_probs_host);

    // IQN: Flat/Hold all 0; Long/Short have ±0.05σ quantile spread
    let mut iqn_host = vec![0.0_f32; B0_SIZE * N_IQN_QUANTILES];
    let long_q = [-0.082_f32, -0.034, -0.001, 0.032, 0.080];  // mean=-0.001, σ≈0.05
    for i in 0..N_IQN_QUANTILES {
        iqn_host[i] = long_q[i];
        iqn_host[2 * N_IQN_QUANTILES + i] = long_q[i];
    }
    let iqn_quantiles = upload(&stream, &iqn_host);

    // 1. argmax picks Flat (or Hold)
    let argmax_pick = launch_argmax_eq_pinned(&stream, &c51_probs, &atom_values, &iqn_quantiles, n_atoms);
    assert!(argmax_pick == 1 || argmax_pick == 3,
        "argmax expected Flat/Hold, got {}", argmax_pick);

    // 2. Thompson: P(Long) ≥ 0.20 over 10 000 seeds in ONE batched launch.
    //    Host-side counts on the returned Vec<u8>; no per-seed dtoh.
    let n_seeds = 10_000_usize;
    let picks = launch_thompson_direction_batched(
        &stream, &c51_probs, &atom_values, &iqn_quantiles, n_atoms, 0, n_seeds,
    );
    let long_count: u32 = picks.iter().filter(|&&d| d == 2).count() as u32;
    let total_active: u32 = picks.iter().filter(|&&d| d == 0 || d == 2).count() as u32;
    let p_long = long_count as f32 / n_seeds as f32;
    let p_active = total_active as f32 / n_seeds as f32;
    println!("Test 0.D: argmax_pick={}, P(Long)={:.5}, P(active)={:.5}", argmax_pick, p_long, p_active);
    assert!(p_long >= 0.20,
        "Thompson P(Long) expected ≥ 0.20, got {} (active total {})", p_long, p_active);
}

/// Phase 0 Test 0.E — Synthetic edge discovery via Q-learning + Thompson.
///
/// Tests an ALGORITHMIC PROPERTY of Thompson exploration (does it discover
/// edge if edge exists?), not a kernel correctness property. CPU is appropriate.
///
/// Setup: 1-state 2-action bandit. Action Long has KNOWN +0.005 edge after
/// tx_cost. Action Flat returns exactly 0. Run Q-learning for 100 iterations
/// with two action selectors:
///   1. Argmax-only: always pick action with highest Q. Both Q[Long] and
///      Q[Flat] start at E=0; strict-`>` tie-break consistently picks Flat
///      → Long is never explored, Q[Long] never updates, the +0.005 edge is
///      never discovered. Assertion: e_long ≤ e_flat + 0.001 (vacuously true
///      at exactly e_long=e_flat=0).
///   2. Thompson: sample s_long ~ p_long (has σ from spread), s_flat ~ p_flat
///      (= δ(0), always 0). P(s_long > 0) = p_long[3] + p_long[4] = 0.30
///      → Long is selected ~30% of iterations, each update pushes p_long
///      mean toward the +0.005 reward target's nearest atom (+0.05).
///      Mean drifts above zero by iter 100 → e_long > e_flat = 0.
///
/// Initial distributions (revised from plan A draft, option 3 — production-realistic):
///   p_long initial = [0.10, 0.20, 0.40, 0.20, 0.10]   (uniform-ish, E=0, has σ;
///                                                      mimics freshly initialized C51 head)
///   p_flat initial = [0, 0, 1, 0, 0]                  (δ(v=0), tx_cost-bare Flat
///                                                      deterministically returns 0)
///
/// Stop condition: if Thompson e_long ≤ e_flat with this setup, the hypothesis
/// is genuinely wrong and reward shaping must change before proceeding to Phase 2.
#[test]
fn test_0e_synthetic_edge_discovery_cpu() {
    use rand::{Rng, SeedableRng};
    use rand::rngs::StdRng;
    use rand_distr::{Normal, Distribution};

    fn sample_categorical(p: &[f32], v: &[f32], rng: &mut StdRng) -> f32 {
        let u: f32 = rng.gen();
        let mut cum = 0.0_f32;
        for i in 0..p.len() {
            cum += p[i];
            if u < cum { return v[i]; }
        }
        v[p.len() - 1]
    }

    fn run_q_learning(thompson: bool, seed: u64) -> (f32, f32) {
        let mut rng = StdRng::seed_from_u64(seed);
        // C51-like distribution: 5 atoms, [-0.1, -0.05, 0, 0.05, 0.1]
        let atoms = [-0.1_f32, -0.05, 0.0, 0.05, 0.1];
        // Long: uniform-ish, E[Q]=0, has variance. Mimics a freshly initialized
        // C51 head before training: spread mass, mean near zero. P(sample > 0) ≈ 0.30.
        let mut p_long = [0.10_f32, 0.20, 0.40, 0.20, 0.10];
        // Flat = δ(v=0), deterministic. Production reality: tx_cost-bare flat returns 0.
        let mut p_flat = [0.0_f32, 0.0, 1.0, 0.0, 0.0];
        // True Long edge: Normal(+0.005, 0.03). Known positive edge ≈ 1/2 of σ.
        let true_long_dist = Normal::new(0.005_f32, 0.03)
            .expect("Normal(+0.005, 0.03) parameters valid");
        let lr = 0.05_f32;
        const N_ITERS: usize = 100;

        for _iter in 0..N_ITERS {
            // Action selection
            let action = if thompson {
                let s_long: f32 = sample_categorical(&p_long, &atoms, &mut rng);
                let s_flat: f32 = sample_categorical(&p_flat, &atoms, &mut rng);
                if s_long > s_flat { 0 } else { 1 }  // 0=long, 1=flat
            } else {
                let e_long: f32 = (0..5).map(|i| p_long[i] * atoms[i]).sum();
                let e_flat: f32 = (0..5).map(|i| p_flat[i] * atoms[i]).sum();
                if e_long > e_flat { 0 } else { 1 }
            };

            // Get reward — Long gets the synthetic edge, Flat is deterministic 0.
            let reward = if action == 0 {
                true_long_dist.sample(&mut rng)
            } else {
                0.0_f32
            };

            // Q-learning update: project reward onto atoms (nearest-bin).
            let target_atom = atoms.iter()
                .enumerate()
                .min_by(|(_, a), (_, b)| {
                    (*a - reward).abs()
                        .partial_cmp(&(*b - reward).abs())
                        .expect("atom distance is finite (no NaN)")
                })
                .map(|(i, _)| i)
                .expect("atoms slice is non-empty");
            let p = if action == 0 { &mut p_long } else { &mut p_flat };
            for i in 0..5 {
                let target = if i == target_atom { 1.0 } else { 0.0 };
                p[i] = (1.0 - lr) * p[i] + lr * target;
            }
            // Renormalize (defensive — the convex blend above already preserves sum=1
            // when both inputs sum to 1, but keep this so a future change to the
            // update rule can't silently drift the total mass).
            let s: f32 = p.iter().sum();
            for i in 0..5 { p[i] /= s; }
        }

        let e_long: f32 = (0..5).map(|i| p_long[i] * atoms[i]).sum();
        let e_flat: f32 = (0..5).map(|i| p_flat[i] * atoms[i]).sum();
        (e_long, e_flat)
    }

    // Argmax-only: never explores Long → Q[Long] stays at 0 = Q[Flat].
    let (e_long_argmax, e_flat_argmax) = run_q_learning(false, 42);
    println!(
        "Test 0.E argmax: e_long={:.6}, e_flat={:.6}",
        e_long_argmax, e_flat_argmax
    );
    assert!(e_long_argmax <= e_flat_argmax + 0.001,
        "argmax: Q[Long]={} should not exceed Q[Flat]={} (no exploration)",
        e_long_argmax, e_flat_argmax);

    // Thompson: ~30% Long-selections → mean drifts toward +0.005 → crosses 0.
    let (e_long_thomp, e_flat_thomp) = run_q_learning(true, 42);
    println!(
        "Test 0.E thompson: e_long={:.6}, e_flat={:.6}",
        e_long_thomp, e_flat_thomp
    );
    assert!(e_long_thomp > e_flat_thomp,
        "Thompson: Q[Long]={} expected to exceed Q[Flat]={} after 100 iters",
        e_long_thomp, e_flat_thomp);
}

/// Phase 0 Test 0.F — Real checkpoint extraction (GPU integration).
///
/// Loads a 4-direction modern-architecture DQN checkpoint, runs a forward
/// pass on a synthetic state vector, and extracts the per-direction C51
/// atom probabilities from the direction branch (`advantage_log_probs[0]`).
/// Then runs the diagnostic σ kernel + argmax + Thompson kernels and
/// asserts the bias-hypothesis structure on a real model.
///
/// ## C51-only scope (IQN inputs zeroed)
///
/// IQN quantiles live exclusively inside the GPU `iqn_dual_head_kernel`
/// (training-side); they're not exposed through any read-only host surface.
/// We supply zero-IQN inputs to the argmax / Thompson kernels — they accept
/// arbitrary `iqn_quantiles` and combine via simple addition, so zero IQN
/// means the joint = C51-only, which is the right semantic for testing the
/// C51-side bias hypothesis on the C51 atoms emitted by the branching head.
///
/// A future Tier-A version moves the test entirely onto the GPU integration
/// path: read the joint distribution directly from the dual-head's
/// GPU-resident output (no CPU forward — CPU is read-only per the GPU-drives
/// principle), then run the diagnostic σ / argmax / Thompson kernels against
/// the same buffers production training reads.
///
/// ## Hypothesis structure (assertions)
///
///   - `σ_C51[FLAT] < 0.01 × σ_C51[LONG]` — Flat C51 distribution collapses
///   - `argmax(E[Q_C51])` picks FLAT (3) or HOLD (1)         — argmax bias reproduces
///   - Thompson over 1 000 trials produces P(LONG+SHORT) ≥ 0.20 — Thompson reverses
///
/// ## Stop condition
///
/// If any structural assertion fails, the bias hypothesis does not match
/// this specific checkpoint at this specific state. Either:
///   - this checkpoint isn't actually converged into the bias state
///     (smoke run, only 5 epochs/fold × 3 folds, val_loss=2.30, no Sharpe),
///   - or distributional learning has collapsed differently here.
///
/// In either case: the recorded measured values determine next steps
/// (re-run on a fully-converged checkpoint vs revisit hypothesis).
#[test]
#[ignore]
fn test_0f_real_checkpoint_extraction() {
    use std::path::PathBuf;
    use ml_core::cuda_autograd::GpuTensor;
    use ml_dqn::{DQN, DQNConfig};

    // ── Locate checkpoint ─────────────────────────────────────────────
    //
    // Search for a modern (4-direction, 4-branch) checkpoint. The plan's
    // step 1 runs `find /mnt/training-data /workspace/output ...`; we mirror
    // that intent locally by checking the smoke-test output path first
    // (Apr 2026 4-direction architecture) and the production output dir
    // second (Apr 9 2026 — 3-direction stale architecture, would fail
    // the `branch_0_out` shape match).
    let workspace_root = PathBuf::from(env!("CARGO_MANIFEST_DIR"))
        .ancestors()
        .find(|p| p.join("Cargo.toml").exists() && p.join("crates").exists())
        .expect("workspace root not found")
        .to_path_buf();

    let candidate_paths = [
        workspace_root.join("target/smoke_multi_fold_output/dqn_fold0_best.safetensors"),
        workspace_root.join("test_data/phase0_converged.safetensors"),
        workspace_root.join("ml/trained_models/dqn_fold0_best.safetensors"),
    ];

    let checkpoint_path = candidate_paths
        .iter()
        .find(|p| p.exists())
        .unwrap_or_else(|| panic!(
            "Test 0.F: no checkpoint found. Tried: {:?}. Run \
             `./scripts/argo-train.sh dqn --baseline --epochs 60 --folds 1 \
             --gpu-pool ci-training-l40s --tag phase0-converged-checkpoint` \
             and `kubectl cp` the result to test_data/phase0_converged.safetensors.",
            candidate_paths,
        ));

    println!("Test 0.F: using checkpoint {}", checkpoint_path.display());

    // ── Build matching DQNConfig from checkpoint metadata ─────────────
    //
    // `DQNConfig::from_safetensors_file` reads the architecture-critical
    // fields from the file's `__metadata__` (state_dim, num_actions,
    // hidden_dims, num_atoms, dueling_hidden_dim, num_order_types,
    // num_urgency_levels, gamma) and fills the rest from
    // `DQNConfig::default()`. Non-architecture fields (LR, epsilon,
    // buffer capacity) don't affect inference. This also enforces
    // `state_dim == ml_core::state_layout::STATE_DIM` (returns Err if
    // the checkpoint is from an old STATE_DIM commit).
    let mut config = DQNConfig::from_safetensors_file(checkpoint_path)
        .unwrap_or_else(|e| panic!(
            "Test 0.F: failed to derive DQNConfig from checkpoint: {e}",
        ));
    // The smoke trainer leaves branch_hidden_dim == dueling_hidden_dim
    // (both 128). `from_safetensors_file` doesn't preserve
    // branch_hidden_dim explicitly (no metadata key), so set it from
    // dueling_hidden_dim to match the saved branch FC shape [128, 64].
    config.branch_hidden_dim = config.dueling_hidden_dim;

    // Sanity gate before constructing the GPU model: the test only makes
    // sense when the checkpoint is the modern 4-direction layout.
    assert_eq!(
        config.num_actions, B0_SIZE,
        "Test 0.F: checkpoint num_actions={} != B0_SIZE={} (4-direction \
         layout). Apr-9 3-direction checkpoints cannot be tested here.",
        config.num_actions, B0_SIZE,
    );
    let ck_num_atoms = config.num_atoms;
    let ck_v_min = config.v_min;
    let ck_v_max = config.v_max;
    println!(
        "Test 0.F: matched config — num_actions={}, num_atoms={}, \
         hidden_dims={:?}, dueling_hidden_dim={}, v_min={}, v_max={}",
        config.num_actions, ck_num_atoms, config.hidden_dims,
        config.dueling_hidden_dim, ck_v_min, ck_v_max,
    );

    // ── Construct DQN + load weights ──────────────────────────────────
    let mut dqn = DQN::new(config).expect("DQN::new with checkpoint-derived config");
    let path_str = checkpoint_path.to_str().expect("checkpoint path utf8");
    dqn.load_from_safetensors(path_str)
        .expect("DQN::load_from_safetensors");
    println!("Test 0.F: weights loaded ({} atoms)", ck_num_atoms);

    // ── Forward on synthetic state (zeros) ────────────────────────────
    //
    // The structural assertions don't depend on the specific state — they
    // measure the model's general C51 distributional shape per direction.
    // Use zeros for a representative-input baseline.
    let stream = Arc::clone(dqn.cuda_stream());
    let state_dim = ml_core::state_layout::STATE_DIM;
    let state_host = vec![0.0_f32; state_dim];
    let state = GpuTensor::from_host(&state_host, vec![1, state_dim], &stream)
        .expect("upload zero state");

    let branching = dqn.branching_q_network.as_ref()
        .expect("branching_q_network must exist after load_from_safetensors");
    let branch_out = branching.forward_branches_eval(&state)
        .expect("branching forward eval");

    // direction = branch index 0 — log_probs shape [1, B0_SIZE, num_atoms].
    let log_probs_branches = branch_out.advantage_log_probs.as_ref()
        .expect("advantage_log_probs Some when distributional+branching");
    let dir_log_probs = log_probs_branches.first()
        .expect("at least one branch (direction) exists");
    assert_eq!(
        dir_log_probs.shape(), &[1_usize, B0_SIZE, ck_num_atoms],
        "Direction-branch log-probs shape mismatch: expected [1, {}, {}], got {:?}",
        B0_SIZE, ck_num_atoms, dir_log_probs.shape(),
    );
    // exp(log_probs) → probs in [0, 1] (each row over atoms sums to 1 by
    // construction of log_softmax).
    let dir_probs_tensor = dir_log_probs.exp(&stream).expect("exp(log_probs)");
    let c51_probs_slice = dir_probs_tensor.data();

    // ── Atom values support [num_atoms] ───────────────────────────────
    let delta = (ck_v_max - ck_v_min) / (ck_num_atoms as f32 - 1.0);
    let atom_values_host: Vec<f32> = (0..ck_num_atoms)
        .map(|i| ck_v_min + (i as f32) * delta)
        .collect();
    let atom_values_tensor = GpuTensor::from_host(
        &atom_values_host, vec![ck_num_atoms], &stream,
    ).expect("upload atom values");
    let atom_values_slice = atom_values_tensor.data();

    // ── Tier-B-prime: zero-IQN buffer for argmax/Thompson kernels ────
    //
    // The kernels read [b0_size, N_IQN_QUANTILES]. With zeros, the IQN
    // contribution to E[Q] and to each Thompson sample is exactly 0, so
    // the joint reduces to C51-only — which is the right semantic for an
    // IQN-disabled checkpoint.
    let iqn_zero_host = vec![0.0_f32; B0_SIZE * N_IQN_QUANTILES];
    let iqn_zero_tensor = GpuTensor::from_host(
        &iqn_zero_host, vec![B0_SIZE, N_IQN_QUANTILES], &stream,
    ).expect("upload zero IQN buffer");
    let iqn_zero_slice = iqn_zero_tensor.data();

    // ── Run σ_C51 diagnostic kernel via mapped pinned f32 ─────────────
    let n_atoms_i32 = i32::try_from(ck_num_atoms).expect("num_atoms fits in i32");
    let sigma_kernel = &kernels(&stream).sigma_c51;

    // Safety: CUDA context active on this thread (we resolved via the same stream).
    let sigma_buf = unsafe { MappedF32Buffer::new(B0_SIZE) }
        .expect("alloc MappedF32Buffer for sigma_c51");
    let sigma_dev_ptr = sigma_buf.dev_ptr;

    unsafe {
        stream.launch_builder(sigma_kernel)
            .arg(c51_probs_slice)
            .arg(atom_values_slice)
            .arg(&(B0_SIZE as i32))
            .arg(&n_atoms_i32)
            .arg(&sigma_dev_ptr)
            .launch(LaunchConfig {
                grid_dim: (1, 1, 1),
                block_dim: (B0_SIZE as u32, 1, 1),
                shared_mem_bytes: 0,
            })
            .expect("launch compute_sigma_c51_test");
    }
    stream.synchronize().expect("synchronize after sigma_c51 launch");

    let sigma_c51 = sigma_buf.read_all();
    println!(
        "Test 0.F σ_C51 per direction (Short=0, Hold=1, Long=2, Flat=3): {:?}",
        sigma_c51,
    );

    // Sanity: every σ ≥ 0 and finite (kernel uses fmaxf(var, 0); a NaN
    // would mean the C51 distribution is degenerate).
    for (d, s) in sigma_c51.iter().enumerate() {
        assert!(s.is_finite() && *s >= 0.0,
            "σ_C51[{d}] = {s} is not a finite non-negative value", d = d);
    }

    // ── argmax over E[Q_C51] (joint with zero IQN ⇒ C51-only) ─────────
    let argmax_pick = launch_argmax_eq_pinned(
        &stream, c51_probs_slice, atom_values_slice, iqn_zero_slice, n_atoms_i32,
    );
    println!("Test 0.F argmax pick (0=Short, 1=Hold, 2=Long, 3=Flat): {}", argmax_pick);

    // ── Thompson over 1 000 seeds, batched ────────────────────────────
    let n_seeds = 1_000_usize;
    let picks = launch_thompson_direction_batched(
        &stream, c51_probs_slice, atom_values_slice, iqn_zero_slice,
        n_atoms_i32, /*base_seed=*/ 0, n_seeds,
    );
    let mut counts = [0_u32; B0_SIZE];
    for d in &picks {
        counts[*d as usize] += 1;
    }
    let active = counts[0] + counts[2]; // Short + Long
    let p_active = active as f32 / n_seeds as f32;
    println!(
        "Test 0.F Thompson counts (S/H/L/F): {:?}, P(active)={:.4}",
        counts, p_active,
    );

    // ── Structural assertions ─────────────────────────────────────────
    //
    // Per spec (Test 0.F section): σ_C51[FLAT] < 0.01 × σ_C51[LONG],
    // argmax picks FLAT or HOLD, P(active) ≥ 0.20 over Thompson trials.
    //
    // No threshold relaxation: if the assertions fail, the measured
    // values printed above are the concrete finding for the next agent.
    assert!(
        sigma_c51[3] < 0.01 * sigma_c51[2],
        "σ_C51[FLAT] = {} should be ≪ σ_C51[LONG] = {} (ratio = {:.4} ≥ 0.01)",
        sigma_c51[3], sigma_c51[2], sigma_c51[3] / sigma_c51[2].max(f32::EPSILON),
    );
    assert!(
        argmax_pick == 1 || argmax_pick == 3,
        "argmax expected FLAT (3) or HOLD (1), got direction={} \
         — bias hypothesis does not reproduce on this checkpoint",
        argmax_pick,
    );
    assert!(
        p_active >= 0.20,
        "Thompson P(LONG+SHORT) expected ≥ 0.20, got {} \
         (counts S/H/L/F = {:?}) — Thompson exploration insufficient",
        p_active, counts,
    );
}

// ── Phase 2 Tests 2.A–2.D: Production kernel direct exercise ──────────────
//
// Plan reference:
//   docs/superpowers/plans/2026-04-27-distributional-rl-thompson-plan-C-phase-2.md
//   tasks 6, 7, 8, 9 (lines 371-531).
//
// These tests exercise the PRODUCTION `experience_action_select` kernel
// directly via inline cudarc fixtures rather than the never-built Plan B
// audit fixture (`GpuExperienceCollector::new_for_test()`). Each test
// uploads controlled inputs (q_values, b_logits_dir, per_sample_support)
// to the GPU, launches the production kernel with the Plan C T2 amended
// ABI, and reads back per-sample factored action picks for assertions.
//
// All four tests are `#[ignore]`-gated because they require a CUDA GPU.

const PROD_EXP_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/experience_kernels.cubin"
));

/// Cached production-kernel handle. Loading the cubin and resolving the
/// function on every test invocation would dominate kernel time. Pattern
/// matches `KERNELS` above.
static PROD_ACTION_SELECT: OnceLock<CudaFunction> = OnceLock::new();

fn prod_action_select(stream: &Arc<CudaStream>) -> &'static CudaFunction {
    PROD_ACTION_SELECT.get_or_init(|| {
        let module = stream
            .context()
            .load_cubin(PROD_EXP_CUBIN.to_vec())
            .expect("load experience_kernels cubin");
        module
            .load_function("experience_action_select")
            .expect("load experience_action_select")
    })
}

/// Branch sizes for the 4-direction production action space.
const PROD_B0: i32 = 4;  // direction (Short, Hold, Long, Flat)
const PROD_B1: i32 = 3;  // magnitude (Small, Half, Full)
const PROD_B2: i32 = 3;  // order type
const PROD_B3: i32 = 3;  // urgency

/// Q-stride (b0 + b1 + b2 + b3 = 13 floats per sample).
const PROD_Q_STRIDE: usize = (PROD_B0 + PROD_B1 + PROD_B2 + PROD_B3) as usize;

/// Decode dir_idx from a factored action index using PROD_B[1..3].
fn decode_dir(action_idx: i32) -> i32 {
    action_idx / (PROD_B1 * PROD_B2 * PROD_B3)
}

/// Decode mag_idx from a factored action index.
fn decode_mag(action_idx: i32) -> i32 {
    (action_idx / (PROD_B2 * PROD_B3)) % PROD_B1
}

/// Decode ord_idx from a factored action index.
fn decode_ord(action_idx: i32) -> i32 {
    (action_idx / PROD_B3) % PROD_B2
}

/// Decode urg_idx from a factored action index.
fn decode_urg(action_idx: i32) -> i32 {
    action_idx % PROD_B3
}

/// Fixture for launching the production `experience_action_select` kernel
/// with controlled inputs. Each call to `launch` produces fresh actions in
/// `actions_buf`; reads back via host download for assertions.
struct ProdActionSelectFixture {
    stream: Arc<CudaStream>,
    n_atoms: i32,
    batch: usize,
    q_values: CudaSlice<f32>,
    b_logits_dir: CudaSlice<f32>,
    per_sample_support: CudaSlice<f32>,
    actions_buf: CudaSlice<i32>,
    q_gaps_buf: CudaSlice<f32>,
    intent_buf: CudaSlice<i32>,
    conv_buf: CudaSlice<f32>,
    mag_conv_buf: CudaSlice<f32>,
}

impl ProdActionSelectFixture {
    /// Build a new fixture. Inputs:
    ///   - `q_values_host`        [batch * PROD_Q_STRIDE]
    ///   - `b_logits_host`        [batch * PROD_B0 * n_atoms] raw C51 logits
    ///   - `support_host`         [batch * PROD_B0 * 3] stride-12 (v_min, v_max, delta_z)
    fn new(
        stream: Arc<CudaStream>,
        batch: usize,
        n_atoms: i32,
        q_values_host: &[f32],
        b_logits_host: &[f32],
        support_host: &[f32],
    ) -> Self {
        assert_eq!(q_values_host.len(), batch * PROD_Q_STRIDE);
        assert_eq!(b_logits_host.len(), batch * (PROD_B0 as usize) * (n_atoms as usize));
        assert_eq!(support_host.len(), batch * (PROD_B0 as usize) * 3);

        let q_values = stream.clone_htod(q_values_host).expect("upload q_values");
        let b_logits_dir = stream.clone_htod(b_logits_host).expect("upload b_logits");
        let per_sample_support = stream.clone_htod(support_host).expect("upload support");
        let actions_buf = stream.alloc_zeros::<i32>(batch).expect("alloc actions");
        let q_gaps_buf = stream.alloc_zeros::<f32>(batch).expect("alloc q_gaps");
        let intent_buf = stream.alloc_zeros::<i32>(batch).expect("alloc intent");
        let conv_buf = stream.alloc_zeros::<f32>(batch).expect("alloc conviction");
        let mag_conv_buf = stream.alloc_zeros::<f32>(batch).expect("alloc mag_conv");

        Self {
            stream,
            n_atoms,
            batch,
            q_values,
            b_logits_dir,
            per_sample_support,
            actions_buf,
            q_gaps_buf,
            intent_buf,
            conv_buf,
            mag_conv_buf,
        }
    }

    /// Launch the kernel and return a host-side copy of `actions_buf`.
    /// `eval_mode` toggles eps_start = eps_end = 0 (eval) vs > 0 (training).
    fn launch(&mut self, eval_mode: bool, timestep: i32) -> Vec<i32> {
        let kernel = prod_action_select(&self.stream);
        let blocks = (self.batch as u32).div_ceil(256);
        let cfg = LaunchConfig {
            grid_dim: (blocks.max(1), 1, 1),
            block_dim: (256, 1, 1),
            shared_mem_bytes: 0,
        };
        let n_i32: i32 = self.batch as i32;
        let null_ptr: u64 = 0;
        let (eps_start, eps_end) = if eval_mode {
            (0.0_f32, 0.0_f32)
        } else {
            // Training: cosine schedule with eps_start=0.5 ... eps_end=0.05.
            // current_epoch=0 → epsilon = eps_end + (eps_start-eps_end)*1 = 0.5.
            (0.5_f32, 0.05_f32)
        };

        unsafe {
            self.stream
                .launch_builder(kernel)
                .arg(&self.q_values)
                .arg(&mut self.actions_buf)
                .arg(&mut self.q_gaps_buf)
                .arg(&eps_start)
                .arg(&eps_end)
                .arg(&0_i32)              // current_epoch
                .arg(&1_i32)              // total_epochs
                .arg(&n_i32)
                .arg(&PROD_B0)
                .arg(&PROD_B1)
                .arg(&PROD_B2)
                .arg(&PROD_B3)
                .arg(&0.0_f32)            // q_gap_threshold
                .arg(&null_ptr)           // portfolio_states (NULL)
                .arg(&0_i32)              // min_hold_bars
                .arg(&0.0_f32)            // max_position
                .arg(&1.0_f32)            // eps_exp_mult
                .arg(&1.0_f32)            // eps_ord_mult
                .arg(&1.0_f32)            // eps_urg_mult
                .arg(&timestep)
                .arg(&null_ptr)           // per_sample_epsilon (NULL → cosine)
                .arg(&null_ptr)           // isv_signals_ptr (NULL → static)
                .arg(&0_i32)              // contrarian_active
                .arg(&mut self.intent_buf)
                .arg(&mut self.conv_buf)
                .arg(&mut self.mag_conv_buf)
                // Plan C T2 amendment buffers:
                .arg(&self.b_logits_dir)
                .arg(&self.per_sample_support)
                .arg(&null_ptr)           // atom_positions = NULL → linear support
                .arg(&self.n_atoms)
                .launch(cfg)
                .expect("launch experience_action_select");
        }
        self.stream.synchronize().expect("sync after action_select");
        self.stream
            .clone_dtoh(&self.actions_buf)
            .expect("readback actions")
    }
}

/// Helper: upload a single per-direction "peaked-at-target-v" C51 logit
/// pattern across a whole batch. Logit is `peak_logit` at the atom whose
/// linear-support value is closest to `target_v[d]`, 0 elsewhere.
fn fill_peaked_logits(
    batch: usize,
    n_atoms: i32,
    v_min: f32,
    delta_z: f32,
    target_v: &[f32; 4],
    peak_logit: f32,
) -> Vec<f32> {
    let mut out = vec![0.0_f32; batch * (PROD_B0 as usize) * (n_atoms as usize)];
    for i in 0..batch {
        for (d, tv) in target_v.iter().enumerate() {
            let peak_a = (((*tv - v_min) / delta_z).round() as usize)
                .min((n_atoms - 1) as usize);
            let base =
                i * (PROD_B0 as usize) * (n_atoms as usize) + d * (n_atoms as usize);
            out[base + peak_a] = peak_logit;
        }
    }
    out
}

/// Helper: build a per_sample_support tile with uniform linear support.
fn fill_linear_support(batch: usize, v_min: f32, v_max: f32, delta_z: f32) -> Vec<f32> {
    let mut out = vec![0.0_f32; batch * (PROD_B0 as usize) * 3];
    for i in 0..batch {
        for d in 0..(PROD_B0 as usize) {
            let base = i * (PROD_B0 as usize) * 3 + d * 3;
            out[base] = v_min;
            out[base + 1] = v_max;
            out[base + 2] = delta_z;
        }
    }
    out
}

/// Phase 2 Test 2.A — Production action_select kernel mode behavior.
///
/// Verifies the production kernel's two operating modes under the failure
/// pattern from Phase 0 Test 0.D (Flat/Hold collapse via E[Q] argmax bias):
///
///   * Eval mode (eps_start=eps_end=0): 10 calls produce IDENTICAL `dir_idx`
///     per sample (deterministic argmax E[Q]); the argmax is Flat or Hold.
///
///   * Training mode (eps_start=0.5): 10000 samples produce diverse `dir_idx`
///     with count(d ∈ {Long, Short}) ≥ 30% — Thompson sampling explores the
///     directional alternatives despite their lower E[Q].
///
/// Test setup mirrors Phase 0 Test 0.D's failure-mode distribution: Flat
/// = δ(v=0), Hold = δ(v=0), Long/Short = Gaussian centred at v=-0.001 with
/// σ=0.05 (tight, post-tx-cost). With argmax E[Q] this produces
/// deterministic Flat/Hold; with Thompson the spread of Long/Short atoms
/// gives them >0 probability mass at positive v, so Thompson samples them.
#[test]
#[ignore]
fn test_2a_production_kernel_modes() {
    let stream = make_test_stream();
    let n_atoms: i32 = 21;
    let v_min = -0.5_f32;
    let v_max = 0.5_f32;
    let delta_z = (v_max - v_min) / (n_atoms as f32 - 1.0);

    // Build per-direction C51 logits matching Phase 0 Test 0.D failure mode.
    // Convert "probabilities" (Gaussian shape) to logits via log; Flat/Hold
    // get one large logit at the v=0 atom and large-negative elsewhere.
    let mut b_logits_host =
        vec![0.0_f32; (PROD_B0 as usize) * (n_atoms as usize)];
    let center_atom = ((0.0 - v_min) / delta_z).round() as usize;
    // Flat (d=3): δ(v=0) — peaked logit at center
    b_logits_host[3 * n_atoms as usize + center_atom] = 10.0;
    // Hold (d=1): δ(v=0)
    b_logits_host[n_atoms as usize + center_atom] = 10.0;
    // Long (d=2): Gaussian-shaped logits (log of Gaussian density) at v=-0.001 with σ=0.05
    for a in 0..n_atoms as usize {
        let v = v_min + (a as f32) * delta_z;
        let mean = -0.001_f32;
        let sigma = 0.05_f32;
        // Log-Gaussian: -((v-mean)^2)/(2*sigma^2). softmax cancels the constant.
        b_logits_host[2 * n_atoms as usize + a] =
            -((v - mean).powi(2)) / (2.0 * sigma * sigma);
    }
    // Short (d=0): same shape as Long
    for a in 0..n_atoms as usize {
        b_logits_host[a] = b_logits_host[2 * n_atoms as usize + a];
    }

    // Replicate the per-sample tile across the batch.
    let batch: usize = 10_000;
    let mut full_logits = Vec::with_capacity(batch * b_logits_host.len());
    for _ in 0..batch {
        full_logits.extend_from_slice(&b_logits_host);
    }
    let support_host = fill_linear_support(batch, v_min, v_max, delta_z);
    // q_values uniform: direction is decided entirely by Thompson/argmax over
    // C51, mag/ord/urg fall through their Boltzmann floor logic uniformly.
    let q_values_host = vec![0.0_f32; batch * PROD_Q_STRIDE];

    let mut fixture = ProdActionSelectFixture::new(
        stream,
        batch,
        n_atoms,
        &q_values_host,
        &full_logits,
        &support_host,
    );

    // ── Eval mode: 10 calls produce identical dir_idx per sample ─────────
    let eval_runs: Vec<Vec<i32>> = (0..10)
        .map(|t| {
            fixture
                .launch(/*eval_mode=*/ true, /*timestep=*/ t)
                .iter()
                .map(|a| decode_dir(*a))
                .collect()
        })
        .collect();
    for run_idx in 1..eval_runs.len() {
        for sample_idx in 0..batch {
            assert_eq!(
                eval_runs[run_idx][sample_idx], eval_runs[0][sample_idx],
                "Eval mode: dir_idx differs between run 0 and run {} at sample {} \
                 — kernel is non-deterministic in eval mode",
                run_idx, sample_idx,
            );
        }
    }

    // Eval picks must be Hold (1) or Flat (3) — argmax E[Q] over the
    // failure-mode distribution. (Hold/Flat have E[Q]=0, Long/Short have
    // E[Q] slightly negative due to the -0.001 mean.)
    let eval_picks = &eval_runs[0];
    let mut eval_counts = [0_u32; PROD_B0 as usize];
    for d in eval_picks {
        eval_counts[*d as usize] += 1;
    }
    let eval_passive = eval_counts[1] + eval_counts[3]; // Hold + Flat
    println!(
        "Test 2.A eval counts (S/H/L/F): {:?} (passive={}/{})",
        eval_counts, eval_passive, batch,
    );
    assert_eq!(
        eval_passive, batch as u32,
        "Eval mode: every sample must pick Hold or Flat under the failure \
         distribution (argmax E[Q] bias). Got counts {:?}",
        eval_counts,
    );

    // ── Training mode: ≥30% Long+Short across 10 000 samples ─────────────
    let train_actions = fixture.launch(/*eval_mode=*/ false, /*timestep=*/ 0);
    let mut train_counts = [0_u32; PROD_B0 as usize];
    for action in &train_actions {
        let d = decode_dir(*action) as usize;
        train_counts[d] += 1;
    }
    let active = train_counts[0] + train_counts[2]; // Short + Long
    let p_active = active as f32 / batch as f32;
    println!(
        "Test 2.A train counts (S/H/L/F): {:?} (active={:.4})",
        train_counts, p_active,
    );
    assert!(
        p_active >= 0.30,
        "Training mode: P(Long+Short) = {:.4} expected ≥ 0.30 — Thompson \
         sampling should explore directional alternatives. Counts S/H/L/F = {:?}",
        p_active, train_counts,
    );
}

// ── Test 2.B helper: standalone v2 kernel launcher ────────────────────────

/// Cached handle for the v2 standalone direction Thompson kernel.
static STANDALONE_V2: OnceLock<CudaFunction> = OnceLock::new();

fn standalone_thompson_v2(stream: &Arc<CudaStream>) -> &'static CudaFunction {
    STANDALONE_V2.get_or_init(|| {
        let module = stream
            .context()
            .load_cubin(TEST_KERNEL_CUBIN.to_vec())
            .expect("load thompson_test_kernel cubin (v2)");
        module
            .load_function("direction_thompson_v2_test")
            .expect("load direction_thompson_v2_test")
    })
}

/// Launch the standalone v2 direction Thompson test kernel and return the
/// per-seed direction picks. Inputs are SINGLE-tile (the v2 wrapper uses
/// `i=0` internally — see `compute_atom_values_test`):
///   - `b_logits`     [b0_size * n_atoms]
///   - `support_tile` [b0_size * 3]   stride-3 (v_min, v_max, delta_z)
fn launch_standalone_v2(
    stream: &Arc<CudaStream>,
    b_logits_tile: &[f32],
    support_tile: &[f32],
    n_atoms: i32,
    base_seed: u32,
    n_seeds: usize,
) -> Vec<i32> {
    assert_eq!(b_logits_tile.len(), (PROD_B0 as usize) * (n_atoms as usize));
    assert_eq!(support_tile.len(), (PROD_B0 as usize) * 3);
    let kernel = standalone_thompson_v2(stream);
    let logits_buf = stream.clone_htod(b_logits_tile).expect("upload v2 logits");
    let support_buf = stream.clone_htod(support_tile).expect("upload v2 support");

    let buffer = unsafe { MappedI32Buffer::new(n_seeds) }
        .expect("alloc MappedI32Buffer for v2 thompson seeds");

    let n_seeds_i32 = i32::try_from(n_seeds).expect("n_seeds fits in i32");
    let block_dim_x: u32 = 256;
    let grid_dim_x: u32 = (n_seeds as u32).div_ceil(block_dim_x);
    let null_atom_positions: u64 = 0;
    let dev_ptr = buffer.dev_ptr;

    unsafe {
        stream
            .launch_builder(kernel)
            .arg(&logits_buf)
            .arg(&support_buf)
            .arg(&null_atom_positions)
            .arg(&PROD_B0)
            .arg(&n_atoms)
            .arg(&base_seed)
            .arg(&n_seeds_i32)
            .arg(&dev_ptr)
            .launch(LaunchConfig {
                grid_dim: (grid_dim_x.max(1), 1, 1),
                block_dim: (block_dim_x, 1, 1),
                shared_mem_bytes: 0,
            })
            .expect("launch direction_thompson_v2_test");
    }
    stream
        .synchronize()
        .expect("synchronize after standalone v2 launch");
    buffer.read_all()
}

/// Two-sample Kolmogorov-Smirnov-style histogram distance for a discrete
/// 4-bin distribution. Returns `max_d |F1(d) - F2(d)|` where F is the
/// empirical CDF.
fn ks_distance_4bin(hist1: &[u32; 4], hist2: &[u32; 4]) -> f32 {
    let n1: u32 = hist1.iter().sum();
    let n2: u32 = hist2.iter().sum();
    let mut cum1 = 0_u32;
    let mut cum2 = 0_u32;
    let mut max_d = 0.0_f32;
    for i in 0..4 {
        cum1 += hist1[i];
        cum2 += hist2[i];
        let d = ((cum1 as f32 / n1 as f32) - (cum2 as f32 / n2 as f32)).abs();
        if d > max_d {
            max_d = d;
        }
    }
    max_d
}

/// Phase 2 Test 2.B — Production kernel matches Phase 0 standalone wrapper
/// in distribution.
///
/// Bit-equivalent comparison is impossible without modifying the standalone
/// kernel signature: the production direction Thompson uses Philox keyed on
/// `(i, timestep, rng_ctr)`; the standalone v2 uses LCG keyed on
/// `base_seed + seed_idx`. Path (a) (shared uniform array) requires
/// editing `direction_thompson_v2_test` to take a `float* uniforms` instead
/// of generating its own — out of scope for a 1-2-hour dispatch.
///
/// Path (b) — KOLMOGOROV-SMIRNOV-STYLE DISTRIBUTION COMPARISON. Run both
/// kernels on identical inputs over 100 000 samples / seeds. The two output
/// histograms over directions {Short, Hold, Long, Flat} must be statistically
/// indistinguishable: KS statistic ≤ 0.02 (≈ 4 * sqrt(2/100000), well above
/// the noise floor for matching distributions).
///
/// The shared input is the Phase 0 Test 0.D failure mode (Flat/Hold = δ(0),
/// Long/Short = log-Gaussian at v=-0.001, σ=0.05). Both kernels must reach
/// the same active-direction frequency under Thompson. A KS distance >0.02
/// would localise to either: (i) the inverse-CDF math diverges (would
/// shift mass between bins by tens of %); or (ii) the softmax/atom-value
/// helpers diverge (same effect). RNG quality contributes <0.005 noise at
/// n=100k.
#[test]
#[ignore]
fn test_2b_production_matches_standalone() {
    let stream = make_test_stream();
    let n_atoms: i32 = 21;
    let v_min = -0.5_f32;
    let v_max = 0.5_f32;
    let delta_z = (v_max - v_min) / (n_atoms as f32 - 1.0);

    // Build a single per-direction logits tile matching Phase 0 Test 0.D.
    let mut tile = vec![0.0_f32; (PROD_B0 as usize) * (n_atoms as usize)];
    let center_atom = ((0.0 - v_min) / delta_z).round() as usize;
    tile[3 * n_atoms as usize + center_atom] = 10.0; // Flat δ(v=0)
    tile[n_atoms as usize + center_atom] = 10.0;     // Hold δ(v=0)
    for a in 0..n_atoms as usize {
        let v = v_min + (a as f32) * delta_z;
        let mean = -0.001_f32;
        let sigma = 0.05_f32;
        let logit = -((v - mean).powi(2)) / (2.0 * sigma * sigma);
        tile[2 * n_atoms as usize + a] = logit; // Long
        tile[a] = logit;                         // Short
    }
    let support_tile = {
        let mut s = vec![0.0_f32; (PROD_B0 as usize) * 3];
        for d in 0..(PROD_B0 as usize) {
            s[d * 3] = v_min;
            s[d * 3 + 1] = v_max;
            s[d * 3 + 2] = delta_z;
        }
        s
    };

    // ── Production kernel ────────────────────────────────────────────────
    //
    // Replicate the single tile across the production-kernel batch. Each
    // thread reads the same tile but draws Philox(i, timestep=0, ctr).
    let batch: usize = 100_000;
    let mut full_logits = Vec::with_capacity(batch * tile.len());
    for _ in 0..batch {
        full_logits.extend_from_slice(&tile);
    }
    let mut full_support = Vec::with_capacity(batch * support_tile.len());
    for _ in 0..batch {
        full_support.extend_from_slice(&support_tile);
    }
    let q_values = vec![0.0_f32; batch * PROD_Q_STRIDE];
    let mut fixture = ProdActionSelectFixture::new(
        Arc::clone(&stream),
        batch,
        n_atoms,
        &q_values,
        &full_logits,
        &full_support,
    );
    let prod_actions = fixture.launch(/*eval_mode=*/ false, /*timestep=*/ 0);
    let mut prod_hist = [0_u32; 4];
    for action in &prod_actions {
        let d = decode_dir(*action) as usize;
        prod_hist[d] += 1;
    }

    // ── Standalone v2 kernel ─────────────────────────────────────────────
    //
    // Standalone takes the SINGLE tile + support tile and runs each thread
    // with `seed_idx` keying its LCG. 100 000 seeds matches the production
    // batch size for histogram parity.
    let n_seeds: usize = 100_000;
    let standalone_picks = launch_standalone_v2(
        &stream,
        &tile,
        &support_tile,
        n_atoms,
        /*base_seed=*/ 0,
        n_seeds,
    );
    let mut stand_hist = [0_u32; 4];
    for d in &standalone_picks {
        stand_hist[*d as usize] += 1;
    }

    println!(
        "Test 2.B production hist (S/H/L/F): {:?} (n={})",
        prod_hist, batch,
    );
    println!(
        "Test 2.B standalone hist (S/H/L/F): {:?} (n={})",
        stand_hist, n_seeds,
    );

    let ks = ks_distance_4bin(&prod_hist, &stand_hist);
    println!("Test 2.B KS distance: {:.5}", ks);

    // KS threshold 0.02 — at n=100k the empirical CDF noise floor for
    // matching distributions is ~4*sqrt(2/n) ≈ 0.018. Allow a small
    // additional margin (0.02) for f32 rounding differences between the
    // two kernels' independent softmax / inverse-CDF passes.
    //
    // Algorithm divergences (different inverse-CDF math, different softmax
    // normalisation, etc.) would shift mass between bins by O(10%), so KS
    // would be O(0.1) — well above 0.02. RNG implementation differences
    // alone (Philox vs LCG, both good random sources) contribute < 0.005.
    assert!(
        ks <= 0.02,
        "KS distance prod vs standalone v2 = {:.5} exceeds 0.02 — kernels' \
         direction-Thompson distributions diverge.\n  prod = {:?}\n  stand = {:?}",
        ks, prod_hist, stand_hist,
    );

    // Sanity: both histograms must have non-trivial active-direction mass.
    // If both kernels collapsed to the same mode, KS would still be ~0 but
    // the test wouldn't actually verify Thompson behaviour. With the Phase
    // 0 0.D failure setup we expect P(Long+Short) ≥ 0.20 in both.
    let prod_active = (prod_hist[0] + prod_hist[2]) as f32 / batch as f32;
    let stand_active = (stand_hist[0] + stand_hist[2]) as f32 / n_seeds as f32;
    assert!(
        prod_active >= 0.20,
        "Production active frac {:.4} < 0.20 — Thompson should explore",
        prod_active,
    );
    assert!(
        stand_active >= 0.20,
        "Standalone active frac {:.4} < 0.20 — Thompson should explore",
        stand_active,
    );
}

/// Phase 2 Test 2.C — Magnitude/order/urgency picks unaffected by direction
/// Thompson (behavioral parity vs hand-computed Boltzmann reference).
///
/// The plan-prescribed approach (snapshot pre-T2 kernel output) is impossible
/// — T2 has already landed (commit 52a2663a2). Replacement strategy (ii)
/// from the dispatch brief: assert mag/ord/urg picks against an analytical
/// reference computed in Rust from the same Q-values.
///
/// The mag/ord/urg branches are unchanged across Plan C T2: Boltzmann softmax
/// with adaptive tau = max(q_range, floor) where floor = 0.01 (mag) or 0.5
/// (ord/urg). In EVAL mode (eps=0) all three branches still take the
/// Boltzmann path (NOT strict-argmax — the kernel deliberately keeps
/// stochasticity in eval to avoid lock-in).
///
/// Setup:
/// - C51 logits forcing direction = Long (d=2) deterministically. This avoids
///   the kernel's Hold/Flat → mag_idx=0 short-circuit which would mask
///   the Boltzmann sampling on the magnitude branch.
/// - Magnitude Q = [0.0, 0.0, 1.0] — peak at Full (mag=2). q_range=1.0,
///   tau=max(1.0, 0.01)=1.0. Boltzmann: P(Full) = 1/(1+2/e) ≈ 0.577,
///   P(Small) = P(Half) = e^-1/(1+2/e) ≈ 0.211 each.
/// - Order Q = [1.0, 0.0, 0.0] — peak at Market (ord=0). q_range=1.0,
///   tau=max(1.0, 0.5)=1.0. Same Boltzmann shape: P(Market) ≈ 0.577.
/// - Urgency Q = [0.0, 1.0, 0.0] — peak at urg=1. q_range=1.0, tau=1.0.
///   Same Boltzmann shape: P(urg=1) ≈ 0.577.
///
/// Assertions:
/// - dir_idx is Long (2) for every sample (eval argmax E[Q] over the
///   peaked C51 logits).
/// - mag/ord/urg picks each follow the analytical Boltzmann distribution
///   within finite-sample tolerance (±5% on P(best)).
#[test]
#[ignore]
fn test_2c_mag_ord_urg_branches_unaffected() {
    let stream = make_test_stream();
    let n_atoms: i32 = 21;
    let v_min = -1.0_f32;
    let v_max = 1.0_f32;
    let delta_z = (v_max - v_min) / (n_atoms as f32 - 1.0);

    // C51 logits: peak Long at v=+0.8 atom (E[Q] ≈ +0.8); other directions
    // peaked at v=-0.5 (E[Q] ≈ -0.5). Eval argmax E[Q] picks Long.
    let mut tile = vec![0.0_f32; (PROD_B0 as usize) * (n_atoms as usize)];
    let long_peak = (((0.8 - v_min) / delta_z).round() as usize)
        .min((n_atoms - 1) as usize);
    let other_peak = (((-0.5 - v_min) / delta_z).round() as usize)
        .min((n_atoms - 1) as usize);
    tile[0 * n_atoms as usize + other_peak] = 10.0;          // Short
    tile[1 * n_atoms as usize + other_peak] = 10.0;          // Hold
    tile[2 * n_atoms as usize + long_peak] = 10.0;           // Long (best)
    tile[3 * n_atoms as usize + other_peak] = 10.0;          // Flat

    let batch: usize = 8192;
    let mut full_logits = Vec::with_capacity(batch * tile.len());
    for _ in 0..batch {
        full_logits.extend_from_slice(&tile);
    }
    let support_host = fill_linear_support(batch, v_min, v_max, delta_z);

    // q_values: [direction[4], magnitude[3], order[3], urgency[3]] per sample.
    let mut q_values_host = vec![0.0_f32; batch * PROD_Q_STRIDE];
    for i in 0..batch {
        let base = i * PROD_Q_STRIDE;
        // Direction Q (unused — direction is decided by C51 Thompson/argmax)
        // Magnitude Q (b1=3): [0.0, 0.0, 1.0] — peak at Full
        q_values_host[base + 4 + 0] = 0.0;
        q_values_host[base + 4 + 1] = 0.0;
        q_values_host[base + 4 + 2] = 1.0;
        // Order Q (b2=3): [1.0, 0.0, 0.0] — peak at Market
        q_values_host[base + 7 + 0] = 1.0;
        q_values_host[base + 7 + 1] = 0.0;
        q_values_host[base + 7 + 2] = 0.0;
        // Urgency Q (b3=3): [0.0, 1.0, 0.0] — peak at urg=1
        q_values_host[base + 10 + 0] = 0.0;
        q_values_host[base + 10 + 1] = 1.0;
        q_values_host[base + 10 + 2] = 0.0;
    }

    let mut fixture = ProdActionSelectFixture::new(
        stream,
        batch,
        n_atoms,
        &q_values_host,
        &full_logits,
        &support_host,
    );

    let actions = fixture.launch(/*eval_mode=*/ true, /*timestep=*/ 0);

    // ── Hand-computed Boltzmann reference ────────────────────────────────
    //
    // For each branch with peak Q at one bin and zeros elsewhere:
    //   tau = max(q_range, floor)
    //   P(best) = exp(0) / sum_a exp((q_a - q_max) / tau)
    //           = 1.0 / (1.0 + 2 * exp(-1))
    //           = 1.0 / (1.0 + 2/e)
    //           ≈ 0.5767
    //   P(non-best) = exp(-1) / (1 + 2/e) ≈ 0.2117 each
    //
    // tau = q_range = 1.0 in all three branches (q_max=1, q_min=0; range
    // strictly above the 0.01 / 0.5 floors).
    let p_best_boltzmann = 1.0_f32 / (1.0_f32 + 2.0_f32 * (-1.0_f32).exp());
    let p_other_boltzmann = (-1.0_f32).exp() / (1.0_f32 + 2.0_f32 * (-1.0_f32).exp());

    // ── Histograms ────────────────────────────────────────────────────────
    let mut dir_hist = [0_u32; PROD_B0 as usize];
    let mut mag_hist = [0_u32; PROD_B1 as usize];
    let mut ord_hist = [0_u32; PROD_B2 as usize];
    let mut urg_hist = [0_u32; PROD_B3 as usize];
    for action in &actions {
        dir_hist[decode_dir(*action) as usize] += 1;
        mag_hist[decode_mag(*action) as usize] += 1;
        ord_hist[decode_ord(*action) as usize] += 1;
        urg_hist[decode_urg(*action) as usize] += 1;
    }

    let p = |h: u32| h as f32 / batch as f32;
    println!(
        "Test 2.C dir hist: {:?} (expect Long=1.0)",
        dir_hist,
    );
    println!(
        "Test 2.C mag hist: {:?} P(Full)={:.4} P(Small)={:.4} P(Half)={:.4} \
         (Boltzmann ref: P(best)={:.4}, P(other)={:.4})",
        mag_hist, p(mag_hist[2]), p(mag_hist[0]), p(mag_hist[1]),
        p_best_boltzmann, p_other_boltzmann,
    );
    println!(
        "Test 2.C ord hist: {:?} P(Market)={:.4} (Boltzmann ref: {:.4})",
        ord_hist, p(ord_hist[0]), p_best_boltzmann,
    );
    println!(
        "Test 2.C urg hist: {:?} P(urg=1)={:.4} (Boltzmann ref: {:.4})",
        urg_hist, p(urg_hist[1]), p_best_boltzmann,
    );

    // Direction: every sample picks Long (eval argmax E[Q]).
    assert_eq!(
        dir_hist[2], batch as u32,
        "All samples must pick Long (d=2) under the controlled C51 logits — \
         got dir hist {:?}",
        dir_hist,
    );

    // Tolerance: at batch=8192, 1-σ on a Bernoulli proportion is
    // sqrt(p*(1-p)/n) ≈ 0.0055 for p≈0.58. Use ±5% absolute tolerance
    // (covers ~9σ — safely above any finite-sample noise).
    const TOL: f32 = 0.05;

    let p_full = p(mag_hist[2]);
    let p_small = p(mag_hist[0]);
    let p_half = p(mag_hist[1]);
    assert!(
        (p_full - p_best_boltzmann).abs() < TOL,
        "Magnitude P(Full)={:.4} deviates from Boltzmann reference {:.4} by > {:.2}",
        p_full, p_best_boltzmann, TOL,
    );
    assert!(
        (p_small - p_other_boltzmann).abs() < TOL,
        "Magnitude P(Small)={:.4} deviates from Boltzmann reference {:.4} by > {:.2}",
        p_small, p_other_boltzmann, TOL,
    );
    assert!(
        (p_half - p_other_boltzmann).abs() < TOL,
        "Magnitude P(Half)={:.4} deviates from Boltzmann reference {:.4} by > {:.2}",
        p_half, p_other_boltzmann, TOL,
    );

    let p_market = p(ord_hist[0]);
    assert!(
        (p_market - p_best_boltzmann).abs() < TOL,
        "Order P(Market)={:.4} deviates from Boltzmann reference {:.4} by > {:.2}",
        p_market, p_best_boltzmann, TOL,
    );

    let p_urg1 = p(urg_hist[1]);
    assert!(
        (p_urg1 - p_best_boltzmann).abs() < TOL,
        "Urgency P(urg=1)={:.4} deviates from Boltzmann reference {:.4} by > {:.2}",
        p_urg1, p_best_boltzmann, TOL,
    );
}

/// Phase 2 Test 2.D — Real-batch e2e: synthetic batch with eval-mode argmax
/// matches Rust ground-truth E[Q] argmax per sample; train-mode active_frac
/// > 40%.
///
/// The plan-prescribed approach reuses the Phase 0 Test 0.F converged
/// checkpoint loader. Wiring up that loader (DQN::new + safetensors load
/// + branching forward eval) is heavyweight for a 1-2-hour dispatch and
/// is already exercised by Test 0.F itself. Replacement strategy from the
/// dispatch brief: synthetic batch via direct buffer write, with random
/// per-(sample, direction) C51 logits and per-sample adaptive support
/// covering a realistic v_min/v_max range.
///
/// Setup:
/// - 256 distinct samples; per-sample `b_logits` filled with deterministic
///   pseudorandom values via a hash of `(sample_idx, direction, atom)`.
/// - Per-sample adaptive support: each sample gets its own (v_min, v_max,
///   delta_z) drawn from realistic ranges seen in production
///   (v_min ∈ [-1.0, -0.2], v_max ∈ [0.2, 1.0]).
///
/// Assertions:
/// - Eval mode: per-sample dir_idx EXACTLY matches the argmax over Rust
///   ground-truth E[Q]_d = sum_a softmax(b_logits[i, d])[a] * atom_vals[i, d, a].
/// - Train mode: count(d ∈ {Long=2, Short=0}) / batch > 0.40 — Thompson
///   over realistic non-degenerate distributions explores actively.
#[test]
#[ignore]
fn test_2d_real_batch_e2e() {
    let stream = make_test_stream();
    let n_atoms: i32 = 21;
    let batch: usize = 256;

    // Deterministic pseudorandom: simple integer hash for reproducibility.
    fn hash_to_unit(a: u32, b: u32, c: u32) -> f32 {
        let mut h = a.wrapping_mul(0x9E3779B1);
        h ^= b.wrapping_add(0x85EBCA6B).wrapping_mul(0xC2B2AE35);
        h ^= c.wrapping_add(0x27D4EB2F).wrapping_mul(0x165667B1);
        h ^= h >> 16;
        // Map to [-1, 1].
        let u = (h as f32) / (u32::MAX as f32);
        2.0_f32 * u - 1.0
    }

    // Per-sample adaptive support — realistic [v_min, v_max] ranges.
    let mut support_host = vec![0.0_f32; batch * (PROD_B0 as usize) * 3];
    for i in 0..batch {
        let v_min_offset = hash_to_unit(i as u32, 0xAA, 0).abs() * 0.8;     // [0, 0.8]
        let v_max_offset = hash_to_unit(i as u32, 0xBB, 0).abs() * 0.8;     // [0, 0.8]
        let v_min = -0.2_f32 - v_min_offset;
        let v_max = 0.2_f32 + v_max_offset;
        let delta_z = (v_max - v_min) / (n_atoms as f32 - 1.0);
        for d in 0..(PROD_B0 as usize) {
            let base = i * (PROD_B0 as usize) * 3 + d * 3;
            support_host[base] = v_min;
            support_host[base + 1] = v_max;
            support_host[base + 2] = delta_z;
        }
    }

    // Per-(sample, direction, atom) random logits in [-1, 1] — matches the
    // post-Xavier-init scale of fresh C51 head outputs.
    let mut b_logits_host =
        vec![0.0_f32; batch * (PROD_B0 as usize) * (n_atoms as usize)];
    for i in 0..batch {
        for d in 0..(PROD_B0 as usize) {
            for a in 0..(n_atoms as usize) {
                let base = i * (PROD_B0 as usize) * (n_atoms as usize)
                    + d * (n_atoms as usize) + a;
                b_logits_host[base] = hash_to_unit(i as u32, d as u32, a as u32);
            }
        }
    }

    // q_values uniform; mag/ord/urg fall through Boltzmann.
    let q_values_host = vec![0.0_f32; batch * PROD_Q_STRIDE];

    // ── Compute Rust ground-truth E[Q] argmax per sample ─────────────────
    let mut ground_truth_dir = vec![0_i32; batch];
    let mut ground_truth_eq = vec![[0.0_f32; 4]; batch];
    for i in 0..batch {
        let support_base = i * (PROD_B0 as usize) * 3;
        let v_min = support_host[support_base];
        let delta_z = support_host[support_base + 2];
        let mut best_q = f32::NEG_INFINITY;
        let mut best_d = 0_i32;
        for d in 0..(PROD_B0 as i32) {
            let logits_base = i * (PROD_B0 as usize) * (n_atoms as usize)
                + (d as usize) * (n_atoms as usize);
            let logits =
                &b_logits_host[logits_base..logits_base + n_atoms as usize];
            // Numerically-stable softmax (matches softmax_c51_inline).
            let max_logit = logits
                .iter()
                .cloned()
                .fold(f32::NEG_INFINITY, f32::max);
            let mut sum = 0.0_f32;
            let mut probs = vec![0.0_f32; n_atoms as usize];
            for a in 0..(n_atoms as usize) {
                probs[a] = (logits[a] - max_logit).exp();
                sum += probs[a];
            }
            let inv_sum = if sum > 0.0 { 1.0 / sum } else { 1.0 / n_atoms as f32 };
            for p in probs.iter_mut() {
                *p *= inv_sum;
            }
            // Linear support (atom_positions = NULL in our launch).
            let mut e = 0.0_f32;
            for a in 0..(n_atoms as usize) {
                let atom_v = v_min + (a as f32) * delta_z;
                e += probs[a] * atom_v;
            }
            ground_truth_eq[i][d as usize] = e;
            if e > best_q {
                best_q = e;
                best_d = d;
            }
        }
        ground_truth_dir[i] = best_d;
    }

    // ── Eval mode: dir_idx must match ground truth per sample ────────────
    let mut fixture = ProdActionSelectFixture::new(
        Arc::clone(&stream),
        batch,
        n_atoms,
        &q_values_host,
        &b_logits_host,
        &support_host,
    );
    let eval_actions = fixture.launch(/*eval_mode=*/ true, /*timestep=*/ 0);
    let mut mismatches = 0_usize;
    for i in 0..batch {
        let kernel_dir = decode_dir(eval_actions[i]);
        if kernel_dir != ground_truth_dir[i] {
            if mismatches < 5 {
                eprintln!(
                    "Test 2.D mismatch sample {}: kernel={} gt={} E[Q]={:?}",
                    i, kernel_dir, ground_truth_dir[i], ground_truth_eq[i],
                );
            }
            mismatches += 1;
        }
    }
    assert_eq!(
        mismatches, 0,
        "Eval mode: {} of {} samples have dir_idx ≠ ground-truth argmax E[Q]",
        mismatches, batch,
    );

    // ── Train mode: P(Long+Short) > 0.40 across the batch ────────────────
    //
    // Realistic random Xavier-init logits produce non-degenerate per-direction
    // distributions; Thompson should sample the active directions ≥ 40% of the
    // time. (Per spec Task 9: "count(d ∈ {Long, Short}) / batch_size > 0.40".)
    let train_actions = fixture.launch(/*eval_mode=*/ false, /*timestep=*/ 0);
    let mut train_counts = [0_u32; PROD_B0 as usize];
    for action in &train_actions {
        train_counts[decode_dir(*action) as usize] += 1;
    }
    let active = train_counts[0] + train_counts[2];
    let p_active = active as f32 / batch as f32;
    println!(
        "Test 2.D train counts (S/H/L/F): {:?} active={:.4}",
        train_counts, p_active,
    );
    assert!(
        p_active > 0.40,
        "Train mode active_frac {:.4} expected > 0.40 (Thompson on realistic \
         random logits should explore Long+Short ≥ 40%). Counts: {:?}",
        p_active, train_counts,
    );
}
