//! C51 distributional Q-head for the integrated RL trainer (Phase C).
//!
//! Sits on the shared Mamba2+CfC encoder's `h_t` output. Produces per-action
//! atom logits of shape `[B, N_ACTIONS, Q_N_ATOMS]`. The loss is categorical
//! cross-entropy between the current and Bellman-projected target
//! distributions; γ and target-net τ are sourced from
//! `ISV[RL_GAMMA_INDEX]` / `ISV[RL_TARGET_TAU_INDEX]` per
//! `pearl_controller_anchors_isv_driven` and
//! `feedback_isv_for_adaptive_bounds`.
//!
//! Per `pearl_thompson_for_distributional_action_selection`, rollout-time
//! action selection samples atoms via Thompson; the Bellman bootstrap uses
//! argmax of `E[Q]` (handled inside the loss kernel by the caller-provided
//! `actions_taken` + target distribution). This struct OWNS the device
//! weights + cached kernel handles only; rollout / training-step glue lands
//! in Phase E.
//!
//! ## Architecture
//!
//! ```text
//!   logits[b, a, z] = b[a, z] + Σ_c W[a, z, c] * h_t[b, c]
//! ```
//!
//! where `c` ranges over the encoder hidden dim (`HIDDEN_DIM=128`),
//! `a` over `N_ACTIONS=9` and `z` over `Q_N_ATOMS=21`. The output is RAW
//! logits — the softmax-over-atoms is fused into the loss kernel
//! (`dqn_distributional_q.cu::dqn_distributional_q_bwd`). The target
//! network has identical shape; Phase E's soft-update controller drives
//! the blend `W_target ← (1-τ) W_target + τ W` using τ from
//! `ISV[RL_TARGET_TAU_INDEX]`.
//!
//! ## Initialisation
//!
//! Per `pearl_scoped_init_seed_for_reproducibility`, `DqnHead::new`
//! installs a `scoped_init_seed` guard before drawing any samples so the
//! GPU init path stays deterministic for the given seed. Xavier uniform
//! scaled by `0.01` so the initial atom logits sit near zero
//! (post-softmax atoms ≈ uniform); the target network is initialised to
//! the same Xavier draw so the first Bellman backup sees `W_target == W`.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::cublas::CudaBlas;
use cudarc::cublas::sys as cublas_sys;
use cudarc::cublas::sys::cublasOperation_t;
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut};
use cudarc::driver::sys::CUstream;
use ml_core::cuda_autograd::init::scoped_init_seed;
use ml_core::cuda_autograd::linear::BiasKernels;
use ml_core::device::MlDevice;
use rand::{Rng, SeedableRng};
use rand_chacha::ChaCha8Rng;

use crate::heads::HIDDEN_DIM;
use crate::pinned_mem::MappedF32Buffer;
use crate::rl::common::{N_ACTIONS, Q_N_ATOMS};
use crate::trainer::raw_launch::{RawArgs, raw_launch};

/// Output dimension of the Q-head linear layer: one logit per (action, atom).
const K_OUT: usize = N_ACTIONS * Q_N_ATOMS;

// ── cuBLAS SGEMM helper (mirrors ml-core::cuda_autograd::linear::gemm_ex_f32) ──

/// F32 x F32 -> F32 GEMM via `cublasGemmEx` with F32 internal accumulation.
///
/// # Safety
/// All device pointers must be valid and dimensions must be correct.
#[allow(clippy::too_many_arguments)]
unsafe fn gemm_f32(
    cublas: &CudaBlas,
    transa: cublasOperation_t,
    transb: cublasOperation_t,
    m: i32,
    n: i32,
    k: i32,
    alpha: f32,
    a_ptr: u64,
    lda: i32,
    b_ptr: u64,
    ldb: i32,
    beta: f32,
    c_ptr: u64,
    ldc: i32,
    label: &str,
) -> Result<()> {
    cudarc::cublas::result::gemm_ex(
        *cublas.handle(),
        transa,
        transb,
        m,
        n,
        k,
        (&alpha as *const f32).cast(),
        a_ptr as *const std::ffi::c_void,
        cublas_sys::cudaDataType_t::CUDA_R_32F,
        lda,
        b_ptr as *const std::ffi::c_void,
        cublas_sys::cudaDataType_t::CUDA_R_32F,
        ldb,
        (&beta as *const f32).cast(),
        c_ptr as *mut std::ffi::c_void,
        cublas_sys::cudaDataType_t::CUDA_R_32F,
        ldc,
        cublas_sys::cublasComputeType_t::CUBLAS_COMPUTE_32F,
        cublas_sys::cublasGemmAlgo_t::CUBLAS_GEMM_DFALT,
    )
    .map_err(|e| anyhow::anyhow!("cublasGemmEx {label}: {e:?}"))?;
    Ok(())
}

const DQN_HEAD_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/dqn_distributional_q.cubin"
));
const PPO_LOSS_REDUCE_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/ppo_loss_reduce_b.cubin"
));
const DQN_TARGET_SOFT_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/dqn_target_soft_update.cubin"
));
const BELLMAN_PROJ_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/bellman_target_projection.cubin"
));
const SATURATION_REDUCE_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/rl_bellman_target_saturation_reduce.cubin"
));
const Q_DISTRIBUTION_STATS_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/rl_q_distribution_stats.cubin"
));

/// Construction config for [`DqnHead`].
#[derive(Clone, Debug)]
pub struct DqnHeadConfig {
    /// Encoder hidden dimension `h_t [B, HIDDEN_DIM]` feeding the head.
    /// Matches the production CfC trunk's output width (currently 128).
    /// Stored at construction so the head can carry shape-checks across
    /// later launches without re-deriving from the encoder config.
    pub hidden_dim: usize,
    /// Random seed for Xavier init. Reproducibility-critical per
    /// `pearl_scoped_init_seed_for_reproducibility`.
    pub seed: u64,
}

/// Block size for the occupancy-optimised `dqn_distributional_q_bwd`
/// kernel. Matches `BWD_BLOCK_SIZE` in `dqn_distributional_q.cu`.
const BWD_BLOCK_SIZE: u32 = 256;

