#![allow(unsafe_code)]

//! cuBLAS SGEMM-based batched backward pass for the DQN trainer.
//!
//! Replaces the 1-warp-per-sample `dqn_backward_kernel` with cuBLAS matrix
//! multiplications that process the entire batch in a single GEMM call per
//! layer.  This eliminates `atomicAdd` serialization on the 288K gradient
//! buffer and raises GPU occupancy from 1.56% to >60% on H100.
//!
//! ## Layout convention
//!
//! All tensors use row-major (C-style) layout, same as the forward pass:
//! - `dY`:  [B, out_dim]  — upstream gradient
//! - `X`:   [B, in_dim]   — saved activation (input to this layer)
//! - `W`:   [out_dim, in_dim] — weight matrix (row-major)
//! - `dW`:  [out_dim, in_dim] — weight gradient accumulator
//!
//! ## cuBLAS row-major → col-major mapping
//!
//! ### Weight gradient  `dW[out, in] = dY^T[out, B] @ X[B, in]`
//!
//! cuBLAS works column-major.  We observe the dual:
//! - `dW'[in, out]`  col-major  = `X'[in, B]` col-major  @  `dY'^T[B, out]`
//! - Row-major `dW[out, in]` shares the same memory as col-major `dW'[in, out]`.
//! - Row-major `X[B, in]`   shares memory with col-major `X'[in, B]`.
//! - Row-major `dY[B, out]` shares memory with col-major `dY'[out, B]`.
//!
//! Therefore:
//! ```text
//! sgemm(N, T, m=in_dim, n=out_dim, k=B,
//!        1.0, X, lda=in_dim,
//!        dY,  ldb=out_dim,
//!        1.0, dW, ldc=in_dim)
//! ```
//! The `beta=1.0` means ACCUMULATE into `dW` (not overwrite), which is required
//! since multiple layers share the same flat `grad_buf`.
//!
//! ### Upstream gradient  `dX[B, in] = dY[B, out] @ W[out, in]`
//!
//! ```text
//! sgemm(N, N, m=in_dim, n=B, k=out_dim,
//!        1.0, W,  lda=in_dim,
//!        dY,  ldb=out_dim,
//!        0.0, dX, ldc=in_dim)
//! ```
//! Here `beta=0.0` because `dX` is a fresh scratch buffer per layer.
//!
//! ## Bias gradient
//!
//! `db[out] = sum_b(dY[b, out])` — a row-wise reduction.
//! Implemented as a custom NVRTC kernel (`bias_grad_reduce_kernel`):
//! each thread owns one output neuron and sums over the batch dimension.
//! Uses `atomicAdd` into `grad_buf[bias_offset + j]` so it accumulates
//! safely alongside the weight gradient GEMM.
//!
//! ## ReLU derivative mask
//!
//! After computing `dX = W @ dY`, gate by saved pre-ReLU activation:
//! `dX[i] *= (activation[i] > 0)`.
//! Implemented as `relu_mask_kernel` (1 thread per element).
//!
//! ## Phase status
//!
//! This module provides the building blocks for Phase 2 Task 2.  The full
//! cuBLAS backward path requires a standalone C51 loss kernel that emits
//! `dL/d_logits` for each branch.  Until that kernel exists, the existing
//! `dqn_backward_kernel` (atomicAdd path) remains active.  The functions
//! here are wired via `GpuDqnTrainer::launch_cublas_backward()` with a
//! TODO gate that enables them once the C51 loss kernel is ready.

use std::mem::ManuallyDrop;
use std::sync::Arc;

use cudarc::cublas::result as cublas_result;
use cudarc::cublas::sys as cublas_sys;
use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg};

use crate::MLError;
use super::gpu_dqn_trainer::GpuDqnTrainConfig;

// ── cuBLAS handle wrapper ─────────────────────────────────────────────────────

/// Wrapper to make `cublasHandle_t` Send + Sync.
///
/// Safety reasoning is identical to `CublasForward`: the handle is owned
/// exclusively by `CublasBackward`, which is in turn owned by `GpuDqnTrainer`
/// — a single-threaded owner in practice.
struct SendSyncCublasHandle(cublas_sys::cublasHandle_t);

unsafe impl Send for SendSyncCublasHandle {}
unsafe impl Sync for SendSyncCublasHandle {}

impl Drop for SendSyncCublasHandle {
    fn drop(&mut self) {
        unsafe {
            let _ = cublas_result::destroy_handle(self.0);
        }
    }
}

// ── CublasBackward context ────────────────────────────────────────────────────

