// rl_bellman_target_saturation_reduce.cu — B-9 cross-batch tree-reduce of
// per-batch saturation tallies produced by `bellman_target_projection` and
// `bellman_fused_select_project`. Writes per-step rates + extremes to ISV.
//
// Inputs (consumed AFTER one of the two bellman variants has run):
//   sat_top_per_batch  [B]  — per-block count of (atom_z) where t_z > V_MAX_eff
//   sat_bot_per_batch  [B]  — per-block count of (atom_z) where t_z < V_MIN_eff
//   max_pre_per_batch  [B]  — per-block max(t_z)  pre-clamp
//   min_pre_per_batch  [B]  — per-block min(t_z)  pre-clamp
//
// Outputs (single ISV write each, by thread 0):
//   ISV[RL_Q_TARGET_TOP_SATURATION_RATE_INDEX = 726]  Σ(top) / (B × Q_N_ATOMS)
//   ISV[RL_Q_TARGET_BOT_SATURATION_RATE_INDEX = 727]  Σ(bot) / (B × Q_N_ATOMS)
//   ISV[RL_Q_TARGET_MAX_PRE_PROJ_INDEX        = 728]  max over batches
//   ISV[RL_Q_TARGET_MIN_PRE_PROJ_INDEX        = 729]  min over batches
//
// Per `feedback_no_atomicadd.md`: single-block kernel, shared-mem
// tree-reduce. Caller launches with grid_dim=(1,1,1), block_dim=(N,1,1)
// where N is a power of 2 with N <= max(1, B/2) and N <= 256, and
// shared bytes = 4 * N * sizeof(float).
//
// Block size rationale: each thread iterates B/N elements via grid-stride;
// the tree-reduce after the strided gather is over N threads.

#include <cuda_runtime.h>

#define RL_Q_TARGET_TOP_SATURATION_RATE_INDEX  726
#define RL_Q_TARGET_BOT_SATURATION_RATE_INDEX  727
#define RL_Q_TARGET_MAX_PRE_PROJ_INDEX         728
#define RL_Q_TARGET_MIN_PRE_PROJ_INDEX         729

#define Q_N_ATOMS                              21

extern "C" __global__ void rl_bellman_target_saturation_reduce(
    float*       __restrict__ isv,
    const float* __restrict__ sat_top_per_batch,
    const float* __restrict__ sat_bot_per_batch,
    const float* __restrict__ max_pre_per_batch,
    const float* __restrict__ min_pre_per_batch,
    int                       B
) {
    extern __shared__ float smem[];
    float* s_top = smem;
    float* s_bot = smem +     blockDim.x;
    float* s_max = smem + 2 * blockDim.x;
    float* s_min = smem + 3 * blockDim.x;
    const int tid = threadIdx.x;

    // Grid-stride gather. Padded identity values for max/min reductions
    // (+/- infinity) ensure threads with no work don't pollute the
    // result.
    float l_top = 0.0f, l_bot = 0.0f;
    float l_max = -INFINITY, l_min = INFINITY;
    for (int b = tid; b < B; b += blockDim.x) {
        l_top += sat_top_per_batch[b];
        l_bot += sat_bot_per_batch[b];
        l_max  = fmaxf(l_max, max_pre_per_batch[b]);
        l_min  = fminf(l_min, min_pre_per_batch[b]);
    }
    s_top[tid] = l_top;
    s_bot[tid] = l_bot;
    s_max[tid] = l_max;
    s_min[tid] = l_min;
    __syncthreads();

    // Power-of-2 tree reduce — blockDim.x guaranteed power of 2 by caller.
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_top[tid] += s_top[tid + stride];
            s_bot[tid] += s_bot[tid + stride];
            s_max[tid]  = fmaxf(s_max[tid], s_max[tid + stride]);
            s_min[tid]  = fminf(s_min[tid], s_min[tid + stride]);
        }
        __syncthreads();
    }

    if (tid == 0) {
        const float denom = (float)(B * Q_N_ATOMS);
        isv[RL_Q_TARGET_TOP_SATURATION_RATE_INDEX] = s_top[0] / denom;
        isv[RL_Q_TARGET_BOT_SATURATION_RATE_INDEX] = s_bot[0] / denom;
        isv[RL_Q_TARGET_MAX_PRE_PROJ_INDEX]        = s_max[0];
        isv[RL_Q_TARGET_MIN_PRE_PROJ_INDEX]        = s_min[0];
    }
}
