/**
 * MSE loss gradient kernel — FULLY DETERMINISTIC.
 *
 * Zero atomicAdd: one thread per (b, j), loops over 4 branches.
 * Same restructure as c51_grad_kernel for deterministic training.
 *
 * For each sample [b], atom [j], branch [d]:
 *   d_logit_j = td_error * is_weight * p_j * (z_j - E[Q])
 *
 * Launch config: grid=(ceil(batch_size*num_atoms/256), 1, 1), block=(256, 1, 1).
 */

extern "C" __global__ void mse_grad_kernel(
    const float* __restrict__ save_probs,       // [B, 4, NA] softmax probs
    const float* __restrict__ save_eq_td,       // [B, 4, NA] layout: [td_error, E_Q, 0, ...]
    const float* __restrict__ is_weights,                // [B] f32 (bf16 overflows to Inf)
    const int*   __restrict__ actions,          // [B] factored
    float* __restrict__ d_value_logits,         // [B, NA] f32
    float* __restrict__ d_adv_logits,           // [B, (B0+B1+B2+B3)*NA] f32
    int batch_size,
    int num_atoms,
    int b0_size, int b1_size, int b2_size, int b3_size,
    int total_branch_atoms,
    const float* __restrict__ v_range_buf,       /* [2] adaptive C51 z-support: [v_min, v_max] */
    const float* __restrict__ branch_scales)     /* [B, 4] per-sample per-branch gradient scale */
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total_elems = batch_size * num_atoms;
    if (tid >= total_elems) return;

    float v_min = v_range_buf[0];
    float v_max = v_range_buf[1];

    int j = tid % num_atoms;
    int b = tid / num_atoms;

    float isw = fminf(is_weights[b], 10.0f);
    float inv_batch = 1.0f / (float)batch_size;
    float delta_z = (num_atoms > 1) ? (v_max - v_min) / (float)(num_atoms - 1) : 0.0f;
    float z_j = v_min + (float)j * delta_z;

    /* Factored action decode (shared across branches) */
    int factored = actions[b];
    int max_action = b0_size * b1_size * b2_size * b3_size;
    if (factored < 0 || factored >= max_action) factored = 0;
    int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size };
    int branch_actions[4];
    branch_actions[0] = factored / (b1_size * b2_size * b3_size);
    branch_actions[1] = (factored / (b2_size * b3_size)) % b1_size;
    branch_actions[2] = (factored / b3_size) % b2_size;
    branch_actions[3] = factored % b3_size;

    float d_val_sum = 0.0f;
    int branch_base = 0;

    for (int d = 0; d < 4; d++) {
        int prob_idx = b * 4 * num_atoms + d * num_atoms + j;
        float p_j = (float)save_probs[prob_idx];

        int eq_base = (b * 4 + d) * num_atoms;
        float td_error = (float)save_eq_td[eq_base + 0];
        float e_q      = (float)save_eq_td[eq_base + 1];

        float d_combined = inv_batch * isw * td_error * p_j * (z_j - e_q);

        /* Per-branch entropy boost — inv_batch for batch-independent scaling */
        if (d == 1 || d == 2) {
            float ent_weight = (d == 1) ? 0.005f : 0.003f;
            float lp_approx = fmaxf(logf(fmaxf(p_j, 1e-8f)), -10.0f);
            d_combined += inv_batch * ent_weight * (1.0f + lp_approx);
        }

        d_val_sum += d_combined;

        float branch_scale = branch_scales[b * 4 + d];

        int A_d = branch_sizes[d];
        float inv_A = 1.0f / (float)A_d;
        int a_d = branch_actions[d];

        /* d_adv: plain write — unique slot per thread */
        for (int a = 0; a < A_d; a++) {
            float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
            float grad_val = branch_scale * d_combined * dueling_grad;
            int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j;
            d_adv_logits[adv_idx] = grad_val;
        }
        branch_base += batch_size * A_d * num_atoms;
    }

    /* Single deterministic write */
    d_value_logits[b * num_atoms + j] = d_val_sum;
}