/// cuBLAS-based batched backward pass context for the branching dueling DQN.
///
/// Holds the cuBLAS handle bound to the trainer's stream plus two small NVRTC
/// kernels for the operations that cuBLAS cannot express directly (bias
/// gradient reduction and ReLU derivative masking).
///
/// ## Usage
///
/// 1. Call `backward_fc_layer()` in reverse layer order, passing the upstream
///    gradient `dY`, the saved forward activation `X`, the weight matrix `W`,
///    and gradient accumulator offsets into the flat `grad_buf`.
/// 2. After each weight layer (except the first), call `relu_mask()` to gate
///    the computed upstream gradient by the saved pre-ReLU activation.
/// 3. For the final full backward, call `backward_full()` which orchestrates
///    all layers automatically.
#[allow(missing_debug_implementations)] // handle is not Debug
pub struct CublasBackward {
    handle: SendSyncCublasHandle,

    /// `relu_mask_kernel(dx, activation, n)` — element-wise ReLU derivative gate.
    relu_mask_kernel: CudaFunction,

    /// `bias_grad_reduce_kernel(dy, db, out_dim, batch_size)` — reduce dY over batch.
    bias_grad_kernel: CudaFunction,

    // ── Network dimensions (baked at construction) ──
    batch_size: usize,
    state_dim: usize,
    shared_h1: usize,
    shared_h2: usize,
    value_h: usize,
    adv_h: usize,
    num_atoms: usize,
    branch_0_size: usize,
    branch_1_size: usize,
    branch_2_size: usize,
}

impl CublasBackward {
    /// Create a new cuBLAS backward context.
    ///
    /// Compiles the `relu_mask_kernel` and `bias_grad_reduce_kernel` from
    /// inline NVRTC source.  Binds the cuBLAS handle to the provided stream
    /// so all GEMMs execute in-order on the same stream as the forward pass.
    pub fn new(
        stream: &Arc<CudaStream>,
        config: &GpuDqnTrainConfig,
    ) -> Result<Self, MLError> {
        // ── Create cuBLAS handle ────────────────────────────────────
        let raw_handle = cublas_result::create_handle()
            .map_err(|e| MLError::ModelError(format!("cublasCreate (backward): {e:?}")))?;

        // Bind handle to stream — all backward GEMMs execute on this stream.
        // SAFETY: same cast as CublasForward::new — CUstream pointers are
        // identical across cudarc::driver::sys and cudarc::cublas::sys.
        unsafe {
            let cu_stream = stream.cu_stream() as *mut cublas_sys::CUstream_st;
            cublas_result::set_stream(raw_handle, cu_stream)
                .map_err(|e| MLError::ModelError(format!("cublasSetStream (backward): {e:?}")))?;
        }

        // ── Compile helper kernels ──────────────────────────────────
        let (relu_mask_kernel, bias_grad_kernel) = compile_backward_kernels(stream)?;

        Ok(Self {
            handle: SendSyncCublasHandle(raw_handle),
            relu_mask_kernel,
            bias_grad_kernel,
            batch_size: config.batch_size,
            state_dim: config.state_dim,
            shared_h1: config.shared_h1,
            shared_h2: config.shared_h2,
            value_h: config.value_h,
            adv_h: config.adv_h,
            num_atoms: config.num_atoms,
            branch_0_size: config.branch_0_size,
            branch_1_size: config.branch_1_size,
            branch_2_size: config.branch_2_size,
        })
    }

