//! IQN (Implicit Quantile Network) distributional Q-head for the
//! integrated RL trainer — complementary to the existing C51 head.
//!
//! Sits on the shared Mamba2+CfC encoder's `h_t` output alongside
//! the C51 head. Produces per-action quantile values of shape
//! `[B, N_TAU, N_ACTIONS]` for `N_TAU` quantile fractions sampled
//! from U(0,1). The ensemble action selection combines both heads:
//!
//! ```text
//!   E_Q = α × E_C51 + (1-α) × E_IQN
//! ```
//!
//! where α is read from `ISV[RL_IQN_ENSEMBLE_ALPHA_INDEX=544]`.
//!
//! ## Architecture (cuBLAS-accelerated pipeline)
//!
//! The forward pass is split into custom kernels + cuBLAS SGEMMs:
//!
//! ```text
//!   1. rl_iqn_tau_cos_features (custom):
//!      Sample tau ~ U(0,1), compute cos_features[j] = cos((j+1)*pi*tau)
//!      → [B*N_TAU, EMBED_DIM]
//!
//!   2. cuBLAS SGEMM:
//!      embed_out = cos_features @ W_embed → [B*N_TAU, HIDDEN_DIM]
//!
//!   3. rl_iqn_relu_hadamard (custom):
//!      combined = h_t ⊙ ReLU(embed_out + b_embed) → [B*N_TAU, HIDDEN_DIM]
//!
//!   4. cuBLAS SGEMM:
//!      q_raw = combined @ W_out → [B*N_TAU, N_ACTIONS]
//!
//!   5. rl_iqn_bias_add_q (custom):
//!      q_values = q_raw + b_out → [B*N_TAU, N_ACTIONS]
//! ```
//!
//! The backward pass uses one cuBLAS SGEMM for the largest matmul:
//!
//! ```text
//!   1. cuBLAS SGEMM:
//!      grad_combined = grad_q @ W_out^T → [B*N_TAU, HIDDEN_DIM]
//!
//!   2. rl_iqn_backward (custom):
//!      Per-batch backward accumulation. Recomputes phi(tau) and combined
//!      inline from (tau, W_embed, b_embed, h_t). Reads grad_combined
//!      from cuBLAS. Produces per-batch grad_w_out/b_out/w_embed/b_embed.
//! ```
//!
//! ## Loss
//!
//! Quantile Huber loss (Dabney et al. 2018):
//!
//! ```text
//!   δ_{ij} = target_q[j,a*] - online_q[i,a*]
//!   ρ_{τ_i}(δ_{ij}) = |τ_i - 1(δ<0)| × Huber(δ, κ=1.0)
//!   L = (1/N_TAU²) Σ_i Σ_j ρ_{τ_i}(δ_{ij})
//! ```
//!
//! ## Initialisation
//!
//! Per `pearl_scoped_init_seed_for_reproducibility`, `IqnHead::new`
//! installs a `scoped_init_seed` guard before drawing Xavier samples.
//! Xavier uniform scaled by 0.01 (same as C51 head) so initial Q
//! outputs sit near zero.
//!
//! ## Constraints honoured
//!
//!   * `feedback_no_atomicadd.md` — loss kernel uses block tree-reduce.
//!   * `feedback_no_htod_htoh_only_mapped_pinned.md` — weight uploads
//!     stage through `MappedF32Buffer`.
//!   * `feedback_no_nvrtc.md` — pre-compiled cubins via `build.rs`.
//!   * `feedback_isv_for_adaptive_bounds.md` — N_TAU, ensemble α, and LR
//!     live in ISV slots 543-545.
//!   * `feedback_cpu_is_read_only.md` — forward/backward are GPU kernels
//!     + cuBLAS SGEMMs.

use std::sync::Arc;

use anyhow::{Context, Result};
use cudarc::cublas::CudaBlas;
use cudarc::cublas::sys::{self as cublas_sys, cublasOperation_t};
use cudarc::driver::{
    CudaFunction, CudaModule, CudaSlice, CudaStream,
};
use cudarc::driver::sys::CUstream;
use ml_core::cuda_autograd::init::scoped_init_seed;
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;
use crate::trainer::raw_launch::{RawArgs, raw_launch};