/// C51 distributional Q-head. Owns device weights (`w_d`, `b_d`) for the
/// online network plus a parallel target network (`w_target_d`,
/// `b_target_d`) used as the bootstrap source for the categorical Bellman
/// backup. The target network is soft-updated by the Phase E training
/// loop (τ read from `ISV[RL_TARGET_TAU_INDEX]`); this struct exposes
/// the buffers so the update kernel can write them.
///
/// Per `feedback_no_stubs`, the soft-update kernel itself is NOT exposed
/// from this struct in Phase C — Phase E lands the `soft_update_target`
/// method in the same commit as the training loop wiring. The target
/// buffers live here now because the Bellman loss kernel READS them,
/// which is real wiring, not a placeholder.
pub struct DqnHead {
    #[allow(dead_code)]
    cfg: DqnHeadConfig,
    stream: Arc<CudaStream>,
    /// Cached raw `CUstream` handle -- avoids repeated `.cu_stream()`
    /// calls on the hot path (same pattern as `IntegratedTrainer`).
    raw_stream: CUstream,

    _module: Arc<CudaModule>,
    /// F4.1 reducer cubin (`ppo_loss_reduce_b.cu`) kept alive while
    /// `loss_mean_reduce_fn` references it.
    _loss_reduce_module: Arc<CudaModule>,
    /// Forward kernel: `dqn_distributional_q_fwd` from
    /// `cuda/dqn_distributional_q.cu`. Produces raw atom logits.
    pub fwd_fn: CudaFunction,
    /// Backward kernel: `dqn_distributional_q_bwd`. Computes categorical
    /// CE loss + `∂L/∂logits` given a pre-projected target distribution.
    /// Writes per-batch CE to `loss_per_batch[B]` only (no atomicAdd to
    /// any scalar — that's `mean_reduce_b_f32` below).
    pub bwd_fn: CudaFunction,
    /// F4.1 (2026-05-31) generic `[B] → [1]` block tree-reduce mean,
    /// shared with the PPO surrogate path (same kernel name, same
    /// cubin). Used to produce the scalar `loss_out` diagnostic for
    /// the LR controller; replaces the prior in-kernel atomicAdd.
    pub loss_mean_reduce_fn: CudaFunction,
    /// Phase E.2 backward chain: `dqn_grad_w_b_h_t`. Maps the `grad_logits`
    /// output of `bwd_fn` into per-batch `grad_w` / `grad_b` scratch buffers
    /// + an OVERWRITE `grad_h_t` slot for the integrated trainer to fold
    /// into the encoder grad accumulator. Per-batch scratch — cross-batch
    /// reduction is the caller's responsibility via `reduce_axis0`.
    pub grad_w_b_h_t_fn: CudaFunction,

    /// Bellman target projection kernel handle (Phase E.2-DEFER). Reads
    /// γ from ISV[RL_GAMMA_INDEX=400], projects the target-net atom
    /// distribution onto the discrete support [V_MIN, V_MAX] with
    /// linear interpolation. See `cuda/bellman_target_projection.cu`.
    pub bellman_proj_fn: CudaFunction,
    /// Per-batch action-row extractor kernel. Slices a
    /// `[B × N_ACTIONS × Q_N_ATOMS]` target-net output down to the
    /// `[B × Q_N_ATOMS]` shape that `bellman_target_projection` consumes.
    /// See `cuda/bellman_target_projection.cu::dqn_select_action_atoms`.
    pub select_action_atoms_fn: CudaFunction,
    /// Fused `dqn_select_action_atoms` + `bellman_target_projection` in
    /// a single launch. Eliminates the intermediate
    /// `action_logits[B, Q_N_ATOMS]` global memory round-trip by staging
    /// the selected atom row in shared memory. Hot-path replacement for
    /// the two sequential launches; the standalone kernels are retained
    /// for potential standalone use.
    pub bellman_fused_fn: CudaFunction,
    /// Owns the Bellman projection cubin lifetime (all three kernels —
    /// `bellman_target_projection`, `dqn_select_action_atoms`, and
    /// `bellman_fused_select_project` — live in the same translation unit).
    _bellman_module: Arc<CudaModule>,
    /// B-9 (2026-06-01): cross-batch tree-reduce of per-batch
    /// Bellman-target saturation tallies. Reads the four `[B] f32`
    /// scratch buffers populated by either bellman variant and writes
    /// ISV slots 726-729 (top/bot saturation rate + max/min pre-proj).
    /// Pure observability — does not perturb training dynamics.
    pub saturation_reduce_fn: CudaFunction,
    _saturation_reduce_module: Arc<CudaModule>,

    /// B-10 (2026-06-01) G1: Q-distribution informativeness diagnostic.
    /// Two entry points in the same cubin:
    /// - `rl_q_distribution_per_batch`: per-block per-action softmax over
    ///   atoms + entropy + E_Q computation. Grid `(B, 1, 1)` × block
    ///   `(Q_N_ATOMS, 1, 1)`. Writes 3 per-batch scratch arrays.
    /// - `rl_q_distribution_reduce`: cross-batch tree-reduce, single-block,
    ///   writes ISV slots 730 / 731 / 732. Pure observability.
    pub q_distribution_per_batch_fn: CudaFunction,
    pub q_distribution_reduce_fn: CudaFunction,
    _q_distribution_stats_module: Arc<CudaModule>,

    /// Phase R5: element-wise target-net soft update kernel handle.
    /// `target[i] = (1 - τ) · target[i] + τ · current[i]` reading τ
    /// from `ISV[RL_TARGET_TAU_INDEX = 401]`. Closes defect #4 from
    /// the flawed Phase F+G arc (target network was never soft-updated).
    pub target_soft_update_fn: CudaFunction,
    _target_soft_update_module: Arc<CudaModule>,

    /// Online-network weights `[N_ACTIONS × Q_N_ATOMS, HIDDEN_DIM]`,
    /// row-major. Each "row" indexes one `(action, atom)` output slot.
    pub w_d: CudaSlice<f32>,
    /// Online-network biases `[N_ACTIONS × Q_N_ATOMS]`.
    pub b_d: CudaSlice<f32>,
    /// Target-network weights, same shape as `w_d`. Soft-updated by the
    /// Phase E training loop using τ from `ISV[RL_TARGET_TAU_INDEX]`.
    pub w_target_d: CudaSlice<f32>,
    /// Target-network biases, same shape as `b_d`.
    pub b_target_d: CudaSlice<f32>,

