/**
 * Ensemble Q-Network kernels — aggregate and diversity loss.
 *
 * Two 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).
 *
 * 2. ensemble_diversity_kernel — pairwise KL divergence between softmax
 *    distributions of C51 head logits (diversity regularization).
 *    Uses warp-level → block-level → atomicAdd-per-block hierarchical reduction
 *    matching the pattern in dqn_grad_norm_kernel.
 *
 * 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, sum_sq = 0.0f;
    for (int k = 0; k < K; k++) {
        float q = head_q_values[(long long)k * total + idx];
        sum += q;
        sum_sq += q * q;
    }
    float mean = sum / (float)K;
    mean_q[idx] = mean;
    /* Var = E[X^2] - E[X]^2, clamped to 0 for numerical stability */
    float v = sum_sq / (float)K - mean * mean;
    var_q[idx] = (v > 0.0f) ? v : 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 2: ENSEMBLE DIVERSITY LOSS (KL DIVERGENCE)
 *
 * Computes pairwise KL divergence between softmax distributions of K heads'
 * C51 value logits. KL is averaged over all K*(K-1)/2 head pairs and
 * all B samples. Written as a scalar into diversity_loss[0].
 *
 * Uses hierarchical reduction: warp reduce → block shared memory → per-block
 * atomicAdd. This matches dqn_grad_norm_kernel and avoids serialization at
 * a single atomic address.
 *
 * head_logits layout: [K * B * num_atoms] — head k starts at k*B*num_atoms.
 *
 * Each thread processes one (sample, head_i, head_j) pair's KL contribution.
 * Grid sizing: grid_x = ceil(B * num_pairs / 256), where num_pairs = K*(K-1)/2.
 *
 * Launch config: grid=(ceil(B*num_pairs/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void ensemble_diversity_kernel(
    const float* __restrict__ head_logits, /* [K * B * num_atoms] */
    float* __restrict__ diversity_loss,    /* [1] output: accumulated KL */
    int K,
    int B,
    int num_atoms
) {
    /* Number of ordered pairs (i < j): K*(K-1)/2 */
    int num_pairs = K * (K - 1) / 2;
    int total_work = B * num_pairs;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    float kl_sum = 0.0f;

    if (idx < total_work) {
        /* Decode sample index and pair index */
        int sample = idx / num_pairs;
        int pair   = idx % num_pairs;

        /* Decode ordered pair (i, j) with i < j from linear pair index.
         * Triangular number inversion: pair maps to (i, j) with j > i.
         *   pair = i*(K-1) - i*(i-1)/2 + (j - i - 1)   [0-indexed] */
        int hi = 0, hj = 1;
        int p = 0;
        for (int i = 0; i < K - 1; i++) {
            for (int j = i + 1; j < K; j++) {
                if (p == pair) { hi = i; hj = j; }
                p++;
            }
        }

        /* Pointers to head i and head j logits for this sample */
        const float* logits_i = head_logits + (long long)hi * B * num_atoms + sample * num_atoms;
        const float* logits_j = head_logits + (long long)hj * B * num_atoms + sample * num_atoms;

        /* Softmax of logits_i */
        float max_i = logits_i[0];
        for (int a = 1; a < num_atoms; a++) max_i = fmaxf(max_i, logits_i[a]);
        float sum_i = 0.0f;
        for (int a = 0; a < num_atoms; a++) sum_i += expf(logits_i[a] - max_i);

        /* Softmax of logits_j */
        float max_j = logits_j[0];
        for (int a = 1; a < num_atoms; a++) max_j = fmaxf(max_j, logits_j[a]);
        float sum_j = 0.0f;
        for (int a = 0; a < num_atoms; a++) sum_j += expf(logits_j[a] - max_j);

        /* KL(p_i || p_j) = sum_a p_i[a] * log(p_i[a] / p_j[a]) */
        float kl = 0.0f;
        for (int a = 0; a < num_atoms; a++) {
            float pi = expf(logits_i[a] - max_i) / sum_i;
            float pj = expf(logits_j[a] - max_j) / sum_j;
            /* Symmetrized KL (Jensen–Shannon style): (KL(i||j) + KL(j||i)) / 2 */
            float log_ratio_ij = logf(pi / (pj + 1e-8f) + 1e-8f);
            float log_ratio_ji = logf(pj / (pi + 1e-8f) + 1e-8f);
            kl += 0.5f * (pi * log_ratio_ij + pj * log_ratio_ji);
        }
        kl_sum = kl;
    }

    /* ── Hierarchical reduction: warp → block → atomicAdd ──────────── */
    /* Warp-level reduction via shuffle (no shared memory) */
    for (int offset = 16; offset > 0; offset >>= 1)
        kl_sum += __shfl_xor_sync(0xFFFFFFFF, kl_sum, offset);

    /* Block-level cross-warp reduction via shared memory (padded, same as grad_norm) */
    __shared__ float warp_sums[16]; /* 8 warps × 2 stride (padded) */
    int warp_id   = threadIdx.x / 32;
    int warp_lane = threadIdx.x % 32;
    if (warp_lane == 0) warp_sums[warp_id * 2] = kl_sum;
    __syncthreads();

    /* First warp reduces across warps */
    if (warp_id == 0) {
        float val = (warp_lane < blockDim.x / 32) ? warp_sums[warp_lane * 2] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        if (warp_lane == 0)
            atomicAdd(diversity_loss, val);
    }
}

/* ══════════════════════════════════════════════════════════════════════
 * 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] */
    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) — numerically stable */
    float max0 = -1e30f;
    for (int j = 0; j < NA; j++) {
        float v = head_logits[0 * stride + b * NA + j];
        if (v > max0) 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];
            if (v > maxk) 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;
}
