// rl_q_distribution_stats.cu — B-10 G1: Q-distribution informativeness
// diagnostic. Detects C51 Q-collapse (target → one-hot at center atom
// under EWMA-narrowed atom support) via Q output STRUCTURE rather than
// loss magnitude. A bit-flat `l_q ≈ 0.002` is a known false-positive of
// adaptive-support categorical DQN (per literature 2026-06-01 perplexity:
// "near-zero categorical cross-entropy under adaptive support is a known
// Q-collapse symptom"), so we measure Q's information content directly:
//
//   * `q_dist_entropy_mean` — mean over (b, a) of −Σ_z p log p (where p
//     is softmax over Q_N_ATOMS atoms for that action). Near 0 ⇒ Q
//     concentrates on one atom (could be informative or collapsed-at-edge).
//     Near ln(Q_N_ATOMS) ≈ 3.045 ⇒ near-uniform Q (information-free).
//
//   * `q_value_range_mean` — mean over b of (max_a E_Q(s,a) − min_a E_Q(s,a)).
//     E_Q is the expected return Σ_z p · atom_value[z]. Range near 0 ⇒
//     Q has no preference among actions.
//
//   * `q_value_abs_max` — per-step max over (b, a) of |E_Q|. Surfaces
//     whether Q's expected values are scale-consistent with realized
//     returns or have drifted.
//
// Two-kernel pattern per `pearl_no_atomicadd` (no atomicAdd anywhere)
// and per the B-9 saturation-reduce precedent:
//   1. `rl_q_distribution_per_batch` — grid (B, 1, 1), block (Q_N_ATOMS, 1, 1).
//      Each block processes one batch element across all N_ACTIONS actions
//      sequentially; computes per-action softmax + entropy + E_Q in shared
//      mem; writes per-batch entropy_mean / range / abs_max to 3 [B]
//      scratch arrays.
//   2. `rl_q_distribution_reduce` — single-block tree-reduce over the [B]
//      scratch arrays; writes ISV slots 730 / 731 / 732.
//
// Inputs (online Q logits ONLY — target Q is not consumed by Q-Thompson
// or Q→π distill, so its informativeness is irrelevant to the cascade
// we're diagnosing):
//   q_logits        [B × N_ACTIONS × Q_N_ATOMS] — raw atom logits.
//   atom_supports   [Q_N_ATOMS]                 — atom value array
//                                                 (V_MIN + i·Δz, ISV-driven).
//   B                int                         — batch size.
//
// Outputs (per-batch scratch consumed by reducer below):
//   entropy_per_batch    [B]
//   range_per_batch      [B]
//   abs_max_per_batch    [B]
//
// Per `feedback_cpu_is_read_only`: no HtoD, no host compute. Per
// `feedback_no_nvrtc`: pre-compiled cubin. Per
// `feedback_no_htod_htoh_only_mapped_pinned`: device-only reads/writes.

#include <cuda_runtime.h>
#include <math_constants.h>

#define Q_N_ATOMS                              21
#define N_ACTIONS                              11

#define RL_Q_DIST_ENTROPY_MEAN_INDEX           730
#define RL_Q_VALUE_RANGE_MEAN_INDEX            731
#define RL_Q_VALUE_ABS_MAX_INDEX               732