    // ── cuBLAS infrastructure ────────────────────────────────────────
    /// cuBLAS handle for SGEMM-based forward / backward passes.
    /// Replaces the hand-written `dqn_distributional_q_fwd` and
    /// `dqn_grad_w_b_h_t` kernels: cuBLAS tensor-core SGEMM is 5-10x
    /// faster at b=256 for the [B, 128] x [128, 231] matrix shapes.
    pub cublas: CudaBlas,
    /// Pre-allocated cuBLAS workspace (8 MiB). Prevents per-call
    /// cudaMalloc inside cuBLAS that would break CUDA Graph capture.
    _cublas_workspace: CudaSlice<u8>,
    /// Bias-add and reduce-sum-axis0 kernels from ml-core.
    bias_kernels: BiasKernels,
}

impl DqnHead {
    /// Allocate device weights, load the cubin, and cache the forward /
    /// backward kernel handles. Online and target networks start with
    /// IDENTICAL Xavier draws so the first Bellman backup sees a
    /// zero-divergence target — the soft-update controller's bootstrap
    /// observation (τ ≈ 0.005) then carries the online weights forward.
    pub fn new(dev: &MlDevice, cfg: DqnHeadConfig) -> Result<Self> {
        let stream: Arc<CudaStream> = dev.cuda_stream().context("dqn_head stream")?.clone();
        let ctx = dev.cuda_context().context("dqn_head ctx")?;
        let module = ctx
            .load_cubin(DQN_HEAD_CUBIN.to_vec())
            .context("load dqn_distributional_q cubin")?;
        let fwd_fn = module
            .load_function("dqn_distributional_q_fwd")
            .context("load dqn_distributional_q_fwd")?;
        let bwd_fn = module
            .load_function("dqn_distributional_q_bwd")
            .context("load dqn_distributional_q_bwd")?;
        let loss_reduce_module = ctx
            .load_cubin(PPO_LOSS_REDUCE_CUBIN.to_vec())
            .context("load ppo_loss_reduce_b cubin (for mean_reduce_b_f32)")?;
        let loss_mean_reduce_fn = loss_reduce_module
            .load_function("mean_reduce_b_f32")
            .context("load mean_reduce_b_f32")?;
        let grad_w_b_h_t_fn = module
            .load_function("dqn_grad_w_b_h_t")
            .context("load dqn_grad_w_b_h_t")?;
        let bellman_module = ctx
            .load_cubin(BELLMAN_PROJ_CUBIN.to_vec())
            .context("load bellman_target_projection cubin")?;
        let bellman_proj_fn = bellman_module
            .load_function("bellman_target_projection")
            .context("load bellman_target_projection")?;
        let select_action_atoms_fn = bellman_module
            .load_function("dqn_select_action_atoms")
            .context("load dqn_select_action_atoms")?;
        let bellman_fused_fn = bellman_module
            .load_function("bellman_fused_select_project")
            .context("load bellman_fused_select_project")?;

        // B-9 (2026-06-01): C51 Bellman-target saturation reducer.
        let saturation_reduce_module = ctx
            .load_cubin(SATURATION_REDUCE_CUBIN.to_vec())
            .context("load rl_bellman_target_saturation_reduce cubin")?;
        let saturation_reduce_fn = saturation_reduce_module
            .load_function("rl_bellman_target_saturation_reduce")
            .context("load rl_bellman_target_saturation_reduce")?;

        // B-10 (2026-06-01): Q-distribution informativeness diagnostic.
        let q_distribution_stats_module = ctx
            .load_cubin(Q_DISTRIBUTION_STATS_CUBIN.to_vec())
            .context("load rl_q_distribution_stats cubin")?;
        let q_distribution_per_batch_fn = q_distribution_stats_module
            .load_function("rl_q_distribution_per_batch")
            .context("load rl_q_distribution_per_batch")?;
        let q_distribution_reduce_fn = q_distribution_stats_module
            .load_function("rl_q_distribution_reduce")
            .context("load rl_q_distribution_reduce")?;

        // Phase R5: target soft-update kernel.
        let target_soft_update_module = ctx
            .load_cubin(DQN_TARGET_SOFT_UPDATE_CUBIN.to_vec())
            .context("load dqn_target_soft_update cubin")?;
        let target_soft_update_fn = target_soft_update_module
            .load_function("dqn_target_soft_update")
            .context("load dqn_target_soft_update")?;

        // Per `pearl_scoped_init_seed_for_reproducibility`: install the
        // scoped seed guard BEFORE drawing any Xavier samples.
        let _seed_guard = scoped_init_seed(cfg.seed);
        let mut rng = ChaCha8Rng::seed_from_u64(cfg.seed);

        let n_out = N_ACTIONS * Q_N_ATOMS;
        let n_in = cfg.hidden_dim;
        // Xavier uniform scaled by 0.01: initial atom logits sit near 0,
        // softmax-over-atoms ≈ uniform (1/Q_N_ATOMS each), so the head
        // is uninformed at step 0 and learning is driven by the Bellman
        // gradient rather than init noise.
        let scale = 0.01_f32 * (6.0_f32 / (n_in + n_out) as f32).sqrt();
        let w_host: Vec<f32> = (0..n_out * n_in)
            .map(|_| rng.gen_range(-scale..scale))
            .collect();
        let b_host: Vec<f32> = vec![0.0_f32; n_out];

        let w_d = upload(&stream, &w_host)?;
        let b_d = upload(&stream, &b_host)?;
        // Target network: identical init so step-0 Bellman backup has
        // zero target-net divergence (the divergence-driven τ controller
        // then bootstraps from this observation).
        let w_target_d = upload(&stream, &w_host)?;
        let b_target_d = upload(&stream, &b_host)?;

        // cuBLAS handle + pre-allocated workspace (same pattern as
        // Mamba2Block::new). 8 MiB workspace prevents per-call
        // cudaMalloc that would break CUDA Graph stream capture.
        let cublas = CudaBlas::new(Arc::clone(&stream))
            .context("DqnHead: cuBLAS init")?;
        // Determinism foundation (spec 2026-06-02 §2.B): disable
        // non-deterministic GEMM algorithms. Same fix as Mamba2Block,
        // applied to the DQN distributional-Q forward + backward GEMMs.
        let cublas = crate::cublas_determinism::apply_deterministic_math_mode(cublas)
            .context("DqnHead: cublas determinism setup")?;
        const CUBLAS_WORKSPACE_BYTES: usize = 8 * 1024 * 1024;
        let cublas_workspace = stream
            .alloc_zeros::<u8>(CUBLAS_WORKSPACE_BYTES)
            .context("DqnHead: cuBLAS workspace alloc")?;
        unsafe {
            let ws_ptr = cublas_workspace.raw_ptr();
            cudarc::cublas::sys::cublasSetWorkspace_v2(
                *cublas.handle(),
                ws_ptr as *mut std::ffi::c_void,
                CUBLAS_WORKSPACE_BYTES,
            )
            .result()
            .map_err(|e| anyhow::anyhow!("DqnHead: cublasSetWorkspace_v2: {e:?}"))?;
        }

        // Bias-add and reduce-sum-axis0 kernel handles from ml-core.
        let bias_kernels = BiasKernels::shared(&stream)
            .map_err(|e| anyhow::anyhow!("DqnHead: BiasKernels init: {e}"))?;

        let raw_stream = stream.cu_stream();

        Ok(Self {
            cfg,
            stream,
            raw_stream,
            _module: module,
            _loss_reduce_module: loss_reduce_module,
            loss_mean_reduce_fn,
            fwd_fn,
            bwd_fn,
            grad_w_b_h_t_fn,
            bellman_proj_fn,
            select_action_atoms_fn,
            bellman_fused_fn,
            _bellman_module: bellman_module,
            saturation_reduce_fn,
            _saturation_reduce_module: saturation_reduce_module,
            q_distribution_per_batch_fn,
            q_distribution_reduce_fn,
            _q_distribution_stats_module: q_distribution_stats_module,
            target_soft_update_fn,
            _target_soft_update_module: target_soft_update_module,
            w_d,
            b_d,
            w_target_d,
            b_target_d,
            cublas,
            _cublas_workspace: cublas_workspace,
            bias_kernels,
        })
    }

