// rl_iqn_backward.cu — IQN backward pass: hybrid cuBLAS + custom kernel.
//
// The backward pass is split into:
//
//   1. cuBLAS SGEMM (from Rust):
//      grad_combined[B*N_TAU, HIDDEN_DIM] = grad_q[B*N_TAU, N_ACTIONS]
//                                           @ W_out^T[N_ACTIONS, HIDDEN_DIM]
//      Replaces the per-thread inner loop over N_ACTIONS.
//
//   2. rl_iqn_backward (custom kernel, this file):
//      Per-batch backward accumulation. Recomputes phi(tau) and combined
//      from (h_t, tau, W_embed, b_embed) inline (same as the original
//      monolithic kernel). Uses the cuBLAS-computed grad_combined instead
//      of recomputing it from grad_q × W_out.
//
//      Produces per-batch gradients for:
//        - grad_w_out_pb  [B, HIDDEN_DIM * N_ACTIONS]
//        - grad_b_out_pb  [B, N_ACTIONS]
//        - grad_w_embed_pb [B, EMBED_DIM * HIDDEN_DIM]
//        - grad_b_embed_pb [B, HIDDEN_DIM]
//
//   3. rl_iqn_bwd_relu_hadamard (custom kernel, this file):
//      Element-wise backward through hadamard + ReLU. Used when the
//      forward cache provides embed_pre_relu (alternative to inline
//      recompute). Kept for future use when forward caching is added.
//
// Block layout for rl_iqn_backward:
//   grid  = (B, 1, 1)
//   block = (HIDDEN_DIM, 1, 1)
//   One thread per hidden dim; each thread loops over N_TAU.
//
// Per `feedback_no_atomicadd.md`: no atomicAdd.
// Per `feedback_cpu_is_read_only.md`: all compute on GPU.

#define HIDDEN_DIM   128
#define N_ACTIONS    11
#define N_TAU_MAX    64
#define EMBED_DIM    64
#define PI_F         3.14159265f