/// Quantile embedding dimension — number of cosine basis functions.
/// Matches the `EMBED_DIM` define in `cuda/rl_iqn_forward.cu`.
pub const EMBED_DIM: usize = 64;

/// Default number of quantile samples per forward pass. ISV-overridable
/// via `RL_IQN_N_TAU_INDEX`. Matches the `N_TAU` default in the kernel.
const DEFAULT_N_TAU: usize = 32;

/// cuBLAS workspace size — 8 MiB, prevents internal cudaMalloc during
/// CUDA Graph stream capture.
const CUBLAS_WORKSPACE_BYTES: usize = 8 * 1024 * 1024;

const IQN_FWD_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/rl_iqn_forward.cubin"
));
const IQN_LOSS_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/rl_iqn_loss.cubin"
));
const IQN_BWD_CUBIN: &[u8] = include_bytes!(concat!(
    env!("OUT_DIR"),
    "/rl_iqn_backward.cubin"
));

/// Construction config for [`IqnHead`].
#[derive(Clone, Debug)]
pub struct IqnHeadConfig {
    /// Encoder hidden dimension `h_t [B, HIDDEN_DIM]` feeding the head.
    pub hidden_dim: usize,
    /// Random seed for Xavier init. Reproducibility-critical per
    /// `pearl_scoped_init_seed_for_reproducibility`.
    pub seed: u64,
    /// Number of quantile samples per forward pass. Defaults to
    /// `DEFAULT_N_TAU` (32); overridden at runtime by
    /// `ISV[RL_IQN_N_TAU_INDEX]` in the training loop.
    pub n_tau: usize,
    /// Maximum batch size for pre-allocated scratch buffers.
    /// Must be >= any b_size passed to forward/backward.
    pub max_batch_size: usize,
}

impl Default for IqnHeadConfig {
    fn default() -> Self {
        Self {
            hidden_dim: HIDDEN_DIM,
            seed: 0x19A,
            n_tau: DEFAULT_N_TAU,
            max_batch_size: 256,
        }
    }
}

// ── raw pointer helpers (same pattern as ml-core linear.rs) ──────────

/// F32 × F32 → F32 GEMM via `cublasGemmEx` with F32 internal accumulation.
///
/// Uses `CUBLAS_GEMM_DFALT` algo — deterministic, safe for CUDA Graph capture.
///
/// # Safety
/// All device pointers must be valid and dimensions must be correct.
unsafe fn gemm_ex_f32(
    cublas: &CudaBlas,
    transa: cublasOperation_t,
    transb: cublasOperation_t,
    m: i32,
    n: i32,
    k: i32,
    a_ptr: u64,
    lda: i32,
    b_ptr: u64,
    ldb: i32,
    c_ptr: u64,
    ldc: i32,
    label: &str,
) -> Result<()> {
    let alpha = 1.0_f32;
    let beta = 0.0_f32;
    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 IQN {label}: {e:?}"))?;
    Ok(())
}

/// IQN distributional Q-head with cuBLAS-accelerated matmuls.
///
/// Owns device weights for the quantile embedding (`w_embed`, `b_embed`)
/// and output projection (`w_out`, `b_out`), plus a target network copy
/// for Bellman bootstrapping.
pub struct IqnHead {
    cfg: IqnHeadConfig,
    stream: Arc<CudaStream>,
    raw_stream: CUstream,

    // ── cuBLAS ───────────────────────────────────────────────────────
    cublas: CudaBlas,
    _cublas_workspace: CudaSlice<u8>,

    // ── Forward kernels (split pipeline) ─────────────────────────────
    _fwd_module: Arc<CudaModule>,
    /// Stage 1: tau sampling + cosine basis computation. READ-ONLY on
    /// `prng_state` per determinism Phase 2.6 fix (2026-06-02).
    pub tau_cos_features_fn: CudaFunction,
    /// Determinism Phase 2.6 sibling: advance the per-batch xorshift32
    /// PRNG state by 8 steps. Run AFTER `tau_cos_features_fn` so the
    /// sampling sees a stable read across all (tau_idx) blocks. The
    /// kernel-launch ordering on a single stream is the grid-wide
    /// barrier (replaces the racy "tau_idx==0 block writes prng_state
    /// at end of monolithic kernel" pattern that produced run-dependent
    /// τ values for tau_idx > 0 pre-fix).
    pub advance_prng_state_fn: CudaFunction,
    /// Stage 3: bias-add → ReLU → hadamard product with h_t.
    pub relu_hadamard_fn: CudaFunction,
    /// Stage 5: bias addition to cuBLAS output projection result.
    pub bias_add_q_fn: CudaFunction,
    /// Expected-Q reduction: mean over tau dimension.
    pub expected_q_fn: CudaFunction,