    /// Phase R5: launch `dqn_target_soft_update` on both the weight
    /// and bias tensors. Element-wise
    /// `target[i] = (1 − τ) · target[i] + τ · current[i]` reading τ
    /// from `ISV[RL_TARGET_TAU_INDEX = 401]`. Closes defect #4 from
    /// the flawed Phase F+G arc (no target-net soft update).
    ///
    /// Should be called once per training step (R6 wires this from
    /// `step_with_lobsim` after the Q-head Adam update so the soft
    /// update reflects the latest online weights).
    pub fn soft_update_target(&mut self, isv_dev_ptr: &u64) -> Result<()> {
        let func = self.target_soft_update_fn.cu_function();
        // Weights.
        {
            let n_w = self.w_d.len();
            let n_w_i = n_w as i32;
            let mut args = RawArgs::new();
            args.push_ptr(self.w_d.raw_ptr());
            args.push_ptr(self.w_target_d.raw_ptr());
            args.push_ptr(*isv_dev_ptr);
            args.push_i32(n_w_i);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    func,
                    (((n_w as u32) + 255) / 256, 1, 1), (256, 1, 1), 0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                ).map_err(|e| anyhow::anyhow!("dqn_target_soft_update (weights): {:?}", e))?;
            }
        }
        // Biases.
        {
            let n_b = self.b_d.len();
            let n_b_i = n_b as i32;
            let mut args = RawArgs::new();
            args.push_ptr(self.b_d.raw_ptr());
            args.push_ptr(self.b_target_d.raw_ptr());
            args.push_ptr(*isv_dev_ptr);
            args.push_i32(n_b_i);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    func,
                    (((n_b as u32) + 31) / 32, 1, 1), (32, 1, 1), 0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                ).map_err(|e| anyhow::anyhow!("dqn_target_soft_update (biases): {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Phase E.2 forward via cuBLAS SGEMM:
    ///   `logits[B, K_OUT] = h_t[B, HIDDEN_DIM] @ W^T[HIDDEN_DIM, K_OUT] + bias[K_OUT]`
    ///
    /// Replaces the hand-written `dqn_distributional_q_fwd` kernel whose
    /// per-thread serial loop over HIDDEN_DIM=128 consumed 18.1% of GPU
    /// time. cuBLAS tensor-core SGEMM is 5-10x faster at b=256.
    ///
    /// The trainer pre-allocates `logits_out` so the same memory is
    /// reused across step() calls.
    pub fn forward(
        &self,
        h_t: &CudaSlice<f32>,
        b_size: usize,
        logits_out: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(logits_out.len(), b_size * K_OUT);
        self.forward_gemm(&self.w_d, &self.b_d, h_t, b_size, logits_out)
            .context("dqn forward (online)")
    }

    /// cuBLAS SGEMM forward shared by online and target network paths.
    ///
    /// Computes `logits[B, K_OUT] = h_t[B, HIDDEN_DIM] @ W^T + bias`:
    ///   - cuBLAS column-major: C_col[K_OUT, B] = W_col^T[K_OUT, HIDDEN_DIM] @ X_col[HIDDEN_DIM, B]
    ///     transA=T, transB=N, m=K_OUT, n=B, k=HIDDEN_DIM
    ///   - Then add_bias_2d kernel broadcasts bias[K_OUT] into logits[B, K_OUT].
    fn forward_gemm(
        &self,
        w: &CudaSlice<f32>,
        b: &CudaSlice<f32>,
        h_t: &CudaSlice<f32>,
        b_size: usize,
        logits_out: &mut CudaSlice<f32>,
    ) -> Result<()> {
        let w_ptr = w.raw_ptr();
        let h_ptr = h_t.raw_ptr();
        let out_ptr = logits_out.raw_ptr();
        // SGEMM: logits = h_t @ W^T (alpha=1, beta=0 overwrites logits_out)
        unsafe {
            gemm_f32(
                &self.cublas,
                cublasOperation_t::CUBLAS_OP_T, // transA: W[K_OUT, HIDDEN_DIM] stored row-major = col[HIDDEN_DIM, K_OUT]
                cublasOperation_t::CUBLAS_OP_N, // transB: h_t[B, HIDDEN_DIM] stored row-major = col[HIDDEN_DIM, B]
                K_OUT as i32,                    // m
                b_size as i32,                   // n
                HIDDEN_DIM as i32,               // k
                1.0,                             // alpha
                w_ptr,                           // A = W
                HIDDEN_DIM as i32,               // lda (W row-major [K_OUT, HIDDEN_DIM] → col lda = HIDDEN_DIM)
                h_ptr,                           // B = h_t
                HIDDEN_DIM as i32,               // ldb
                0.0,                             // beta
                out_ptr,                         // C = logits_out
                K_OUT as i32,                    // ldc
                "dqn_fwd_sgemm",
            )?;
        }
        // Broadcast bias: logits_out[b, j] += bias[j]
        self.add_bias(logits_out, b, b_size)?;
        Ok(())
    }

    /// Launch the `add_bias_2d_kernel` from ml-core's BiasKernels.
    fn add_bias(
        &self,
        y: &mut CudaSlice<f32>,
        bias: &CudaSlice<f32>,
        rows: usize,
    ) -> Result<()> {
        let total = rows * K_OUT;
        let threads = 256_u32;
        let blocks = ((total as u32) + threads - 1) / threads;
        let rows_i32 = rows as i32;
        let cols_i32 = K_OUT as i32;
        let mut args = RawArgs::new();
        args.push_ptr(y.raw_ptr());
        args.push_ptr(bias.raw_ptr());
        args.push_i32(rows_i32);
        args.push_i32(cols_i32);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.bias_kernels.add_fn.cu_function(),
                (blocks, 1, 1), (threads, 1, 1), 0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("dqn add_bias_2d: {:?}", e))?;
        }
        Ok(())
    }

    /// Phase E.2 backward chain stage 1 — categorical CE on the taken
    /// action's atom distribution. Writes `loss_per_batch [B]` and
    /// `grad_logits [B × N_ACTIONS × Q_N_ATOMS]`.
    ///
    /// F4.1 (2026-05-31): the cross-batch scalar `loss_out` is now
    /// produced by a separate `mean_reduce_b_f32` block tree-reduce
    /// kernel that the caller invokes on `loss_per_batch`. The prior
    /// in-kernel `atomicAdd(loss_out, ce/B)` is gone (per
    /// `feedback_no_atomicadd`).
    #[allow(clippy::too_many_arguments)]
    pub fn backward_logits(
        &self,
        logits: &CudaSlice<f32>,
        target_dist: &CudaSlice<f32>,
        actions_taken: &CudaSlice<i32>,
        b_size: usize,
        loss_per_batch: &mut CudaSlice<f32>,
        grad_logits: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(logits.len(), b_size * N_ACTIONS * Q_N_ATOMS);
        debug_assert_eq!(target_dist.len(), b_size * Q_N_ATOMS);
        debug_assert_eq!(actions_taken.len(), b_size);
        debug_assert_eq!(
            loss_per_batch.len(),
            b_size,
            "R7d: per-sample CE buffer must match batch size for PER priority update"
        );
        debug_assert_eq!(grad_logits.len(), b_size * N_ACTIONS * Q_N_ATOMS);

        let b_i = b_size as i32;
        let mut args = RawArgs::new();
        args.push_ptr(logits.raw_ptr());
        args.push_ptr(target_dist.raw_ptr());
        args.push_ptr(actions_taken.raw_ptr());
        args.push_i32(b_i);
        args.push_ptr(loss_per_batch.raw_ptr());
        args.push_ptr(grad_logits.raw_ptr());
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.bwd_fn.cu_function(),
                (b_size as u32, 1, 1), (BWD_BLOCK_SIZE, 1, 1), 0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("dqn_distributional_q_bwd: {:?}", e))?;
        }
        Ok(())
    }

    /// F4.1 (2026-05-31) block tree-reduce `loss_per_batch [B] → scalar
    /// mean` written to `loss_out_dev_ptr`. Replaces the prior in-kernel
    /// `atomicAdd(loss_out, ce/B)` from `dqn_distributional_q_bwd`. Per
    /// `feedback_no_atomicadd`: single block, 256 threads, grid-strided
    /// over B, shared-memory tree-reduce, single-writer store.
    pub fn reduce_loss_to_scalar(
        &self,
        loss_per_batch: &CudaSlice<f32>,
        b_size: usize,
        loss_out_dev_ptr: &u64,
    ) -> Result<()> {
        debug_assert_eq!(loss_per_batch.len(), b_size);
        let b_i = b_size as i32;
        let mut args = RawArgs::new();
        args.push_ptr(loss_per_batch.raw_ptr());
        args.push_i32(b_i);
        args.push_ptr(*loss_out_dev_ptr);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.loss_mean_reduce_fn.cu_function(),
                (1, 1, 1),
                (256, 1, 1),
                0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("mean_reduce_b_f32 (dqn loss): {:?}", e))?;
        }
        Ok(())
    }

    /// Phase E.2 backward chain stage 2 (LEGACY — kept for reference;
    /// production path is [`backward_gemm`]).
    ///
    /// Propagates `grad_logits` into per-batch `grad_w_per_batch`,
    /// `grad_b_per_batch`, and `grad_h_t` via the hand-written
    /// `dqn_grad_w_b_h_t` kernel. Callers must still reduce
    /// `grad_w_per_batch` / `grad_b_per_batch` along axis 0 via
    /// `reduce_axis0`. Superseded by [`backward_gemm`] which produces
    /// the reduced grad_w / grad_b directly via cuBLAS SGEMM +
    /// reduce-sum kernel.
    #[allow(clippy::too_many_arguments)]
    pub fn backward_to_w_b_h(
        &self,
        h_t: &CudaSlice<f32>,
        grad_logits: &CudaSlice<f32>,
        b_size: usize,
        grad_w_per_batch: &mut CudaSlice<f32>,
        grad_b_per_batch: &mut CudaSlice<f32>,
        grad_h_t: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(grad_logits.len(), b_size * K_OUT);
        debug_assert_eq!(
            grad_w_per_batch.len(),
            b_size * K_OUT * HIDDEN_DIM
        );
        debug_assert_eq!(grad_b_per_batch.len(), b_size * K_OUT);
        debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM);

        let b_i = b_size as i32;
        let smem = (K_OUT * std::mem::size_of::<f32>()) as u32;
        let mut args = RawArgs::new();
        args.push_ptr(self.w_d.raw_ptr());
        args.push_ptr(h_t.raw_ptr());
        args.push_ptr(grad_logits.raw_ptr());
        args.push_i32(b_i);
        args.push_ptr(grad_w_per_batch.raw_ptr());
        args.push_ptr(grad_b_per_batch.raw_ptr());
        args.push_ptr(grad_h_t.raw_ptr());
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.grad_w_b_h_t_fn.cu_function(),
                (b_size as u32, 1, 1), (HIDDEN_DIM as u32, 1, 1), smem,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("dqn_grad_w_b_h_t: {:?}", e))?;
        }
        Ok(())
    }

    /// cuBLAS backward: compute grad_w, grad_b, and grad_h_t from
    /// `grad_logits` in three SGEMM + one reduce-sum launch. Replaces
    /// the hand-written `dqn_grad_w_b_h_t` kernel + `reduce_axis0`.
    ///
    /// Mathematics:
    ///   grad_h_t [B, HIDDEN_DIM] = grad_logits [B, K_OUT] @ W [K_OUT, HIDDEN_DIM]
    ///   grad_w   [K_OUT, HIDDEN_DIM] = grad_logits^T [K_OUT, B] @ h_t [B, HIDDEN_DIM]
    ///   grad_b   [K_OUT] = sum(grad_logits [B, K_OUT], axis=0)
    ///
    /// Benefits:
    ///   - Eliminates the B*K_OUT*HIDDEN_DIM per-batch scratch (7.5 MiB at B=256).
    ///   - Eliminates two `reduce_axis0` kernel launches.
    ///   - cuBLAS SGEMM is 5-10x faster than the hand-written kernel for
    ///     these shapes ([B=256, 128] x [128, 231]).
    ///   - grad_w / grad_b are produced directly reduced, ready for Adam.
    #[allow(clippy::too_many_arguments)]
    pub fn backward_gemm(
        &self,
        h_t: &CudaSlice<f32>,
        grad_logits: &CudaSlice<f32>,
        b_size: usize,
        grad_w: &mut CudaSlice<f32>,
        grad_b: &mut CudaSlice<f32>,
        grad_h_t: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(grad_logits.len(), b_size * K_OUT);
        debug_assert_eq!(grad_w.len(), K_OUT * HIDDEN_DIM);
        debug_assert_eq!(grad_b.len(), K_OUT);
        debug_assert_eq!(grad_h_t.len(), b_size * HIDDEN_DIM);

        let gl_ptr = grad_logits.raw_ptr();
        let ht_ptr = h_t.raw_ptr();
        let w_ptr = self.w_d.raw_ptr();

        // ── grad_h_t [B, HIDDEN_DIM] = grad_logits [B, K_OUT] @ W [K_OUT, HIDDEN_DIM] ──
        // Column-major: grad_h_t_col[HIDDEN_DIM, B] = W_col[HIDDEN_DIM, K_OUT] @ gl_col[K_OUT, B]
        //   transA=N, transB=N, m=HIDDEN_DIM, n=B, k=K_OUT
        let gh_ptr = grad_h_t.raw_ptr();
        unsafe {
            gemm_f32(
                &self.cublas,
                cublasOperation_t::CUBLAS_OP_N, // W[K_OUT, HIDDEN_DIM] row-major = col[HIDDEN_DIM, K_OUT]
                cublasOperation_t::CUBLAS_OP_N, // gl[B, K_OUT] row-major = col[K_OUT, B]
                HIDDEN_DIM as i32,              // m
                b_size as i32,                  // n
                K_OUT as i32,                   // k
                1.0,                            // alpha
                w_ptr,                          // A = W
                HIDDEN_DIM as i32,              // lda
                gl_ptr,                         // B = grad_logits
                K_OUT as i32,                   // ldb
                0.0,                            // beta
                gh_ptr,                         // C = grad_h_t
                HIDDEN_DIM as i32,              // ldc
                "dqn_bwd_grad_h_t",
            )?;
        }

        // ── grad_w [K_OUT, HIDDEN_DIM] = grad_logits^T [K_OUT, B] @ h_t [B, HIDDEN_DIM] ──
        // Column-major: grad_w_col[HIDDEN_DIM, K_OUT] = h_t_col[HIDDEN_DIM, B] @ gl_col[K_OUT, B]^T
        //   transA=N, transB=T, m=HIDDEN_DIM, n=K_OUT, k=B
        let gw_ptr = grad_w.raw_ptr();
        unsafe {
            gemm_f32(
                &self.cublas,
                cublasOperation_t::CUBLAS_OP_N, // h_t[B, HIDDEN_DIM] row-major = col[HIDDEN_DIM, B]
                cublasOperation_t::CUBLAS_OP_T, // gl[B, K_OUT] row-major = col[K_OUT, B]; transposed
                HIDDEN_DIM as i32,              // m
                K_OUT as i32,                   // n
                b_size as i32,                  // k
                1.0,                            // alpha
                ht_ptr,                         // A = h_t
                HIDDEN_DIM as i32,              // lda
                gl_ptr,                         // B = grad_logits
                K_OUT as i32,                   // ldb
                0.0,                            // beta
                gw_ptr,                         // C = grad_w
                HIDDEN_DIM as i32,              // ldc
                "dqn_bwd_grad_w",
            )?;
        }

        // ── grad_b [K_OUT] = sum(grad_logits [B, K_OUT], axis=0) ──────
        let threads = 256_u32;
        let blocks = ((K_OUT as u32) + threads - 1) / threads;
        let rows_i32 = b_size as i32;
        let cols_i32 = K_OUT as i32;
        let mut args = RawArgs::new();
        args.push_ptr(grad_logits.raw_ptr());
        args.push_ptr(grad_b.raw_ptr());
        args.push_i32(rows_i32);
        args.push_i32(cols_i32);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.bias_kernels.reduce_fn.cu_function(),
                (blocks, 1, 1), (threads, 1, 1), 0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("dqn_bwd_grad_b reduce_sum_axis0: {:?}", e))?;
        }
        Ok(())
    }

    /// Target-network forward via cuBLAS SGEMM (same math as online
    /// forward, different weight/bias buffers). Used by the Bellman
    /// target-bootstrap path.
    pub fn forward_target(
        &self,
        h_t: &CudaSlice<f32>,
        b_size: usize,
        logits_out: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(h_t.len(), b_size * HIDDEN_DIM);
        debug_assert_eq!(logits_out.len(), b_size * K_OUT);
        self.forward_gemm(&self.w_target_d, &self.b_target_d, h_t, b_size, logits_out)
            .context("dqn forward (target)")
    }

    /// Extract per-batch action-row atoms from a full target-net output
    /// `[B × N_ACTIONS × Q_N_ATOMS]` into the shape
    /// `[B × Q_N_ATOMS]` expected by
    /// [`project_bellman_target`]. Caller-provided `actions_d` selects
    /// the action index per batch — typically the argmax-Q action of the
    /// NEXT state for the standard distributional double-DQN target.
    pub fn select_action_atoms(
        &self,
        full_logits_d: &CudaSlice<f32>,
        actions_d: &CudaSlice<i32>,
        b_size: usize,
        action_logits_out: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(full_logits_d.len(), b_size * N_ACTIONS * Q_N_ATOMS);
        debug_assert_eq!(actions_d.len(), b_size);
        debug_assert_eq!(action_logits_out.len(), b_size * Q_N_ATOMS);

        let b_i = b_size as i32;
        let mut args = RawArgs::new();
        args.push_ptr(full_logits_d.raw_ptr());
        args.push_ptr(actions_d.raw_ptr());
        args.push_i32(b_i);
        args.push_ptr(action_logits_out.raw_ptr());
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.select_action_atoms_fn.cu_function(),
                (b_size as u32, 1, 1), (Q_N_ATOMS as u32, 1, 1), 0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("dqn_select_action_atoms: {:?}", e))?;
        }
        Ok(())
    }

    /// Phase E.2-DEFER Bellman categorical projection. Reads γ from
    /// `ISV[RL_GAMMA_INDEX=400]` and emits the projected target
    /// distribution `target_dist [B × Q_N_ATOMS]` consumed by
    /// [`backward_logits`].
    ///
    /// `target_logits` must come from a SINGLE action — typically the
    /// argmax-Q action of `s_{t+1}` per the standard double-DQN /
    /// distributional Q convention. The caller is responsible for
    /// selecting + materialising that single-action atom row from the
    /// target-net output. For the synthetic toy-bandit path the trainer
    /// uses the action that the policy actually took.
    #[allow(clippy::too_many_arguments)]
    pub fn project_bellman_target(
        &self,
        target_logits_d: &CudaSlice<f32>,
        rewards_d: &CudaSlice<f32>,
        dones_d: &CudaSlice<f32>,
        n_step_gammas_d: &CudaSlice<f32>,
        isv_dev_ptr: &u64,
        b_size: usize,
        target_dist_d: &mut CudaSlice<f32>,
        // B-9 (2026-06-01): per-batch saturation tally scratch outputs;
        // consumed by `launch_saturation_reduce` after this call.
        sat_top_per_batch_d: &mut CudaSlice<f32>,
        sat_bot_per_batch_d: &mut CudaSlice<f32>,
        max_pre_per_batch_d: &mut CudaSlice<f32>,
        min_pre_per_batch_d: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(target_logits_d.len(), b_size * Q_N_ATOMS);
        debug_assert_eq!(rewards_d.len(), b_size);
        debug_assert_eq!(dones_d.len(), b_size);
        debug_assert_eq!(n_step_gammas_d.len(), b_size);
        debug_assert_eq!(target_dist_d.len(), b_size * Q_N_ATOMS);
        debug_assert_eq!(sat_top_per_batch_d.len(), b_size);
        debug_assert_eq!(sat_bot_per_batch_d.len(), b_size);
        debug_assert_eq!(max_pre_per_batch_d.len(), b_size);
        debug_assert_eq!(min_pre_per_batch_d.len(), b_size);

        let b_i = b_size as i32;
        let mut args = RawArgs::new();
        args.push_ptr(target_logits_d.raw_ptr());
        args.push_ptr(rewards_d.raw_ptr());
        args.push_ptr(dones_d.raw_ptr());
        args.push_ptr(n_step_gammas_d.raw_ptr());
        args.push_ptr(*isv_dev_ptr);
        args.push_i32(b_i);
        args.push_ptr(target_dist_d.raw_ptr());
        args.push_ptr(sat_top_per_batch_d.raw_ptr());
        args.push_ptr(sat_bot_per_batch_d.raw_ptr());
        args.push_ptr(max_pre_per_batch_d.raw_ptr());
        args.push_ptr(min_pre_per_batch_d.raw_ptr());
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.bellman_proj_fn.cu_function(),
                (b_size as u32, 1, 1), (Q_N_ATOMS as u32, 1, 1), 0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("bellman_target_projection: {:?}", e))?;
        }
        Ok(())
    }

    /// Fused `select_action_atoms` + `project_bellman_target` in a single
    /// kernel launch. Reads the selected action's atom row from
    /// `full_logits_d [B × N_ACTIONS × Q_N_ATOMS]` into shared memory,
    /// then runs the Bellman categorical projection in-place — eliminating
    /// the intermediate `[B × Q_N_ATOMS]` global memory buffer and one
    /// kernel launch overhead.
    ///
    /// Produces identical output to the sequential
    /// `select_action_atoms` → `project_bellman_target` pair.
    #[allow(clippy::too_many_arguments)]
    pub fn fused_select_and_project_bellman(
        &self,
        full_logits_d: &CudaSlice<f32>,
        actions_d: &CudaSlice<i32>,
        rewards_d: &CudaSlice<f32>,
        dones_d: &CudaSlice<f32>,
        n_step_gammas_d: &CudaSlice<f32>,
        isv_dev_ptr: &u64,
        b_size: usize,
        target_dist_d: &mut CudaSlice<f32>,
        // B-9 (2026-06-01): per-batch saturation tally scratch outputs;
        // consumed by `launch_saturation_reduce` after this call.
        sat_top_per_batch_d: &mut CudaSlice<f32>,
        sat_bot_per_batch_d: &mut CudaSlice<f32>,
        max_pre_per_batch_d: &mut CudaSlice<f32>,
        min_pre_per_batch_d: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(full_logits_d.len(), b_size * N_ACTIONS * Q_N_ATOMS);
        debug_assert_eq!(actions_d.len(), b_size);
        debug_assert_eq!(rewards_d.len(), b_size);
        debug_assert_eq!(dones_d.len(), b_size);
        debug_assert_eq!(n_step_gammas_d.len(), b_size);
        debug_assert_eq!(target_dist_d.len(), b_size * Q_N_ATOMS);
        debug_assert_eq!(sat_top_per_batch_d.len(), b_size);
        debug_assert_eq!(sat_bot_per_batch_d.len(), b_size);
        debug_assert_eq!(max_pre_per_batch_d.len(), b_size);
        debug_assert_eq!(min_pre_per_batch_d.len(), b_size);

        let b_i = b_size as i32;
        let mut args = RawArgs::new();
        args.push_ptr(full_logits_d.raw_ptr());
        args.push_ptr(actions_d.raw_ptr());
        args.push_ptr(rewards_d.raw_ptr());
        args.push_ptr(dones_d.raw_ptr());
        args.push_ptr(n_step_gammas_d.raw_ptr());
        args.push_ptr(*isv_dev_ptr);
        args.push_i32(b_i);
        args.push_ptr(target_dist_d.raw_ptr());
        args.push_ptr(sat_top_per_batch_d.raw_ptr());
        args.push_ptr(sat_bot_per_batch_d.raw_ptr());
        args.push_ptr(max_pre_per_batch_d.raw_ptr());
        args.push_ptr(min_pre_per_batch_d.raw_ptr());
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.bellman_fused_fn.cu_function(),
                (b_size as u32, 1, 1), (Q_N_ATOMS as u32, 1, 1), 0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("bellman_fused_select_project: {:?}", e))?;
        }
        Ok(())
    }

    /// B-9 (2026-06-01): launch the cross-batch saturation reducer
    /// kernel. Reads the four `[B] f32` scratch buffers populated by
    /// either `project_bellman_target` or `fused_select_and_project_bellman`
    /// and writes per-step rates + extremes to ISV slots 726-729. Call
    /// this AFTER one of the two bellman variants runs in a given step.
    #[allow(clippy::too_many_arguments)]
    pub fn launch_saturation_reduce(
        &self,
        isv_dev_ptr: &u64,
        b_size: usize,
        sat_top_per_batch_d: &CudaSlice<f32>,
        sat_bot_per_batch_d: &CudaSlice<f32>,
        max_pre_per_batch_d: &CudaSlice<f32>,
        min_pre_per_batch_d: &CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(sat_top_per_batch_d.len(), b_size);
        debug_assert_eq!(sat_bot_per_batch_d.len(), b_size);
        debug_assert_eq!(max_pre_per_batch_d.len(), b_size);
        debug_assert_eq!(min_pre_per_batch_d.len(), b_size);

        // Power-of-2 block size, capped at 256 and bounded above by
        // next_power_of_two(B). Identity-padded shared mem ensures
        // threads with no work don't pollute the tree-reduce.
        let block_dim = (b_size as u32).next_power_of_two().clamp(1, 256);
        let smem = (4 * block_dim as usize * std::mem::size_of::<f32>()) as u32;
        let b_i = b_size as i32;
        let mut args = RawArgs::new();
        args.push_ptr(*isv_dev_ptr);
        args.push_ptr(sat_top_per_batch_d.raw_ptr());
        args.push_ptr(sat_bot_per_batch_d.raw_ptr());
        args.push_ptr(max_pre_per_batch_d.raw_ptr());
        args.push_ptr(min_pre_per_batch_d.raw_ptr());
        args.push_i32(b_i);
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.saturation_reduce_fn.cu_function(),
                (1, 1, 1), (block_dim, 1, 1), smem,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("rl_bellman_target_saturation_reduce: {:?}", e))?;
        }
        Ok(())
    }

    /// B-10 (2026-06-01) G1: launch the Q-distribution informativeness
    /// diagnostic — two-kernel pattern (per-batch tally + cross-batch
    /// reduce) consuming online Q logits + the atom support buffer.
    /// Writes ISV slots 730 / 731 / 732. Pure observability; safe to
    /// call after each Q forward.
    #[allow(clippy::too_many_arguments)]
    pub fn launch_q_distribution_stats(
        &self,
        isv_dev_ptr: &u64,
        q_logits_d: &CudaSlice<f32>,
        atom_supports_d: &CudaSlice<f32>,
        b_size: usize,
        // Per-batch scratch (overwritten each call).
        q_ent_per_batch_d: &mut CudaSlice<f32>,
        q_range_per_batch_d: &mut CudaSlice<f32>,
        q_abs_max_per_batch_d: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(q_logits_d.len(), b_size * N_ACTIONS * Q_N_ATOMS);
        debug_assert_eq!(atom_supports_d.len(), Q_N_ATOMS);
        debug_assert_eq!(q_ent_per_batch_d.len(), b_size);
        debug_assert_eq!(q_range_per_batch_d.len(), b_size);
        debug_assert_eq!(q_abs_max_per_batch_d.len(), b_size);

        let b_i = b_size as i32;

        // ── Pass 1: per-batch tally (grid (B,1,1) × block (Q_N_ATOMS,1,1)).
        {
            let mut args = RawArgs::new();
            args.push_ptr(q_logits_d.raw_ptr());
            args.push_ptr(atom_supports_d.raw_ptr());
            args.push_i32(b_i);
            args.push_ptr(q_ent_per_batch_d.raw_ptr());
            args.push_ptr(q_range_per_batch_d.raw_ptr());
            args.push_ptr(q_abs_max_per_batch_d.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.q_distribution_per_batch_fn.cu_function(),
                    (b_size as u32, 1, 1), (Q_N_ATOMS as u32, 1, 1), 0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                ).map_err(|e| anyhow::anyhow!("rl_q_distribution_per_batch: {:?}", e))?;
            }
        }

        // ── Pass 2: cross-batch tree reduce (single block).
        {
            let block_dim = (b_size as u32).next_power_of_two().clamp(1, 256);
            let smem = (3 * block_dim as usize * std::mem::size_of::<f32>()) as u32;
            let mut args = RawArgs::new();
            args.push_ptr(*isv_dev_ptr);
            args.push_ptr(q_ent_per_batch_d.raw_ptr());
            args.push_ptr(q_range_per_batch_d.raw_ptr());
            args.push_ptr(q_abs_max_per_batch_d.raw_ptr());
            args.push_i32(b_i);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.q_distribution_reduce_fn.cu_function(),
                    (1, 1, 1), (block_dim, 1, 1), smem,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                ).map_err(|e| anyhow::anyhow!("rl_q_distribution_reduce: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Stream used to launch all kernels owned by this head. Phase E's
    /// training loop reads this when scheduling the soft-update kernel.
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Encoder hidden dim this head was built for. Used by Phase E
    /// shape-checking when assembling the launch config.
    pub fn hidden_dim(&self) -> usize {
        self.cfg.hidden_dim
    }
}

// ── pinned-staging upload helper (mirrors aux_heads.rs::upload) ──

fn upload(stream: &Arc<CudaStream>, host: &[f32]) -> Result<CudaSlice<f32>> {
    let n = host.len();
    let staging = unsafe { MappedF32Buffer::new(n) }
        .map_err(|e| anyhow::anyhow!("dqn_head upload staging: {e}"))?;
    staging.write_from_slice(host);
    let mut dst = stream
        .alloc_zeros::<f32>(n)
        .context("dqn_head upload alloc")?;
    if n > 0 {
        let nbytes = n * std::mem::size_of::<f32>();
        unsafe {
            let (dst_ptr, _g) = dst.device_ptr_mut(stream);
            cudarc::driver::result::memcpy_dtod_async(
                dst_ptr,
                staging.dev_ptr,
                nbytes,
                stream.cu_stream(),
            )
            .context("dqn_head upload DtoD")?;
        }
    }
    Ok(dst)
}