    /// Rebind the cuBLAS handle to a (potentially new) stream.
    ///
    /// Must be called before CUDA Graph capture when the stream changes,
    /// matching the convention in `CublasForward::set_stream`.
    pub fn set_stream(&self, stream: &CudaStream) -> Result<(), MLError> {
        unsafe {
            let cu_stream = stream.cu_stream() as *mut cublas_sys::CUstream_st;
            cublas_result::set_stream(self.handle.0, cu_stream)
                .map_err(|e| MLError::ModelError(format!("cublasSetStream (backward): {e:?}")))
        }
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Public building blocks
    // ══════════════════════════════════════════════════════════════════════════

    /// Backward through one fully-connected layer using cuBLAS SGEMM.
    ///
    /// Computes:
    /// - `dW += dY^T @ X`   (weight gradient, accumulated with beta=1.0)
    /// - `db += sum_b(dY)`  (bias gradient, accumulated via atomicAdd)
    /// - `dX  = dY  @ W^T`  (upstream gradient, overwritten with beta=0.0)
    ///
    /// This does NOT apply the ReLU derivative mask; call `relu_mask()` after
    /// if the layer had a ReLU activation.
    ///
    /// ## Arguments
    ///
    /// - `dy`      : raw u64 device pointer to `dY[B, out_dim]` (upstream gradient)
    /// - `x`       : raw u64 device pointer to `X[B, in_dim]`  (saved input activation)
    /// - `w`       : raw u64 device pointer to `W[out_dim, in_dim]` (weight matrix)
    /// - `dw`      : raw u64 device pointer to `dW[out_dim, in_dim]` in `grad_buf`
    /// - `db`      : raw u64 device pointer to `db[out_dim]` in `grad_buf`
    /// - `dx`      : raw u64 device pointer to `dX[B, in_dim]` scratch (may be 0
    ///               if upstream gradient is not needed, e.g. input layer)
    /// - `out_dim` : number of neurons in this layer's output
    /// - `in_dim`  : number of neurons in this layer's input
    /// - `batch`   : batch size B
    #[allow(clippy::too_many_arguments)]
    pub fn backward_fc_layer(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,
        x: u64,
        w: u64,
        dw: u64,
        db: u64,
        dx: u64,           // 0 means skip upstream gradient computation
        out_dim: usize,
        in_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        // ── Weight gradient: dW[out, in] += dY^T[out, B] @ X[B, in] ──
        //
        // Row-major dW[out,in] viewed col-major as dW'[in,out]:
        //   dW'[in,out] = X'[in,B] @ dY'^T[B,out]
        // => sgemm(N, T, m=in_dim, n=out_dim, k=B, X, lda=in_dim, dY, ldb=out_dim, dW, ldc=in_dim)
        // beta=1.0 to ACCUMULATE (shared grad_buf with bias kernel)
        unsafe {
            let alpha = 1.0_f32;
            let beta_acc = 1.0_f32; // accumulate
            cublas_result::sgemm(
                self.handle.0,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,  // no transpose on X'
                cublas_sys::cublasOperation_t::CUBLAS_OP_T,  // transpose dY'
                in_dim as i32,    // m: rows of op(X') = in_dim
                out_dim as i32,   // n: cols of op(dY'^T) = out_dim
                batch as i32,     // k: inner dim = B
                &alpha,
                x as *const f32,
                in_dim as i32,    // lda: leading dim of X' = in_dim
                dy as *const f32,
                out_dim as i32,   // ldb: leading dim of dY' = out_dim
                &beta_acc,
                dw as *mut f32,
                in_dim as i32,    // ldc: leading dim of dW' = in_dim
            )
            .map_err(|e| MLError::ModelError(format!("backward sgemm dW: {e:?}")))?;
        }

        // ── Bias gradient: db[out] += sum_b(dY[b, out]) ──────────────
        self.launch_bias_grad(stream, dy, db, out_dim, batch)?;

        // ── Upstream gradient: dX[B, in] = dY[B, out] @ W[out, in] ──
        // Skip when dx == 0 (input layer — states are not trainable).
        if dx != 0 {
            // Row-major dX[B,in] viewed col-major as dX'[in,B]:
            //   dX'[in,B] = W'[in,out] @ dY'[out,B]
            // => sgemm(N, N, m=in_dim, n=B, k=out_dim, W, lda=in_dim, dY, ldb=out_dim, dX, ldc=in_dim)
            // beta=0.0: overwrite dX scratch for this layer
            unsafe {
                let alpha = 1.0_f32;
                let beta_zero = 0.0_f32;
                cublas_result::sgemm(
                    self.handle.0,
                    cublas_sys::cublasOperation_t::CUBLAS_OP_N,  // no transpose on W'
                    cublas_sys::cublasOperation_t::CUBLAS_OP_N,  // no transpose on dY'
                    in_dim as i32,    // m: rows of W' = in_dim
                    batch as i32,     // n: cols of dY' = B
                    out_dim as i32,   // k: inner dim = out_dim
                    &alpha,
                    w as *const f32,
                    in_dim as i32,    // lda: leading dim of W' = in_dim
                    dy as *const f32,
                    out_dim as i32,   // ldb: leading dim of dY' = out_dim
                    &beta_zero,
                    dx as *mut f32,
                    in_dim as i32,    // ldc: leading dim of dX' = in_dim
                )
                .map_err(|e| MLError::ModelError(format!("backward sgemm dX: {e:?}")))?;
            }
        }

        Ok(())
    }

    /// Apply the ReLU derivative mask in-place: `dx[i] *= (activation[i] > 0)`.
    ///
    /// Called after `backward_fc_layer()` for layers that used ReLU activation.
    /// The `activation` argument is the saved post-ReLU value from the forward
    /// pass (i.e. `save_h_*`).  For ReLU, `activation[i] > 0` iff the pre-ReLU
    /// value was positive, so we can use the post-ReLU value directly.
    ///
    /// Grid: `ceil(n / 256)`, Block: 256.
    pub fn relu_mask(
        &self,
        stream: &Arc<CudaStream>,
        dx: u64,
        activation: u64,
        n: usize,
    ) -> Result<(), MLError> {
        let n_i32 = n as i32;
        let blocks = ((n + 255) / 256) as u32;

        unsafe {
            stream
                .launch_builder(&self.relu_mask_kernel)
                .arg(&dx)
                .arg(&activation)
                .arg(&n_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("relu_mask_kernel: {e}")))?;
        }

        Ok(())
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Full backward pass
    // ══════════════════════════════════════════════════════════════════════════

    /// Full backward pass through the branching dueling network via cuBLAS SGEMM.
    ///
    /// Requires pre-computed `dL/d_logits` for each head from the C51 loss kernel.
    /// Until the standalone C51 loss kernel exists (Phase 2 Task 2), this function
    /// is gated behind a TODO comment in `GpuDqnTrainer::launch_cublas_backward`.
    ///
    /// ## Backward order (reverse of forward)
    ///
    /// For each branch d in {0, 1, 2}:
    ///   1. Branch output layer:  `d_adv_logits_d` → `h_bd`
    ///   2. Branch FC layer:      `h_bd` gradient  → `h_s2` (accumulated)
    ///
    /// Value head:
    ///   3. Value output layer:  `d_value_logits` → `h_v`
    ///   4. Value FC layer:      `h_v` gradient   → `h_s2` (accumulated)
    ///
    /// Shared trunk (using accumulated `d_h_s2`):
    ///   5. Shared layer 2: `d_h_s2` → `h_s1`
    ///   6. Shared layer 1: `d_h_s1` → (states — skip, not trainable)
    ///
    /// ## Arguments
    ///
    /// - `d_value_logits`   : `[B, NA]` gradient from C51 loss for the value head
    /// - `d_adv_logits`     : three `[B, branch_k * NA]` gradients for each branch
    /// - `states`           : `[B, SD]` raw input states (for last shared layer dW)
    /// - `save_h_s1`        : `[B, SH1]` saved shared layer 1 ReLU output
    /// - `save_h_s2`        : `[B, SH2]` saved shared layer 2 ReLU output
    /// - `save_h_v`         : `[B, VH]` saved value FC ReLU output
    /// - `save_h_b`         : three `[B, AH]` saved branch FC ReLU outputs
    /// - `w_ptrs`           : 20-element array of raw F32 device pointers into
    ///                        `params_buf` (online weights, read-only)
    /// - `grad_buf`         : flat gradient accumulator `[TOTAL_PARAMS]` — must be
    ///                        zeroed before this call
    /// - `scratch_dh`       : scratch buffers for layer-to-layer gradient passing
    #[allow(clippy::too_many_arguments)]
    pub fn backward_full(
        &self,
        stream: &Arc<CudaStream>,
        // dL/d_logits from C51 loss kernel
        d_value_logits: u64,          // [B, NA]
        d_adv_logits: &[u64; 3],      // [B, branch_k * NA] each
        // Saved forward activations
        states: u64,                  // [B, SD]
        save_h_s1: u64,               // [B, SH1]
        save_h_s2: u64,               // [B, SH2]
        save_h_v: u64,                // [B, VH]
        save_h_b: &[u64; 3],          // [B, AH] each
        // Online weight pointers (read-only for W^T computation)
        w_ptrs: &[u64; 20],
        // Flat gradient accumulator (must be zeroed by caller)
        grad_buf_base: u64,
        // Scratch buffers for inter-layer gradients
        scratch_d_h_s2: u64,          // [B, SH2] — accumulated from value + all 3 branches
        scratch_d_h_s1: u64,          // [B, SH1]
        scratch_d_h_v: u64,           // [B, VH]
        scratch_d_h_b: &[u64; 3],     // [B, AH] each
    ) -> Result<(), MLError> {
        let b = self.batch_size;
        let na = self.num_atoms;

        // ── Precomputed grad_buf offsets (match GOFF_* in dqn_training_kernel.cu) ──
        //
        // Layout: [w_s1, b_s1, w_s2, b_s2, w_v1, b_v1, w_v2, b_v2,
        //          w_b0fc, b_b0fc, w_b0out, b_b0out,
        //          w_b1fc, b_b1fc, w_b1out, b_b1out,
        //          w_b2fc, b_b2fc, w_b2out, b_b2out]
        //
        // The w_ptrs array (online weight pointers from params_buf) uses the
        // same GOFF_* layout, so goff_ptr(i) == w_ptrs[i].  We store grad
        // offsets as byte distances from grad_buf_base.  Since param_sizes
        // determines element counts, grad byte offsets are:
        //   goff_bytes[i] = sum_{j<i} param_sizes[j] * sizeof(f32)
        //
        // Rather than recompute these here, we leverage the fact that
        // params_buf and grad_buf have the same GOFF_* layout.  The grad
        // pointer for tensor i is: grad_buf_base + f32_byte_offset(i).
        //
        // We compute these inline from the config dimensions.
        let f32 = std::mem::size_of::<f32>() as u64;
        let sd = self.state_dim as u64;
        let sh1 = self.shared_h1 as u64;
        let sh2 = self.shared_h2 as u64;
        let vh = self.value_h as u64;
        let ah = self.adv_h as u64;
        let na64 = na as u64;
        let b0 = self.branch_0_size as u64;
        let b1 = self.branch_1_size as u64;
        let b2 = self.branch_2_size as u64;

        // GOFF byte offsets (must match compute_param_sizes order)
        let goff_w_s1:    u64 = 0;
        let goff_b_s1:    u64 = goff_w_s1    + sh1 * sd         * f32;
        let goff_w_s2:    u64 = goff_b_s1    + sh1              * f32;
        let goff_b_s2:    u64 = goff_w_s2    + sh2 * sh1        * f32;
        let goff_w_v1:    u64 = goff_b_s2    + sh2              * f32;
        let goff_b_v1:    u64 = goff_w_v1    + vh  * sh2        * f32;
        let goff_w_v2:    u64 = goff_b_v1    + vh               * f32;
        let goff_b_v2:    u64 = goff_w_v2    + na64 * vh        * f32;
        let goff_w_b0fc:  u64 = goff_b_v2    + na64             * f32;
        let goff_b_b0fc:  u64 = goff_w_b0fc  + ah  * sh2        * f32;
        let goff_w_b0out: u64 = goff_b_b0fc  + ah               * f32;
        let goff_b_b0out: u64 = goff_w_b0out + b0  * na64 * ah  * f32;
        let goff_w_b1fc:  u64 = goff_b_b0out + b0  * na64       * f32;
        let goff_b_b1fc:  u64 = goff_w_b1fc  + ah  * sh2        * f32;
        let goff_w_b1out: u64 = goff_b_b1fc  + ah               * f32;
        let goff_b_b1out: u64 = goff_w_b1out + b1  * na64 * ah  * f32;
        let goff_w_b2fc:  u64 = goff_b_b1out + b1  * na64       * f32;
        let goff_b_b2fc:  u64 = goff_w_b2fc  + ah  * sh2        * f32;
        let goff_w_b2out: u64 = goff_b_b2fc  + ah               * f32;
        let goff_b_b2out: u64 = goff_w_b2out + b2  * na64 * ah  * f32;

        let goff_w_bfc  = [goff_w_b0fc,  goff_w_b1fc,  goff_w_b2fc];
        let goff_b_bfc  = [goff_b_b0fc,  goff_b_b1fc,  goff_b_b2fc];
        let goff_w_bout = [goff_w_b0out, goff_w_b1out, goff_w_b2out];
        let goff_b_bout = [goff_b_b0out, goff_b_b1out, goff_b_b2out];

        let branch_n = [
            self.branch_0_size * na,
            self.branch_1_size * na,
            self.branch_2_size * na,
        ];

        // w_ptrs indices:
        //   [0]=w_s1, [1]=b_s1, [2]=w_s2, [3]=b_s2,
        //   [4]=w_v1, [5]=b_v1, [6]=w_v2, [7]=b_v2,
        //   [8]=w_b0fc, [9]=b_b0fc, [10]=w_b0out, [11]=b_b0out,
        //   [12]=w_b1fc, [13]=b_b1fc, [14]=w_b1out, [15]=b_b1out,
        //   [16]=w_b2fc, [17]=b_b2fc, [18]=w_b2out, [19]=b_b2out
        let w_bfc_idx  = [8_usize, 12, 16];
        let w_bout_idx = [10_usize, 14, 18];

        // ══════════════════════════════════════════════════════════════════
        // BRANCHES (all 3, in parallel with shared h_s2 accumulation)
        // ══════════════════════════════════════════════════════════════════

        for d in 0..3 {
            let n_d = branch_n[d];
            let w_out = w_ptrs[w_bout_idx[d]]; // W_bdk_out[n_d, AH]
            let w_fc  = w_ptrs[w_bfc_idx[d]];  // W_bdk_fc[AH, SH2]

            // ── Branch output layer (no ReLU — logit layer) ──────────
            // dY = d_adv_logits[d] [B, n_d]
            // X  = save_h_b[d]     [B, AH]
            // W  = W_bdk_out       [n_d, AH]
            // dW += dY^T @ X,  db += sum(dY),  dX_b = dY @ W^T
            self.backward_fc_layer(
                stream,
                d_adv_logits[d],           // dY [B, n_d]
                save_h_b[d],               // X  [B, AH]
                w_out,                     // W  [n_d, AH]
                grad_buf_base + goff_w_bout[d], // dW
                grad_buf_base + goff_b_bout[d], // db
                scratch_d_h_b[d],          // dX [B, AH]
                n_d,                       // out_dim
                self.adv_h,                // in_dim
                b,
            )?;
            // No ReLU mask needed — branch output is a logit layer.

            // ── Branch FC layer (ReLU activation) ────────────────────
            // dY = scratch_d_h_b[d] [B, AH]
            // X  = save_h_s2        [B, SH2]
            // W  = W_bdk_fc         [AH, SH2]
            // dW += dY^T @ X,  db += sum(dY),  dX_s2 written to scratch_d_h_s2
            //
            // NOTE: We accumulate contributions from all 3 branches into
            // scratch_d_h_s2.  To do this safely, we use beta=1.0 in the
            // backward_fc_layer call.  However, backward_fc_layer always uses
            // beta=0.0 for dX.  We therefore use a temporary per-branch dX
            // and manually add it into scratch_d_h_s2.
            //
            // For simplicity, we write into scratch_d_h_b[d] (reuse — output
            // layer has already consumed it) and then accumulate.
            //
            // Actually, scratch_d_h_b[d] is already holding the branch FC
            // upstream gradient dX that we need to produce.  We write it
            // there, apply ReLU, then accumulate into scratch_d_h_s2.
            //
            // So we need a separate branch_dh_s2 scratch.  For now, use
            // scratch_d_h_b[d] as the per-branch dh_s2 temporary: after the
            // branch out backward we no longer need d_h_b[d] for anything,
            // so we can repurpose it.  But its size is [B, AH], not [B, SH2].
            // We cannot reuse it directly.
            //
            // Solution: run the branch FC backward writing dX into
            // scratch_d_h_s2 with beta=0 first, then for branches 1 and 2
            // we need to add rather than overwrite.  We handle this by
            // setting dx=0 for the fc layer and doing the upstream gradient
            // via a separate GEMM call with beta=1.0 for branches d>0.

            // Apply ReLU mask to branch FC upstream: d_h_b[d] *= (save_h_b[d] > 0)
            self.relu_mask(
                stream,
                scratch_d_h_b[d],          // dx to gate
                save_h_b[d],               // saved post-ReLU activation
                b * self.adv_h,
            )?;

            // dW for branch FC: dW[AH, SH2] += d_h_b[d]^T @ h_s2
            // (dX is computed separately below to allow accumulation)
            self.launch_dw_only(
                stream,
                scratch_d_h_b[d],          // dY [B, AH]
                save_h_s2,                 // X  [B, SH2]
                grad_buf_base + goff_w_bfc[d], // dW [AH, SH2]
                grad_buf_base + goff_b_bfc[d], // db [AH]
                self.adv_h,                // out_dim
                self.shared_h2,            // in_dim
                b,
            )?;

            // dX for branch FC: d_h_s2 += d_h_b[d] @ W_bdk_fc
            // Use beta = if d==0 { 0.0 } else { 1.0 } to accumulate branches.
            let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 };
            self.launch_dx_only(
                stream,
                scratch_d_h_b[d],          // dY [B, AH]
                w_fc,                      // W  [AH, SH2]
                scratch_d_h_s2,            // dX [B, SH2]
                self.adv_h,                // out_dim
                self.shared_h2,            // in_dim
                b,
                beta_s2,
            )?;
        }

        // ══════════════════════════════════════════════════════════════════
        // VALUE HEAD
        // ══════════════════════════════════════════════════════════════════

        // ── Value output layer (no ReLU) ─────────────────────────────
        // dY = d_value_logits [B, NA],  X = save_h_v [B, VH],  W = W_v2 [NA, VH]
        self.backward_fc_layer(
            stream,
            d_value_logits,
            save_h_v,
            w_ptrs[6],                 // W_v2 [NA, VH]
            grad_buf_base + goff_w_v2,
            grad_buf_base + goff_b_v2,
            scratch_d_h_v,             // dX [B, VH]
            na,
            self.value_h,
            b,
        )?;

        // ── Value FC layer (ReLU) ─────────────────────────────────────
        self.relu_mask(stream, scratch_d_h_v, save_h_v, b * self.value_h)?;

        // dW for value FC: dW[VH, SH2] += d_h_v^T @ h_s2
        self.launch_dw_only(
            stream,
            scratch_d_h_v,
            save_h_s2,
            grad_buf_base + goff_w_v1,
            grad_buf_base + goff_b_v1,
            self.value_h,
            self.shared_h2,
            b,
        )?;

        // dX for value FC: accumulated into scratch_d_h_s2 (beta=1.0 — branches already wrote)
        self.launch_dx_only(
            stream,
            scratch_d_h_v,
            w_ptrs[4],                 // W_v1 [VH, SH2]
            scratch_d_h_s2,
            self.value_h,
            self.shared_h2,
            b,
            1.0_f32,                   // accumulate on top of branch contributions
        )?;

        // ══════════════════════════════════════════════════════════════════
        // SHARED TRUNK (backward through accumulated d_h_s2)
        // ══════════════════════════════════════════════════════════════════

        // ── Shared layer 2 (ReLU) ─────────────────────────────────────
        self.relu_mask(stream, scratch_d_h_s2, save_h_s2, b * self.shared_h2)?;

        self.backward_fc_layer(
            stream,
            scratch_d_h_s2,
            save_h_s1,
            w_ptrs[2],                 // W_s2 [SH2, SH1]
            grad_buf_base + goff_w_s2,
            grad_buf_base + goff_b_s2,
            scratch_d_h_s1,            // dX [B, SH1]
            self.shared_h2,
            self.shared_h1,
            b,
        )?;

        // ── Shared layer 1 (ReLU) — input layer, no dX needed ────────
        self.relu_mask(stream, scratch_d_h_s1, save_h_s1, b * self.shared_h1)?;

        // dW for shared layer 1: states is the input
        self.backward_fc_layer(
            stream,
            scratch_d_h_s1,
            states,
            w_ptrs[0],                 // W_s1 [SH1, SD]
            grad_buf_base + goff_w_s1,
            grad_buf_base + goff_b_s1,
            0,                         // dx=0: states are not trainable
            self.shared_h1,
            self.state_dim,
            b,
        )?;

        let _ = goff_b_b2out; // suppress dead_code for last computed offset

        Ok(())
    }

