// rl_ensemble_action_value.cu — combine C51 + IQN expected Q values.
//
// E_ensemble[b,a] = alpha * E_C51[b,a] + (1-alpha) * E_IQN[b,a]
//
// Reads alpha from ISV[RL_IQN_ENSEMBLE_ALPHA_INDEX=544].
// One thread per (batch, action). Output feeds the Q-pi agreement
// diagnostic and can inform ensemble-level action selection in future.
//
// Per `feedback_no_atomicadd.md`: no atomicAdd (not needed here).
// Per `feedback_cpu_is_read_only.md`: alpha from ISV, not host param.
// Per `feedback_isv_for_adaptive_bounds.md`: alpha is ISV-driven.

#include <stdint.h>

#define N_ACTIONS                       11
#define Q_N_ATOMS                       21
#define RL_IQN_ENSEMBLE_ALPHA_INDEX     544

extern "C" __global__ void rl_ensemble_action_value(
    const float* __restrict__ c51_logits,    // [B * N_ACTIONS * Q_N_ATOMS]
    const float* __restrict__ atom_supports, // [Q_N_ATOMS]
    const float* __restrict__ iqn_expected,  // [B * N_ACTIONS]
    const float* __restrict__ isv,
    float*       __restrict__ ensemble_q,    // [B * N_ACTIONS] OUT
    int                       b_size
) {
    const int b = blockIdx.x;
    const int a = threadIdx.x;
    if (b >= b_size || a >= N_ACTIONS) return;

    const float alpha = isv[RL_IQN_ENSEMBLE_ALPHA_INDEX];

    // C51 expected Q for this (batch, action): sum of probs * atoms
    const int q_base = (b * N_ACTIONS + a) * Q_N_ATOMS;
    float max_l = c51_logits[q_base];
    for (int z = 1; z < Q_N_ATOMS; ++z) {
        float v = c51_logits[q_base + z];
        if (v > max_l) max_l = v;
    }
    float sum_exp = 0.0f;
    float e_c51 = 0.0f;
    for (int z = 0; z < Q_N_ATOMS; ++z) {
        float p = expf(c51_logits[q_base + z] - max_l);
        sum_exp += p;
        e_c51 += p * atom_supports[z];
    }
    e_c51 /= (sum_exp + 1e-8f);

    // IQN expected Q (already reduced over tau)
    const float e_iqn = iqn_expected[b * N_ACTIONS + a];

    // Ensemble
    ensemble_q[b * N_ACTIONS + a] = alpha * e_c51 + (1.0f - alpha) * e_iqn;
}
