// rl_q_pi_agree_b.cu — Q–π distributional agreement diagnostic.
//
// Measures cosine similarity between Q's expected-value vector
// E_Q[a] = Σ_z softmax(q_logits[a,z]) × atom_supports[z] and π's
// softmax(pi_logits[a]) across all actions. Folded into a streaming
// EMA at ISV[407].
//
// Cosine similarity ranges [-1, +1]: +1 = Q and π rank actions
// identically, 0 = uncorrelated, -1 = perfectly opposed. This is
// robust to near-uniform distributions where argmax agreement is
// pure noise — cosine measures the full ranking direction, not just
// the noisy peak.
//
// Per `feedback_no_atomicadd`: single-block tree reduce.
// Per `feedback_cpu_is_read_only`: all state in ISV.

#define RL_Q_ARG_VS_PI_AGREE_INDEX 407
#define Q_N_ATOMS                  21
#define N_ACTIONS                   11
#define AGREE_EMA_ALPHA            0.05f

extern "C" __global__ void rl_q_pi_agree_b(
    const float* __restrict__ q_logits,       // [b_size × N_ACTIONS × Q_N_ATOMS]
    const float* __restrict__ pi_logits,      // [b_size × N_ACTIONS]
    const float* __restrict__ atom_supports,  // [Q_N_ATOMS]
    float*       __restrict__ isv,
    int                       b_size
) {
    extern __shared__ float s_cosine[];
    const int tid = threadIdx.x;

    float cosine = 0.0f;
    if (tid < b_size) {
        // Compute E_Q[a] for each action via softmax over atoms.
        float eq[N_ACTIONS];
        for (int a = 0; a < N_ACTIONS; ++a) {
            const int base = (tid * N_ACTIONS + a) * Q_N_ATOMS;
            float max_l = q_logits[base];
            for (int z = 1; z < Q_N_ATOMS; ++z) {
                float v = q_logits[base + z];
                if (v > max_l) max_l = v;
            }
            float sum_exp = 0.0f;
            float weighted = 0.0f;
            for (int z = 0; z < Q_N_ATOMS; ++z) {
                float p = expf(q_logits[base + z] - max_l);
                sum_exp += p;
                weighted += p * atom_supports[z];
            }
            eq[a] = weighted / (sum_exp + 1e-8f);
        }

        // Compute softmax(pi_logits) for this batch.
        const int pi_base = tid * N_ACTIONS;
        float pi_max = pi_logits[pi_base];
        for (int a = 1; a < N_ACTIONS; ++a) {
            float v = pi_logits[pi_base + a];
            if (v > pi_max) pi_max = v;
        }
        float pi_prob[N_ACTIONS];
        float pi_sum = 0.0f;
        for (int a = 0; a < N_ACTIONS; ++a) {
            pi_prob[a] = expf(pi_logits[pi_base + a] - pi_max);
            pi_sum += pi_prob[a];
        }
        for (int a = 0; a < N_ACTIONS; ++a) pi_prob[a] /= (pi_sum + 1e-8f);

        // Cosine similarity between E_Q and pi_prob vectors.
        float dot = 0.0f, norm_q = 0.0f, norm_pi = 0.0f;
        for (int a = 0; a < N_ACTIONS; ++a) {
            dot    += eq[a] * pi_prob[a];
            norm_q += eq[a] * eq[a];
            norm_pi += pi_prob[a] * pi_prob[a];
        }
        cosine = dot / (sqrtf(norm_q * norm_pi) + 1e-8f);
    }
    s_cosine[tid] = cosine;
    __syncthreads();

    // Tree reduce sum.
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) s_cosine[tid] += s_cosine[tid + stride];
        __syncthreads();
    }

    if (tid == 0) {
        const float agree = s_cosine[0] / (float)b_size;
        const float prev  = isv[RL_Q_ARG_VS_PI_AGREE_INDEX];
        float out;
        if (prev == 0.0f) {
            out = agree;
        } else {
            out = (1.0f - AGREE_EMA_ALPHA) * prev + AGREE_EMA_ALPHA * agree;
        }
        isv[RL_Q_ARG_VS_PI_AGREE_INDEX] = out;
    }
}