    // ══════════════════════════════════════════════════════════════════════════
    // Private SGEMM helpers
    // ══════════════════════════════════════════════════════════════════════════

    /// Compute only weight gradient + bias gradient (no upstream dX).
    ///
    /// Used for branch FC layers where the upstream dX is computed separately
    /// via `launch_dx_only` so multiple branches can accumulate into the
    /// shared `d_h_s2` scratch buffer with controllable beta.
    #[allow(clippy::too_many_arguments)]
    fn launch_dw_only(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,
        x: u64,
        dw: u64,
        db: u64,
        out_dim: usize,
        in_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        // dW[out, in] += dY^T @ X  (same mapping as in backward_fc_layer)
        unsafe {
            let alpha = 1.0_f32;
            let beta_acc = 1.0_f32;
            cublas_result::sgemm(
                self.handle.0,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                cublas_sys::cublasOperation_t::CUBLAS_OP_T,
                in_dim as i32,
                out_dim as i32,
                batch as i32,
                &alpha,
                x as *const f32,
                in_dim as i32,
                dy as *const f32,
                out_dim as i32,
                &beta_acc,
                dw as *mut f32,
                in_dim as i32,
            )
            .map_err(|e| MLError::ModelError(format!("backward sgemm dW_only: {e:?}")))?;
        }

        self.launch_bias_grad(stream, dy, db, out_dim, batch)
    }