    _loss_module: Arc<CudaModule>,
    pub loss_fn: CudaFunction,

    // ── Backward kernels ─────────────────────────────────────────────
    _bwd_module: Arc<CudaModule>,
    /// Per-batch backward accumulation with cuBLAS-precomputed grad_combined.
    pub bwd_fn: CudaFunction,
    /// Element-wise backward through hadamard + ReLU (for cached forward path).
    pub bwd_relu_hadamard_fn: CudaFunction,

    // ── Online weights ───────────────────────────────────────────────
    pub w_embed_d: CudaSlice<f32>,
    pub b_embed_d: CudaSlice<f32>,
    pub w_out_d: CudaSlice<f32>,
    pub b_out_d: CudaSlice<f32>,

    // ── Target-network weights ───────────────────────────────────────
    pub w_embed_target_d: CudaSlice<f32>,
    pub b_embed_target_d: CudaSlice<f32>,
    pub w_out_target_d: CudaSlice<f32>,
    pub b_out_target_d: CudaSlice<f32>,

    // ── Pre-allocated scratch for forward/backward (graph-capture safe) ─
    scratch_cos_features: CudaSlice<f32>,
    scratch_embed_out: CudaSlice<f32>,
    scratch_combined: CudaSlice<f32>,
    scratch_embed_pre_relu: CudaSlice<f32>,
    scratch_grad_combined: CudaSlice<f32>,
}

