// rl_dueling_q_loss_and_grad.cu — Phase 4 Bellman loss + decompose backward.
//
// Computes the scalar Huber loss on (target − online_composed_Q[taken])
// AND emits grad_composed[B × N_ACTIONS] in one fused kernel.
//
// Loss (per-batch):
//   δ_b         = target_value[b] − online_composed_Q[b, a_taken[b]]
//   L_b         = Huber(δ_b, κ=1.0) / B    (mean over batch)
//   loss_per_batch[b] = L_b
//
// Gradient w.r.t. online_composed_Q (mostly zero, nonzero at taken):
//   For a_taken: dL/d_composed[b, a_taken] = −(1/B) × Huber'(δ_b)
//   For a ≠ a_taken: dL/d_composed[b, a] = 0
//
// Huber'(δ) = δ                  if |δ| ≤ κ
//           = κ × sign(δ)        if |δ| > κ
//
// Block layout:
//   grid  = (B, 1, 1)
//   block = (N_ACTIONS, 1, 1)
//   Each thread writes one (b, a) cell of grad_composed; thread 0
//   computes loss + the nonzero gradient scalar.
//
// Per feedback_no_atomicadd: sole-writer per cell.

#define N_ACTIONS  11
#define HUBER_KAPPA 1.0f

extern "C" __global__ void rl_dueling_q_loss_and_grad(
    const float* __restrict__ online_composed_q,  // [B × N_ACTIONS]
    const float* __restrict__ target_value,       // [B]
    const int*   __restrict__ actions_taken,      // [B]
    int          B,
    float*       __restrict__ loss_per_batch,     // [B]
    float*       __restrict__ grad_composed       // [B × N_ACTIONS]
) {
    const int b = blockIdx.x;
    const int a = threadIdx.x;
    if (b >= B)          return;
    if (a >= N_ACTIONS)  return;

    // Defensive clamp on actions_taken (trainer should never produce
    // out-of-range, but guard against corrupt indices).
    int a_t = actions_taken[b];
    if (a_t < 0)            a_t = 0;
    if (a_t >= N_ACTIONS)   a_t = 0;

    __shared__ float s_grad_scalar;

    if (a == 0) {
        const float q_taken = online_composed_q[b * N_ACTIONS + a_t];
        const float target  = target_value[b];
        const float delta   = target - q_taken;
        const float abs_d   = fabsf(delta);
        const float inv_B   = 1.0f / (float)B;

        // Huber loss.
        float l;
        if (abs_d <= HUBER_KAPPA) {
            l = 0.5f * delta * delta;
        } else {
            l = HUBER_KAPPA * (abs_d - 0.5f * HUBER_KAPPA);
        }
        loss_per_batch[b] = l * inv_B;

        // Huber'(δ) — gradient of l w.r.t. delta.
        float huber_grad;
        if (abs_d <= HUBER_KAPPA) {
            huber_grad = delta;
        } else {
            huber_grad = HUBER_KAPPA * ((delta > 0.0f) ? 1.0f : -1.0f);
        }
        // dL/d_composed[a_taken] = (dL/dδ) × (dδ/d_composed[a_taken])
        //                        = (−1/B) × huber_grad   (δ = target − Q)
        s_grad_scalar = -inv_B * huber_grad;
    }
    __syncthreads();

    // All threads write grad_composed. Only the taken action gets a
    // nonzero value (CE on a single Q value).
    const int idx = b * N_ACTIONS + a;
    grad_composed[idx] = (a == a_t) ? s_grad_scalar : 0.0f;
}