    /// Compute only upstream gradient dX = dY @ W^T, with configurable beta.
    ///
    /// `beta=0.0` overwrites dX; `beta=1.0` accumulates into dX.  Used to merge
    /// contributions from multiple branches into the shared `d_h_s2` buffer.
    ///
    /// Also used by ensemble diversity backward to skip dW/db for value head layers
    /// (only the upstream gradient d_h_s2 is needed, not the value head weight grads).
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn launch_dx_only(
        &self,
        _stream: &Arc<CudaStream>,
        dy: u64,
        w: u64,
        dx: u64,
        out_dim: usize,
        in_dim: usize,
        batch: usize,
        beta: f32,
    ) -> Result<(), MLError> {
        // dX[B, in] = dY[B, out] @ W[out, in]
        // col-major: dX'[in,B] = W'[in,out] @ dY'[out,B]
        // sgemm(N, N, m=in_dim, n=B, k=out_dim, W, in_dim, dY, out_dim, beta, dX, in_dim)
        unsafe {
            let alpha = 1.0_f32;
            cublas_result::sgemm(
                self.handle.0,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                cublas_sys::cublasOperation_t::CUBLAS_OP_N,
                in_dim as i32,
                batch as i32,
                out_dim as i32,
                &alpha,
                w as *const f32,
                in_dim as i32,
                dy as *const f32,
                out_dim as i32,
                &beta,
                dx as *mut f32,
                in_dim as i32,
            )
            .map_err(|e| MLError::ModelError(format!("backward sgemm dX_only: {e:?}")))?;
        }

        Ok(())
    }