// ─────────────────────────────────────────────────────────────────────
// Per-batch kernel: grid (B, 1, 1), block (Q_N_ATOMS, 1, 1).
// Each block processes ONE batch element by iterating over all
// N_ACTIONS actions sequentially (cheap — N_ACTIONS=11). For each
// action, compute softmax over Q_N_ATOMS atoms (parallel across the
// 21 threads), then entropy and expected value E_Q. Track running
// max_a E_Q, min_a E_Q, max_a |E_Q|, and Σ_a entropy in shared mem.
// Thread 0 writes the 3 per-batch outputs at the end.
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_q_distribution_per_batch(
    const float* __restrict__ q_logits,           // [B × N_ACTIONS × Q_N_ATOMS]
    const float* __restrict__ atom_supports,      // [Q_N_ATOMS]
    int                       B,
    float*       __restrict__ entropy_per_batch,  // [B]
    float*       __restrict__ range_per_batch,    // [B]
    float*       __restrict__ abs_max_per_batch   // [B]
) {
    const int batch = blockIdx.x;
    const int atom  = threadIdx.x;
    if (batch >= B || atom >= Q_N_ATOMS) return;

    // Shared mem for per-action softmax computation. Same layout as
    // `bellman_target_projection.cu` (single per-action reduction).
    __shared__ float s_logits[Q_N_ATOMS];
    __shared__ float s_softmax[Q_N_ATOMS];
    __shared__ float s_max;
    __shared__ float s_sumexp;

    // Per-action accumulators on thread 0; per-block aggregates.
    __shared__ float s_e_q[N_ACTIONS];
    __shared__ float s_entropy_sum;

    // Atom value cached once for this block (read-only, all threads).
    __shared__ float s_atom_value[Q_N_ATOMS];
    s_atom_value[atom] = atom_supports[atom];
    if (atom == 0) {
        s_entropy_sum = 0.0f;
    }
    __syncthreads();

    // Iterate over all actions for this batch element.
    for (int a = 0; a < N_ACTIONS; ++a) {
        // ── 1. Load logits for (batch, a) into shared mem ─────────────
        const long long base = (long long)batch * N_ACTIONS * Q_N_ATOMS
                             + (long long)a   * Q_N_ATOMS;
        s_logits[atom] = q_logits[base + atom];
        __syncthreads();

        // ── 2. Softmax over atoms (numerically-stable max-subtract) ───
        if (atom == 0) {
            float m = s_logits[0];
            #pragma unroll
            for (int z = 1; z < Q_N_ATOMS; ++z) m = fmaxf(m, s_logits[z]);
            s_max = m;
        }
        __syncthreads();

        const float e = expf(s_logits[atom] - s_max);
        s_softmax[atom] = e;
        __syncthreads();

        if (atom == 0) {
            float sum = 0.0f;
            #pragma unroll
            for (int z = 0; z < Q_N_ATOMS; ++z) sum += s_softmax[z];
            s_sumexp = sum;
        }
        __syncthreads();

        const float p = s_softmax[atom] / s_sumexp;

        // ── 3. Reduce to E_Q and entropy on thread 0 ──────────────────
        if (atom == 0) {
            float e_q = 0.0f;
            float ent = 0.0f;
            #pragma unroll
            for (int z = 0; z < Q_N_ATOMS; ++z) {
                const float pz = s_softmax[z] / s_sumexp;
                e_q += pz * s_atom_value[z];
                // entropy: -p log p; safe for p>0 since softmax is positive.
                if (pz > 1e-30f) {
                    ent -= pz * logf(pz);
                }
            }
            s_e_q[a]       = e_q;
            s_entropy_sum += ent;
        }
        __syncthreads();
    }

    // ── 4. Per-batch aggregates: thread 0 writes scratch ──────────────
    if (atom == 0) {
        float max_q   = s_e_q[0];
        float min_q   = s_e_q[0];
        float abs_max = fabsf(s_e_q[0]);
        #pragma unroll
        for (int a = 1; a < N_ACTIONS; ++a) {
            const float q = s_e_q[a];
            max_q   = fmaxf(max_q, q);
            min_q   = fminf(min_q, q);
            abs_max = fmaxf(abs_max, fabsf(q));
        }
        // Per-batch entropy averaged across actions (so 0..ln(Q_N_ATOMS)
        // range is preserved; the cross-batch reducer averages across B).
        entropy_per_batch[batch] = s_entropy_sum / (float)N_ACTIONS;
        range_per_batch[batch]   = max_q - min_q;
        abs_max_per_batch[batch] = abs_max;
    }
}

// ─────────────────────────────────────────────────────────────────────
// Cross-batch reducer: single block, grid-stride loop. Writes the 3
// per-step diag values to ISV. Pattern lifted from
// `rl_bellman_target_saturation_reduce.cu`.
//
// Caller passes block_dim = next_pow2(B).min(256), grid = (1, 1, 1),
// shared bytes = 3 * block_dim * sizeof(float).
// ─────────────────────────────────────────────────────────────────────
extern "C" __global__ void rl_q_distribution_reduce(
    float*       __restrict__ isv,
    const float* __restrict__ entropy_per_batch,
    const float* __restrict__ range_per_batch,
    const float* __restrict__ abs_max_per_batch,
    int                       B
) {
    extern __shared__ float smem[];
    float* s_ent_sum = smem;
    float* s_rng_sum = smem +     blockDim.x;
    float* s_abs_max = smem + 2 * blockDim.x;
    const int tid = threadIdx.x;

    float l_ent = 0.0f, l_rng = 0.0f;
    float l_abs = -CUDART_INF_F;
    for (int b = tid; b < B; b += blockDim.x) {
        l_ent += entropy_per_batch[b];
        l_rng += range_per_batch[b];
        l_abs  = fmaxf(l_abs, abs_max_per_batch[b]);
    }
    s_ent_sum[tid] = l_ent;
    s_rng_sum[tid] = l_rng;
    s_abs_max[tid] = l_abs;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_ent_sum[tid] += s_ent_sum[tid + stride];
            s_rng_sum[tid] += s_rng_sum[tid + stride];
            s_abs_max[tid]  = fmaxf(s_abs_max[tid], s_abs_max[tid + stride]);
        }
        __syncthreads();
    }

    if (tid == 0) {
        const float denom = (float)B;
        isv[RL_Q_DIST_ENTROPY_MEAN_INDEX] = s_ent_sum[0] / denom;
        isv[RL_Q_VALUE_RANGE_MEAN_INDEX]  = s_rng_sum[0] / denom;
        isv[RL_Q_VALUE_ABS_MAX_INDEX]     = s_abs_max[0];
    }
}
