// rl_noisy_linear_backward.cu — Factored noisy linear layer backward pass.
//
// Forward recap:
//   y[b,o] = sum_i (mu_w[o,i] + sigma_w[o,i] * eps_out[o] * eps_in[i]) * x[b,i]
//          + mu_b[o] + sigma_b[o] * eps_out[o]
//
// Backward (chain rule through noisy linear):
//
//   dL/d_mu_w[o,i]    = sum_b grad_y[b,o] * x[b,i]
//   dL/d_sigma_w[o,i] = sum_b grad_y[b,o] * x[b,i] * eps_out[o] * eps_in[i]
//   dL/d_mu_b[o]      = sum_b grad_y[b,o]
//   dL/d_sigma_b[o]   = sum_b grad_y[b,o] * eps_out[o]
//
// Per-batch scratch layout (caller runs reduce_axis0 across batches):
//   grad_mu_w_pb    [B, out_dim * in_dim]
//   grad_sigma_w_pb [B, out_dim * in_dim]
//   grad_mu_b_pb    [B, out_dim]
//   grad_sigma_b_pb [B, out_dim]
//
// Block layout:
//   grid  = (B, 1, 1)          — one block per batch element
//   block = (out_dim, 1, 1)    — one thread per output neuron
//
// Each thread o computes its contribution for batch b and writes to
// per-batch scratch at [b, o, ...]. No cross-thread reduction needed
// since each thread owns its output index.
//
// Per `feedback_no_atomicadd.md`: per-thread writes only.
// Per `feedback_cpu_is_read_only.md`: all compute on GPU.

extern "C" __global__ void rl_noisy_linear_backward(
    const float* __restrict__ x,                // [B, in_dim]
    const float* __restrict__ grad_y,           // [B, out_dim]
    const float* __restrict__ eps_in,           // [in_dim]  f(noise) pre-computed
    const float* __restrict__ eps_out,          // [out_dim] f(noise) pre-computed
    int          B,
    int          in_dim,
    int          out_dim,
    float*       __restrict__ grad_mu_w_pb,     // [B, out_dim * in_dim]
    float*       __restrict__ grad_sigma_w_pb,  // [B, out_dim * in_dim]
    float*       __restrict__ grad_mu_b_pb,     // [B, out_dim]
    float*       __restrict__ grad_sigma_b_pb   // [B, out_dim]
) {
    const int b = blockIdx.x;
    const int o = threadIdx.x;
    if (b >= B) return;
    if (o >= out_dim) return;

    const float g = grad_y[b * out_dim + o];
    const float eps_o = eps_out[o];

    // grad_mu_b[b, o] = grad_y[b, o]
    grad_mu_b_pb[b * out_dim + o] = g;

    // grad_sigma_b[b, o] = grad_y[b, o] * eps_out[o]
    grad_sigma_b_pb[b * out_dim + o] = g * eps_o;

    // Per-weight gradients for this (b, o) row.
    const int w_base = b * out_dim * in_dim + o * in_dim;
    const int x_base = b * in_dim;
    for (int i = 0; i < in_dim; ++i) {
        const float x_i = x[x_base + i];

        // grad_mu_w[b, o, i] = grad_y[b, o] * x[b, i]
        grad_mu_w_pb[w_base + i] = g * x_i;

        // grad_sigma_w[b, o, i] = grad_y[b, o] * x[b, i] * eps_out[o] * eps_in[i]
        grad_sigma_w_pb[w_base + i] = g * x_i * eps_o * eps_in[i];
    }
}