    /// Launch `bias_grad_reduce_kernel`: db[j] += sum_b(dY[b * out_dim + j]).
    fn launch_bias_grad(
        &self,
        stream: &Arc<CudaStream>,
        dy: u64,
        db: u64,
        out_dim: usize,
        batch: usize,
    ) -> Result<(), MLError> {
        let out_dim_i32 = out_dim as i32;
        let batch_i32 = batch as i32;
        let blocks = ((out_dim + 255) / 256) as u32;

        unsafe {
            stream
                .launch_builder(&self.bias_grad_kernel)
                .arg(&dy)
                .arg(&db)
                .arg(&out_dim_i32)
                .arg(&batch_i32)
                .launch(LaunchConfig {
                    grid_dim: (blocks, 1, 1),
                    block_dim: (256, 1, 1),
                    shared_mem_bytes: 0,
                })
                .map_err(|e| MLError::ModelError(format!("bias_grad_reduce_kernel: {e}")))?;
        }

        Ok(())
    }
}

// ── Kernel compilation ────────────────────────────────────────────────────────

/// Compile `relu_mask_kernel` and `bias_grad_reduce_kernel` from inline NVRTC.
///
/// Both kernels are simple element-wise or reduction operations:
/// - `relu_mask_kernel`: 1 thread per element, gating dx by saved activation.
/// - `bias_grad_reduce_kernel`: 1 thread per output neuron, summing over batch.
///
/// These are NOT captured in the CUDA Graph because their launch overhead
/// (< 1 µs per kernel) is negligible compared to the GEMM time.  For CUDA Graph
/// compatibility, they would need to be added to the captured region, but this
/// optimisation is deferred until Phase 2 Task 3.
fn compile_backward_kernels(
    stream: &Arc<CudaStream>,
) -> Result<(CudaFunction, CudaFunction), MLError> {
    let src = r#"
// ReLU derivative mask: dx[i] *= (activation[i] > 0.0f)
//
// `activation` is the SAVED POST-ReLU value from the forward pass.
// For standard ReLU, post-ReLU > 0 iff pre-ReLU > 0, so the saved
// post-ReLU value is a valid proxy for the derivative gate.
//
// Grid: ceil(n / 256), Block: 256
extern "C" __global__ void relu_mask_kernel(
    float* __restrict__ dx,
    const float* __restrict__ activation,
    int n)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    if (activation[i] <= 0.0f) dx[i] = 0.0f;
}