// ─────────────────────────────────────────────────────────────────────
// rl_iqn_backward:
//   Per-batch backward with cuBLAS-precomputed grad_combined.
//
//   Thread c recomputes phi(tau)[c] and combined[c] from (tau, W_embed,
//   b_embed, h_t) for each quantile sample, then accumulates:
//     - grad_w_out[c, a] = Σ_t combined[c] * grad_q[t, a]
//     - grad_phi[c] = grad_combined[c] * h_t[c]       (from cuBLAS)
//     - grad_embed_input[c] = grad_phi[c] * relu_mask
//     - grad_w_embed[i, c] = Σ_t cos((i+1)*pi*tau[t]) * grad_embed_input[c]
//     - grad_b_embed[c] = Σ_t grad_embed_input[c]
//     - grad_b_out[a] = Σ_t grad_q[t, a]  (thread 0)
//
// Inputs:
//   h_t              [B, HIDDEN_DIM]
//   tau              [B, N_TAU]               — saved from forward
//   w_embed          [EMBED_DIM, HIDDEN_DIM]  — online weights
//   b_embed          [HIDDEN_DIM]
//   grad_combined    [B*N_TAU, HIDDEN_DIM]    — from cuBLAS (grad_q @ W_out^T)
//   grad_output      [B, N_TAU, N_ACTIONS]    — from loss
//   B, N_TAU
// Outputs:
//   grad_w_out_pb    [B, HIDDEN_DIM * N_ACTIONS]
//   grad_b_out_pb    [B, N_ACTIONS]
//   grad_w_embed_pb  [B, EMBED_DIM * HIDDEN_DIM]
//   grad_b_embed_pb  [B, HIDDEN_DIM]
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_iqn_backward(
    const float* __restrict__ h_t,              // [B, HIDDEN_DIM]
    const float* __restrict__ tau,              // [B, N_TAU]
    const float* __restrict__ w_embed,          // [EMBED_DIM, HIDDEN_DIM]
    const float* __restrict__ b_embed,          // [HIDDEN_DIM]
    const float* __restrict__ grad_combined,    // [B*N_TAU, HIDDEN_DIM] (cuBLAS)
    const float* __restrict__ grad_output,      // [B, N_TAU, N_ACTIONS]
    int          B,
    int          N_TAU,
    float*       __restrict__ grad_w_out_pb,    // [B, HIDDEN_DIM * N_ACTIONS]
    float*       __restrict__ grad_b_out_pb,    // [B, N_ACTIONS]
    float*       __restrict__ grad_w_embed_pb,  // [B, EMBED_DIM * HIDDEN_DIM]
    float*       __restrict__ grad_b_embed_pb   // [B, HIDDEN_DIM]
) {
    const int batch = blockIdx.x;
    const int c = threadIdx.x;  // hidden dim index
    if (batch >= B) return;
    if (c >= HIDDEN_DIM) return;

    const float h_c = h_t[batch * HIDDEN_DIM + c];

    // Accumulate gradients across all tau for this (batch, c).
    float acc_grad_w_out[N_ACTIONS];
    for (int a = 0; a < N_ACTIONS; ++a) acc_grad_w_out[a] = 0.0f;
    float acc_grad_w_embed[EMBED_DIM];
    for (int i = 0; i < EMBED_DIM; ++i) acc_grad_w_embed[i] = 0.0f;
    float acc_grad_b_embed = 0.0f;

    for (int t = 0; t < N_TAU; ++t) {
        const int row = batch * N_TAU + t;
        const float tau_val = tau[batch * N_TAU + t];

        // Recompute forward: phi(tau)[c]
        float embed_acc = b_embed[c];
        float cos_feats[EMBED_DIM];
        #pragma unroll
        for (int i = 0; i < EMBED_DIM; ++i) {
            cos_feats[i] = cosf((float)(i + 1) * PI_F * tau_val);
            embed_acc += cos_feats[i] * w_embed[i * HIDDEN_DIM + c];
        }
        float phi_c = fmaxf(embed_acc, 0.0f);
        float relu_mask = (embed_acc > 0.0f) ? 1.0f : 0.0f;

        // combined[c] = h_t[c] * phi(tau)[c]
        float combined_c = h_c * phi_c;

        // grad_output for this (batch, tau): [N_ACTIONS]
        const int go_base = row * N_ACTIONS;

        // ─── Grad w.r.t. w_out: dL/dW_out[c,a] += combined_c * grad_Q[a]
        for (int a = 0; a < N_ACTIONS; ++a) {
            acc_grad_w_out[a] += combined_c * grad_output[go_base + a];
        }

        // ─── Grad w.r.t. combined: READ from cuBLAS output ───────────
        float grad_combined_c = grad_combined[row * HIDDEN_DIM + c];

        // ─── Chain through element-wise product: d_combined/d_phi = h_t[c]
        float grad_phi_c = grad_combined_c * h_c;

        // ─── Chain through ReLU: d_phi/d_embed_input = relu_mask
        float grad_embed_input_c = grad_phi_c * relu_mask;

        // ─── Grad w.r.t. w_embed: dL/dW_embed[i,c] += cos_feats[i] * grad_embed_input_c
        for (int i = 0; i < EMBED_DIM; ++i) {
            acc_grad_w_embed[i] += cos_feats[i] * grad_embed_input_c;
        }

        // ─── Grad w.r.t. b_embed: dL/db_embed[c] += grad_embed_input_c
        acc_grad_b_embed += grad_embed_input_c;
    }

    // Write accumulated grad_w_out for this (batch, c).
    const int wo_base = batch * HIDDEN_DIM * N_ACTIONS + c * N_ACTIONS;
    for (int a = 0; a < N_ACTIONS; ++a) {
        grad_w_out_pb[wo_base + a] = acc_grad_w_out[a];
    }

    // Write accumulated grad_w_embed for this (batch, c).
    for (int i = 0; i < EMBED_DIM; ++i) {
        grad_w_embed_pb[batch * EMBED_DIM * HIDDEN_DIM + i * HIDDEN_DIM + c] = acc_grad_w_embed[i];
    }

    // Write grad_b_embed for this (batch, c).
    grad_b_embed_pb[batch * HIDDEN_DIM + c] = acc_grad_b_embed;

    // grad_b_out: thread 0 computes by summing over tau.
    if (c == 0) {
        for (int a = 0; a < N_ACTIONS; ++a) {
            float sum = 0.0f;
            for (int t = 0; t < N_TAU; ++t) {
                sum += grad_output[batch * N_TAU * N_ACTIONS + t * N_ACTIONS + a];
            }
            grad_b_out_pb[batch * N_ACTIONS + a] = sum;
        }
    }
}

// ─────────────────────────────────────────────────────────────────────
// rl_iqn_bwd_relu_hadamard:
//   Element-wise backward through hadamard + ReLU when embed_pre_relu
//   is available from a cached forward pass.
//
//   Grid  = (B*N_TAU, ceil(HIDDEN_DIM / 256), 1)
//   Block = (min(HIDDEN_DIM, 256), 1, 1)
//
// Inputs:
//   grad_combined     [M, HIDDEN_DIM]   — from cuBLAS (grad_q @ W_out^T)
//   h_t               [B, HIDDEN_DIM]   — encoder hidden state
//   embed_pre_relu    [M, HIDDEN_DIM]   — saved from forward (before ReLU)
//   M, B, N_TAU
// Outputs:
//   grad_embed_input  [M, HIDDEN_DIM]   — gradient for embedding input
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_iqn_bwd_relu_hadamard(
    const float* __restrict__ grad_combined,    // [M, HIDDEN_DIM]
    const float* __restrict__ h_t,              // [B, HIDDEN_DIM]
    const float* __restrict__ embed_pre_relu,   // [M, HIDDEN_DIM]
    int          M,
    int          B,
    int          N_TAU,
    float*       __restrict__ grad_embed_input  // [M, HIDDEN_DIM]
) {
    const int row = blockIdx.x;
    const int c = blockIdx.y * blockDim.x + threadIdx.x;
    if (row >= M) return;
    if (c >= HIDDEN_DIM) return;

    const int batch_idx = row / N_TAU;
    const float h_c = h_t[batch_idx * HIDDEN_DIM + c];

    float grad_phi_c = grad_combined[row * HIDDEN_DIM + c] * h_c;
    float relu_mask = (embed_pre_relu[row * HIDDEN_DIM + c] > 0.0f) ? 1.0f : 0.0f;
    grad_embed_input[row * HIDDEN_DIM + c] = grad_phi_c * relu_mask;
}