impl IqnHead {
    /// Allocate device weights, load cubins, and cache kernel handles.
    pub fn new(dev: &MlDevice, cfg: IqnHeadConfig) -> Result<Self> {
        let stream: Arc<CudaStream> = dev.cuda_stream().context("iqn_head stream")?.clone();
        let ctx = dev.cuda_context().context("iqn_head ctx")?;

        // ── cuBLAS handle + pre-allocated workspace ──────────────────
        let cublas = CudaBlas::new(Arc::clone(&stream))
            .map_err(|e| anyhow::anyhow!("IqnHead: cuBLAS init: {e}"))?;
        // Determinism foundation (spec 2026-06-02 §2.B): disable
        // non-deterministic GEMM algorithms. Same fix as Mamba2Block,
        // applied to the IQN τ-conditioned GEMMs.
        let cublas = crate::cublas_determinism::apply_deterministic_math_mode(cublas)
            .map_err(|e| anyhow::anyhow!("IqnHead: cublas determinism setup: {e}"))?;
        let cublas_workspace = stream
            .alloc_zeros::<u8>(CUBLAS_WORKSPACE_BYTES)
            .map_err(|e| anyhow::anyhow!("IqnHead: cuBLAS workspace: {e}"))?;
        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!("IqnHead: cublasSetWorkspace_v2: {e:?}"))?;
        }

        // ── Forward kernel symbols ───────────────────────────────────
        let fwd_module = ctx
            .load_cubin(IQN_FWD_CUBIN.to_vec())
            .context("load rl_iqn_forward cubin")?;
        let tau_cos_features_fn = fwd_module
            .load_function("rl_iqn_tau_cos_features")
            .context("load rl_iqn_tau_cos_features")?;
        // Determinism Phase 2.6 (2026-06-02): companion advance kernel
        // for the now read-only `rl_iqn_tau_cos_features`. Same cubin.
        let advance_prng_state_fn = fwd_module
            .load_function("rl_iqn_advance_prng_state")
            .context("load rl_iqn_advance_prng_state")?;
        let relu_hadamard_fn = fwd_module
            .load_function("rl_iqn_relu_hadamard")
            .context("load rl_iqn_relu_hadamard")?;
        let bias_add_q_fn = fwd_module
            .load_function("rl_iqn_bias_add_q")
            .context("load rl_iqn_bias_add_q")?;
        let expected_q_fn = fwd_module
            .load_function("rl_iqn_expected_q")
            .context("load rl_iqn_expected_q")?;

        let loss_module = ctx
            .load_cubin(IQN_LOSS_CUBIN.to_vec())
            .context("load rl_iqn_loss cubin")?;
        let loss_fn = loss_module
            .load_function("rl_iqn_loss_fwd")
            .context("load rl_iqn_loss_fwd")?;

        // ── Backward kernel symbols ──────────────────────────────────
        let bwd_module = ctx
            .load_cubin(IQN_BWD_CUBIN.to_vec())
            .context("load rl_iqn_backward cubin")?;
        let bwd_fn = bwd_module
            .load_function("rl_iqn_backward")
            .context("load rl_iqn_backward")?;
        let bwd_relu_hadamard_fn = bwd_module
            .load_function("rl_iqn_bwd_relu_hadamard")
            .context("load rl_iqn_bwd_relu_hadamard")?;

        // ── Weight initialisation ────────────────────────────────────
        let _seed_guard = scoped_init_seed(cfg.seed);
        let mut rng = ChaCha8Rng::seed_from_u64(cfg.seed);

        let embed_scale =
            0.01_f32 * (6.0_f32 / (EMBED_DIM + cfg.hidden_dim) as f32).sqrt();
        let w_embed_host: Vec<f32> = (0..EMBED_DIM * cfg.hidden_dim)
            .map(|_| rng.gen_range(-embed_scale..embed_scale))
            .collect();
        let b_embed_host: Vec<f32> = vec![0.0_f32; cfg.hidden_dim];

        let out_scale =
            0.01_f32 * (6.0_f32 / (cfg.hidden_dim + N_ACTIONS) as f32).sqrt();
        let w_out_host: Vec<f32> = (0..cfg.hidden_dim * N_ACTIONS)
            .map(|_| rng.gen_range(-out_scale..out_scale))
            .collect();
        let b_out_host: Vec<f32> = vec![0.0_f32; N_ACTIONS];

        let w_embed_d = upload(&stream, &w_embed_host)?;
        let b_embed_d = upload(&stream, &b_embed_host)?;
        let w_out_d = upload(&stream, &w_out_host)?;
        let b_out_d = upload(&stream, &b_out_host)?;

        let w_embed_target_d = upload(&stream, &w_embed_host)?;
        let b_embed_target_d = upload(&stream, &b_embed_host)?;
        let w_out_target_d = upload(&stream, &w_out_host)?;
        let b_out_target_d = upload(&stream, &b_out_host)?;

        let raw_stream = stream.cu_stream();

        let max_m = cfg.max_batch_size * cfg.n_tau;
        let hd = cfg.hidden_dim;
        let scratch_cos_features = stream.alloc_zeros::<f32>(max_m * EMBED_DIM)
            .context("iqn scratch_cos_features")?;
        let scratch_embed_out = stream.alloc_zeros::<f32>(max_m * hd)
            .context("iqn scratch_embed_out")?;
        let scratch_combined = stream.alloc_zeros::<f32>(max_m * hd)
            .context("iqn scratch_combined")?;
        let scratch_embed_pre_relu = stream.alloc_zeros::<f32>(max_m * hd)
            .context("iqn scratch_embed_pre_relu")?;
        let scratch_grad_combined = stream.alloc_zeros::<f32>(max_m * hd)
            .context("iqn scratch_grad_combined")?;

        Ok(Self {
            cfg,
            stream,
            raw_stream,
            cublas,
            _cublas_workspace: cublas_workspace,
            _fwd_module: fwd_module,
            tau_cos_features_fn,
            advance_prng_state_fn,
            relu_hadamard_fn,
            bias_add_q_fn,
            expected_q_fn,
            _loss_module: loss_module,
            loss_fn,
            _bwd_module: bwd_module,
            bwd_fn,
            bwd_relu_hadamard_fn,
            w_embed_d,
            b_embed_d,
            w_out_d,
            b_out_d,
            w_embed_target_d,
            b_embed_target_d,
            w_out_target_d,
            b_out_target_d,
            scratch_cos_features,
            scratch_embed_out,
            scratch_combined,
            scratch_embed_pre_relu,
            scratch_grad_combined,
        })
    }

    // ── Forward ──────────────────────────────────────────────────────

    /// cuBLAS-accelerated forward pass with online weights.
    pub fn forward(
        &self,
        prng_state: &mut CudaSlice<u32>,
        h_t: &CudaSlice<f32>,
        tau: &mut CudaSlice<f32>,
        b_size: usize,
        n_tau: usize,
        q_values_out: &mut CudaSlice<f32>,
    ) -> Result<()> {
        self.forward_inner(
            prng_state, h_t, tau,
            &self.w_embed_d, &self.b_embed_d,
            &self.w_out_d, &self.b_out_d,
            b_size, n_tau, q_values_out,
        )
    }

    /// cuBLAS-accelerated forward pass with target weights.
    pub fn forward_target(
        &self,
        prng_state: &mut CudaSlice<u32>,
        h_t: &CudaSlice<f32>,
        tau: &mut CudaSlice<f32>,
        b_size: usize,
        n_tau: usize,
        q_values_out: &mut CudaSlice<f32>,
    ) -> Result<()> {
        self.forward_inner(
            prng_state, h_t, tau,
            &self.w_embed_target_d, &self.b_embed_target_d,
            &self.w_out_target_d, &self.b_out_target_d,
            b_size, n_tau, q_values_out,
        )
    }

    /// Internal forward — parameterised over weight slices.
    ///
    /// Pipeline:
    ///   1. `rl_iqn_tau_cos_features` — tau + cos basis
    ///   2. cuBLAS SGEMM: `cos_features @ W_embed` → `embed_out`
    ///   3. `rl_iqn_relu_hadamard` — bias + ReLU + hadamard
    ///   4. cuBLAS SGEMM: `combined @ W_out` → `q_raw`
    ///   5. `rl_iqn_bias_add_q` — add b_out
    #[allow(clippy::too_many_arguments)]
    fn forward_inner(
        &self,
        prng_state: &mut CudaSlice<u32>,
        h_t: &CudaSlice<f32>,
        tau: &mut CudaSlice<f32>,
        w_embed: &CudaSlice<f32>,
        b_embed: &CudaSlice<f32>,
        w_out: &CudaSlice<f32>,
        b_out: &CudaSlice<f32>,
        b_size: usize,
        n_tau: usize,
        q_values_out: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(prng_state.len(), b_size);
        debug_assert_eq!(h_t.len(), b_size * self.cfg.hidden_dim);
        debug_assert_eq!(tau.len(), b_size * n_tau);
        debug_assert_eq!(q_values_out.len(), b_size * n_tau * N_ACTIONS);

        let m = b_size * n_tau;
        let hd = self.cfg.hidden_dim;

        // ── Stage 1: tau sampling + cos features ─────────────────────
        debug_assert!(m * EMBED_DIM <= self.scratch_cos_features.len(),
            "IQN scratch overflow: m={m} > max_m={}", self.scratch_cos_features.len() / EMBED_DIM);
        let cos_ptr = self.scratch_cos_features.raw_ptr();
        let embed_ptr = self.scratch_embed_out.raw_ptr();
        let comb_ptr = self.scratch_combined.raw_ptr();
        let pre_relu_ptr = self.scratch_embed_pre_relu.raw_ptr();
        {
            let mut args = RawArgs::new();
            args.push_ptr(prng_state.raw_ptr());
            args.push_i32(b_size as i32);
            args.push_i32(n_tau as i32);
            args.push_ptr(tau.raw_ptr());
            args.push_ptr(cos_ptr);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.tau_cos_features_fn.cu_function(),
                    (b_size as u32, n_tau as u32, 1),
                    (EMBED_DIM as u32, 1, 1),
                    std::mem::size_of::<f32>() as u32,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                ).map_err(|e| anyhow::anyhow!("rl_iqn_tau_cos_features: {:?}", e))?;
            }
        }

        // ── Stage 1b: advance PRNG state (determinism Phase 2.6 fix).
        // The previous monolithic `rl_iqn_tau_cos_features` had a read-
        // write race: all N_TAU blocks per batch read `prng_state[batch]`
        // at kernel head, and the (tau_idx==0) block wrote back the
        // advanced state at kernel tail — with NO inter-block barrier.
        // Run on the same stream → kernel-launch ordering acts as a
        // grid-wide barrier between the (now read-only) sampling and
        // the (single-thread-per-batch) advance. No race possible.
        // Per `feedback_no_atomicadd.md` + canonical Phase-2 rebuild rule.
        {
            let mut args = RawArgs::new();
            args.push_ptr(prng_state.raw_ptr());
            args.push_i32(b_size as i32);
            let mut ptrs = args.build_arg_ptrs();
            let block_x = 256u32.min(b_size as u32).max(1);
            let grid_x = ((b_size as u32) + block_x - 1) / block_x;
            unsafe {
                raw_launch(
                    self.advance_prng_state_fn.cu_function(),
                    (grid_x, 1, 1),
                    (block_x, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                ).map_err(|e| anyhow::anyhow!("rl_iqn_advance_prng_state: {:?}", e))?;
            }
        }

        // ── Stage 2: cuBLAS SGEMM — embed_out = cos_features @ W_embed ──
        {
            let w_ptr = w_embed.raw_ptr();
            unsafe {
                gemm_ex_f32(
                    &self.cublas,
                    cublasOperation_t::CUBLAS_OP_N,
                    cublasOperation_t::CUBLAS_OP_N,
                    hd as i32, m as i32, EMBED_DIM as i32,
                    w_ptr, hd as i32,
                    cos_ptr, EMBED_DIM as i32,
                    embed_ptr, hd as i32,
                    "embed_fwd",
                )?;
            }
        }

        // ── Stage 3: ReLU + hadamard ─────────────────────────────────
        {
            let block_x = 128u32.min(hd as u32);
            let grid_y = (hd as u32 + block_x - 1) / block_x;
            let mut args = RawArgs::new();
            args.push_ptr(embed_ptr);
            args.push_ptr(b_embed.raw_ptr());
            args.push_ptr(h_t.raw_ptr());
            args.push_i32(m as i32);
            args.push_i32(b_size as i32);
            args.push_i32(n_tau as i32);
            args.push_ptr(comb_ptr);
            args.push_ptr(pre_relu_ptr);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.relu_hadamard_fn.cu_function(),
                    (m as u32, grid_y, 1),
                    (block_x, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                ).map_err(|e| anyhow::anyhow!("rl_iqn_relu_hadamard: {:?}", e))?;
            }
        }

        // ── Stage 4: cuBLAS SGEMM — q_raw = combined @ W_out ────────
        {
            let wo_ptr = w_out.raw_ptr();
            let q_ptr = q_values_out.raw_ptr();
            unsafe {
                gemm_ex_f32(
                    &self.cublas,
                    cublasOperation_t::CUBLAS_OP_N,
                    cublasOperation_t::CUBLAS_OP_N,
                    N_ACTIONS as i32, m as i32, hd as i32,
                    wo_ptr, N_ACTIONS as i32,
                    comb_ptr, hd as i32,
                    q_ptr, N_ACTIONS as i32,
                    "out_proj_fwd",
                )?;
            }
        }

        // ── Stage 5: bias add ────────────────────────────────────────
        {
            let mut args = RawArgs::new();
            args.push_ptr(q_values_out.raw_ptr());
            args.push_ptr(b_out.raw_ptr());
            args.push_i32(m as i32);
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.bias_add_q_fn.cu_function(),
                    (m as u32, 1, 1),
                    (N_ACTIONS as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                ).map_err(|e| anyhow::anyhow!("rl_iqn_bias_add_q: {:?}", e))?;
            }
        }

        Ok(())
    }

    /// Compute expected Q values: `E_IQN[b,a] = mean_τ Q(s,τ,a)`.
    pub fn expected_q(
        &self,
        q_values: &CudaSlice<f32>,
        b_size: usize,
        n_tau: usize,
        expected_q_out: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(q_values.len(), b_size * n_tau * N_ACTIONS);
        debug_assert_eq!(expected_q_out.len(), b_size * N_ACTIONS);

        let mut args = RawArgs::new();
        args.push_ptr(q_values.raw_ptr());
        args.push_i32(b_size as i32);
        args.push_i32(n_tau as i32);
        args.push_ptr(expected_q_out.raw_ptr());
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.expected_q_fn.cu_function(),
                (b_size as u32, 1, 1),
                (N_ACTIONS as u32, 1, 1),
                0,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("rl_iqn_expected_q: {:?}", e))?;
        }
        Ok(())
    }

    /// Quantile Huber loss forward + backward for the taken action.
    #[allow(clippy::too_many_arguments)]
    pub fn compute_loss(
        &self,
        online_q: &CudaSlice<f32>,
        target_q: &CudaSlice<f32>,
        tau_online: &CudaSlice<f32>,
        actions_taken: &CudaSlice<i32>,
        b_size: usize,
        n_tau_online: usize,
        n_tau_target: usize,
        loss_per_batch: &mut CudaSlice<f32>,
        grad_online_q: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(online_q.len(), b_size * n_tau_online * N_ACTIONS);
        debug_assert_eq!(target_q.len(), b_size * n_tau_target * N_ACTIONS);
        debug_assert_eq!(tau_online.len(), b_size * n_tau_online);
        debug_assert_eq!(actions_taken.len(), b_size);
        debug_assert_eq!(loss_per_batch.len(), b_size);
        debug_assert_eq!(grad_online_q.len(), b_size * n_tau_online * N_ACTIONS);

        let mut args = RawArgs::new();
        args.push_ptr(online_q.raw_ptr());
        args.push_ptr(target_q.raw_ptr());
        args.push_ptr(tau_online.raw_ptr());
        args.push_ptr(actions_taken.raw_ptr());
        args.push_i32(b_size as i32);
        args.push_i32(n_tau_online as i32);
        args.push_i32(n_tau_target as i32);
        args.push_ptr(loss_per_batch.raw_ptr());
        args.push_ptr(grad_online_q.raw_ptr());
        let mut ptrs = args.build_arg_ptrs();
        unsafe {
            raw_launch(
                self.loss_fn.cu_function(),
                (b_size as u32, 1, 1),
                (n_tau_online as u32, 1, 1),
                (n_tau_online * std::mem::size_of::<f32>()) as u32,
                self.raw_stream,
                &mut ptrs[..args.len()],
            ).map_err(|e| anyhow::anyhow!("rl_iqn_loss_fwd: {:?}", e))?;
        }
        Ok(())
    }

    // ── Backward ─────────────────────────────────────────────────────

    /// cuBLAS-accelerated backward through the forward pass.
    ///
    /// Pipeline:
    ///   1. cuBLAS SGEMM: `grad_combined = grad_q @ W_out^T`
    ///   2. `rl_iqn_backward`: per-batch accumulation using cuBLAS-
    ///      precomputed `grad_combined` (recomputes phi/combined inline
    ///      from tau + weights, same as original kernel).
    #[allow(clippy::too_many_arguments)]
    pub fn backward(
        &self,
        h_t: &CudaSlice<f32>,
        tau: &CudaSlice<f32>,
        grad_output: &CudaSlice<f32>,
        b_size: usize,
        n_tau: usize,
        grad_w_out_pb: &mut CudaSlice<f32>,
        grad_b_out_pb: &mut CudaSlice<f32>,
        grad_w_embed_pb: &mut CudaSlice<f32>,
        grad_b_embed_pb: &mut CudaSlice<f32>,
    ) -> Result<()> {
        debug_assert_eq!(h_t.len(), b_size * self.cfg.hidden_dim);
        debug_assert_eq!(tau.len(), b_size * n_tau);
        debug_assert_eq!(grad_output.len(), b_size * n_tau * N_ACTIONS);
        debug_assert_eq!(grad_w_out_pb.len(), b_size * self.cfg.hidden_dim * N_ACTIONS);
        debug_assert_eq!(grad_b_out_pb.len(), b_size * N_ACTIONS);
        debug_assert_eq!(grad_w_embed_pb.len(), b_size * EMBED_DIM * self.cfg.hidden_dim);
        debug_assert_eq!(grad_b_embed_pb.len(), b_size * self.cfg.hidden_dim);

        let m = b_size * n_tau;
        let hd = self.cfg.hidden_dim;

        // ── cuBLAS: grad_combined = grad_q @ W_out^T ─────────────────
        // Row-major: C[M, hd] = A[M, N_ACTIONS] @ B^T[N_ACTIONS, hd]
        // cuBLAS col-major: C_col[hd, M] = W_out_col^T[hd, N_ACTIONS] @ grad_q_col[N_ACTIONS, M]
        //   W_out is row-major [hd, N_ACTIONS] = col-major [N_ACTIONS, hd].
        //   transA=T on col-major [N_ACTIONS, hd] → [hd, N_ACTIONS].
        //   transB=N on grad_q col-major [N_ACTIONS, M].
        //   m=hd, n=M, k=N_ACTIONS, lda=N_ACTIONS, ldb=N_ACTIONS, ldc=hd
        debug_assert!(m * hd <= self.scratch_grad_combined.len());
        let gc_ptr = self.scratch_grad_combined.raw_ptr();
        {
            let wo_ptr = self.w_out_d.raw_ptr();
            let gq_ptr = grad_output.raw_ptr();
            unsafe {
                gemm_ex_f32(
                    &self.cublas,
                    cublasOperation_t::CUBLAS_OP_T,
                    cublasOperation_t::CUBLAS_OP_N,
                    hd as i32, m as i32, N_ACTIONS as i32,
                    wo_ptr, N_ACTIONS as i32,
                    gq_ptr, N_ACTIONS as i32,
                    gc_ptr, hd as i32,
                    "grad_combined_bwd",
                )?;
            }
        }

        // ── Per-batch backward kernel ────────────────────────────────
        {
            let mut args = RawArgs::new();
            args.push_ptr(h_t.raw_ptr());
            args.push_ptr(tau.raw_ptr());
            args.push_ptr(self.w_embed_d.raw_ptr());
            args.push_ptr(self.b_embed_d.raw_ptr());
            args.push_ptr(gc_ptr);
            args.push_ptr(grad_output.raw_ptr());
            args.push_i32(b_size as i32);
            args.push_i32(n_tau as i32);
            args.push_ptr(grad_w_out_pb.raw_ptr());
            args.push_ptr(grad_b_out_pb.raw_ptr());
            args.push_ptr(grad_w_embed_pb.raw_ptr());
            args.push_ptr(grad_b_embed_pb.raw_ptr());
            let mut ptrs = args.build_arg_ptrs();
            unsafe {
                raw_launch(
                    self.bwd_fn.cu_function(),
                    (b_size as u32, 1, 1),
                    (self.cfg.hidden_dim as u32, 1, 1),
                    0,
                    self.raw_stream,
                    &mut ptrs[..args.len()],
                ).map_err(|e| anyhow::anyhow!("rl_iqn_backward: {:?}", e))?;
            }
        }
        Ok(())
    }

    /// Stream used to launch all kernels owned by this head.
    pub fn stream(&self) -> &Arc<CudaStream> {
        &self.stream
    }

    /// Encoder hidden dim this head was built for.
    pub fn hidden_dim(&self) -> usize {
        self.cfg.hidden_dim
    }

    /// Number of quantile samples configured at construction.
    pub fn n_tau(&self) -> usize {
        self.cfg.n_tau
    }
}

// ── pinned-staging upload helper (mirrors dqn.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!("iqn_head upload staging: {e}"))?;
    staging.write_from_slice(host);
    let dst = stream
        .alloc_zeros::<f32>(n)
        .context("iqn_head upload alloc")?;
    if n > 0 {
        let nbytes = n * std::mem::size_of::<f32>();
        unsafe {
            let dst_ptr = dst.raw_ptr();
            crate::trainer::raw_launch::raw_memcpy_dtod_async(
                dst_ptr,
                staging.dev_ptr,
                nbytes,
                stream.cu_stream(),
            )
            .map_err(|e| anyhow::anyhow!("iqn_head upload DtoD: {e:?}"))?;
        }
    }
    Ok(dst)
}