// Bias gradient: db[j] += sum_b dy[b * out_dim + j]
//
// Each thread owns one output neuron j and iterates over the batch.
// Uses atomicAdd so that multiple kernel calls (from different layers
// sharing the same grad_buf) accumulate correctly.
//
// Grid: ceil(out_dim / 256), Block: 256
extern "C" __global__ void bias_grad_reduce_kernel(
    const float* __restrict__ dy,
    float* __restrict__ db,
    int out_dim,
    int batch_size)
{
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (j >= out_dim) return;
    float sum = 0.0f;
    for (int b = 0; b < batch_size; b++) {
        sum += dy[b * out_dim + j];
    }
    atomicAdd(&db[j], sum);
}
"#;

    let context = stream.context();
    let ptx = crate::cuda_pipeline::compile_ptx_for_device(src, &context)
        .map_err(|e| MLError::ModelError(format!("backward kernel compilation: {e}")))?;
    let module = context
        .load_module(ptx)
        .map_err(|e| MLError::ModelError(format!("backward kernel module load: {e}")))?;

    let relu_mask = module
        .load_function("relu_mask_kernel")
        .map_err(|e| MLError::ModelError(format!("relu_mask_kernel load: {e}")))?;
    let bias_grad = module
        .load_function("bias_grad_reduce_kernel")
        .map_err(|e| MLError::ModelError(format!("bias_grad_reduce_kernel load: {e}")))?;

    Ok((relu_mask, bias_grad))
}

