/**
 * Ensemble Q-Network kernels -- aggregate + KL gradient.
 *
 * Two live kernels for ensemble training with K independent value/advantage
 * head weight sets sharing a common DQN trunk:
 *
 * 1. ensemble_aggregate_kernel -- compute mean and variance of Q-values across
 *    K heads (uncertainty estimation, used by exploration bonus).
 *
 * 2. ensemble_kl_gradient_kernel -- compute dL_diversity/d_logits for head 0,
 *    fed into grad_buf via SAXPY in apply_ensemble_diversity_backward.
 *    This is the actual gradient that pushes heads apart during training.
 *
 * REMOVED 2026-04-21: ensemble_diversity_kernel (computed pairwise KL as a
 * scalar diagnostic). The result was written to ensemble_diversity_loss_buf
 * and pending_diversity_loss_ptr, but readback_diversity_loss() was never
 * called from anywhere — pure dead code that ran every training step. Its
 * removal eliminates the only remaining training-path atomicAdd outside
 * monitoring kernels.
 *
 * No #include of common_device_functions.cuh is needed here -- the build
 * function in gpu_dqn_trainer.rs prepends it via string concatenation.
 */

/* ======================================================================
 * KERNEL 1: ENSEMBLE Q-VALUE AGGREGATION
 *
 * Computes mean and variance of Q-values from K independent heads.
 *
 * Input layout: head_q_values[K * B * num_actions] -- head k starts at
 *   offset k * B * num_actions (column-major head ordering).
 *
 * Launch config: grid=(ceil(B*num_actions/256), 1, 1), block=(256, 1, 1).
 * ====================================================================== */
extern "C" __global__ void ensemble_aggregate_kernel(
    const float* __restrict__ head_q_values, /* [K * B * num_actions] flat */
    float* __restrict__ mean_q,              /* [B * num_actions] output: mean */
    float* __restrict__ var_q,               /* [B * num_actions] output: variance */
    int K,
    int B,
    int num_actions
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int total = B * num_actions;
    if (idx >= total) return;

    float sum = 0.0f;
    float sum_sq = 0.0f;
    for (int k = 0; k < K; k++) {
        float q = head_q_values[(long long)k * total + idx];
        sum = sum + q;
        sum_sq = sum_sq + q * q;
    }
    float k_bf = (float)K;
    float mean = sum / k_bf;
    mean_q[idx] = mean;
    /* Var = E[X^2] - E[X]^2, clamped to 0 for numerical stability */
    float v = sum_sq / k_bf - mean * mean;
    var_q[idx] = fmaxf(v, 0.0f);
}

/* ======================================================================
 * ENSEMBLE KL GRADIENT KERNEL
 *
 * Computes dL_diversity/d_logits for head 0 by averaging the KL gradient
 * from all other heads:
 *
 *   d_logits_0[b][a] = -diversity_weight * mean_k( softmax(logits_0) - softmax(logits_k) )
 *
 * The NEGATIVE sign encourages head 0 to MAXIMIZE divergence from other heads
 * (diversity regularization). This gradient is added to grad_buf's value logits
 * section via SAXPY, so the graph_adam's single Adam sees it.
 *
 * Grid: (B, 1, 1), Block: (min(NA, 256), 1, 1).
 * ====================================================================== */

extern "C" __global__ void ensemble_kl_gradient_kernel(
    const float* __restrict__ head_logits,  /* [K * B * NA] f32 */
    float* __restrict__ d_logits_0,          /* [B * NA] output: diversity gradient for head 0 */
    float diversity_weight,
    int K, int B, int NA
) {
    int b = blockIdx.x;
    int a = threadIdx.x;
    if (b >= B || a >= NA) return;

    int stride = B * NA;

    /* Softmax of head 0 at (b, a) -- f32 */
    float max0 = -1e30f;
    for (int j = 0; j < NA; j++) {
        float v = head_logits[0 * stride + b * NA + j];
        max0 = fmaxf(max0, v);
    }
    float sum0 = 0.0f;
    for (int j = 0; j < NA; j++)
        sum0 += expf(head_logits[0 * stride + b * NA + j] - max0);
    float p0 = expf(head_logits[0 * stride + b * NA + a] - max0) / (sum0 + 1e-8f);

    /* Average gradient from all other heads: mean_k(p0 - pk) */
    float grad = 0.0f;
    for (int k = 1; k < K; k++) {
        float maxk = -1e30f;
        for (int j = 0; j < NA; j++) {
            float v = head_logits[k * stride + b * NA + j];
            maxk = fmaxf(maxk, v);
        }
        float sumk = 0.0f;
        for (int j = 0; j < NA; j++)
            sumk += expf(head_logits[k * stride + b * NA + j] - maxk);
        float pk = expf(head_logits[k * stride + b * NA + a] - maxk) / (sumk + 1e-8f);

        grad += (p0 - pk);
    }
    grad /= (float)(K - 1);

    /* NEGATIVE: maximize divergence (subtract from loss) */
    d_logits_0[b * NA + a] = -diversity_weight * grad;
}