// ── Raw device pointer helpers ────────────────────────────────────────────────
//
// Mirrors the helpers in batched_forward.rs.  Declared here so this module
// can be compiled independently without `pub use`-ing the forward module's
// private functions.

/// Extract raw F32 device pointer from a CudaSlice (read-only).
pub(crate) fn raw_f32_ptr(slice: &CudaSlice<f32>, stream: &Arc<CudaStream>) -> u64 {
    let (ptr, guard) = slice.device_ptr(stream);
    let _no_drop = ManuallyDrop::new(guard);
    ptr
}

// ── Scratch buffer allocation helper ─────────────────────────────────────────

/// Allocate zero-initialised scratch buffers for the cuBLAS backward pass.
///
/// Called from `GpuDqnTrainer::new()` alongside the other buffer allocations.
/// Returns `(d_h_s2, d_h_s1, d_h_v, d_h_b0, d_h_b1, d_h_b2)`.
pub fn alloc_backward_scratch(
    stream: &Arc<CudaStream>,
    config: &GpuDqnTrainConfig,
) -> Result<(
    CudaSlice<f32>,  // d_h_s2 [B, SH2]
    CudaSlice<f32>,  // d_h_s1 [B, SH1]
    CudaSlice<f32>,  // d_h_v  [B, VH]
    CudaSlice<f32>,  // d_h_b0 [B, AH]
    CudaSlice<f32>,  // d_h_b1 [B, AH]
    CudaSlice<f32>,  // d_h_b2 [B, AH]
), MLError> {
    let b = config.batch_size;
    let alloc = |n: usize| -> Result<CudaSlice<f32>, MLError> {
        stream.alloc_zeros::<f32>(n)
            .map_err(|e| MLError::ModelError(format!("backward scratch alloc [{n}]: {e}")))
    };

    Ok((
        alloc(b * config.shared_h2)?,
        alloc(b * config.shared_h1)?,
        alloc(b * config.value_h)?,
        alloc(b * config.adv_h)?,
        alloc(b * config.adv_h)?,
        alloc(b * config.adv_h)?,
    ))
}
