// rl_ppo_diagnostic_stats_reduce.cu — B-10 G3+G4: PPO advantage
// normalization + surrogate decomposition observability. Consumes 5
// per-batch scratch arrays written by `ppo_clipped_surrogate.cu` (which
// gains 5 new pointer params alongside the existing `loss_per_batch[B]`
// write) and emits 8 ISV slots:
//
//   735 RL_PPO_A_NORM_ABS_MAX_INDEX     — max |A_norm|       across B
//   736 RL_PPO_A_NORM_ABS_MEAN_INDEX    — mean |A_norm|      across B
//   737 RL_PPO_A_UNNORM_ABS_MAX_INDEX   — max |A_unnorm|     across B
//   738 RL_PPO_A_UNNORM_ABS_MEAN_INDEX  — mean |A_unnorm|    across B
//   739 RL_PPO_A_SIGMA_USED_INDEX       — sqrt(isv[612]) — the per-batch
//                                         std actually applied as divisor
//                                         in `rl_advantage_normalize.cu`
//                                         (Schulman 2017 canonical
//                                         per-batch standardization).
//   740 RL_PPO_RATIO_DEV_ABS_MEAN_INDEX — mean |ratio - 1|   across B
//   741 RL_PPO_RATIO_CLIP_RATE_INDEX    — fraction outside [1-ε, 1+ε]
//                                         (sum of 0/1 flags / B)
//   742 RL_PPO_L_SURROGATE_INDEX        — mean clipped surrogate
//                                         (without entropy bonus)
//
// Pattern lifted from `rl_bellman_target_saturation_reduce.cu`: single
// block, grid-stride loop, separate `__shared__` arrays per reduction.
// Per `pearl_no_atomicadd`: no atomicAdd anywhere. Per
// `feedback_no_partial_refactor`: the new scratches in
// `ppo_clipped_surrogate.cu` are written alongside `loss_per_batch`
// (untouched); `ppo_loss_reduce_b.cu` and `ppo_log_ratio_abs_max_b.cu`
// continue to operate on their existing inputs.
//
// Caller launches with block_dim = next_pow2(B).min(256), grid = (1,1,1),
// shared bytes = 7 * block_dim * sizeof(float).

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

#define RL_ADV_VAR_PRE_NORM_INDEX            612
#define RL_PPO_A_NORM_ABS_MAX_INDEX          735
#define RL_PPO_A_NORM_ABS_MEAN_INDEX         736
#define RL_PPO_A_UNNORM_ABS_MAX_INDEX        737
#define RL_PPO_A_UNNORM_ABS_MEAN_INDEX       738
#define RL_PPO_A_SIGMA_USED_INDEX            739
#define RL_PPO_RATIO_DEV_ABS_MEAN_INDEX      740
#define RL_PPO_RATIO_CLIP_RATE_INDEX         741
#define RL_PPO_L_SURROGATE_INDEX             742

extern "C" __global__ void rl_ppo_diagnostic_stats_reduce(
    float*       __restrict__ isv,
    const float* __restrict__ a_norm_per_batch,     // [B] |A_norm|
    const float* __restrict__ ratio_dev_per_batch,  // [B] |ratio - 1|
    const float* __restrict__ ratio_clipped_pb,     // [B] {0.0, 1.0}
    const float* __restrict__ surrogate_per_batch,  // [B] clipped surrogate
    int                       B
) {
    // 5 reductions × block_dim floats. Separate arrays to avoid any
    // cross-talk that could perturb determinism (per spec §3.4 + §4 V3).
    // σ_used is a per-step scalar (read once at end from slot 612);
    // |A_unnorm| stats are reconstructed from |A_norm| × σ_used in the
    // final write — no per-batch unnorm scratch needed.
    extern __shared__ float smem[];
    float* s_an_max = smem;
    float* s_an_sum = smem +     blockDim.x;
    float* s_rd_sum = smem + 2 * blockDim.x;
    float* s_rc_sum = smem + 3 * blockDim.x;
    float* s_su_sum = smem + 4 * blockDim.x;
    const int tid = threadIdx.x;

    // Grid-stride per-thread folds.
    float l_an_max = 0.0f, l_an_sum = 0.0f;
    float l_rd_sum = 0.0f, l_rc_sum = 0.0f, l_su_sum = 0.0f;
    for (int b = tid; b < B; b += blockDim.x) {
        const float an = a_norm_per_batch[b];
        l_an_max  = fmaxf(l_an_max, an);
        l_an_sum += an;
        l_rd_sum += ratio_dev_per_batch[b];
        l_rc_sum += ratio_clipped_pb[b];
        l_su_sum += surrogate_per_batch[b];
    }
    s_an_max[tid] = l_an_max;
    s_an_sum[tid] = l_an_sum;
    s_rd_sum[tid] = l_rd_sum;
    s_rc_sum[tid] = l_rc_sum;
    s_su_sum[tid] = l_su_sum;
    __syncthreads();

    // Power-of-2 tree reduce. blockDim.x is guaranteed power of 2 by caller.
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            s_an_max[tid]  = fmaxf(s_an_max[tid], s_an_max[tid + stride]);
            s_an_sum[tid] += s_an_sum[tid + stride];
            s_rd_sum[tid] += s_rd_sum[tid + stride];
            s_rc_sum[tid] += s_rc_sum[tid + stride];
            s_su_sum[tid] += s_su_sum[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) {
        const float denom = (float)B;
        const float var_pre_norm = isv[RL_ADV_VAR_PRE_NORM_INDEX];
        // σ_used = sqrt(var_pre_norm). Guard against tiny negative
        // residuals from float ε.
        const float sigma_used   = sqrtf(fmaxf(var_pre_norm, 0.0f));
        const float a_norm_max   = s_an_max[0];
        const float a_norm_mean  = s_an_sum[0] / denom;
        isv[RL_PPO_A_NORM_ABS_MAX_INDEX]     = a_norm_max;
        isv[RL_PPO_A_NORM_ABS_MEAN_INDEX]    = a_norm_mean;
        // |A_unnorm| = |A_norm| × σ_used (σ_used is constant across batch
        // by construction in rl_advantage_normalize); reconstruction
        // preserves max/mean structure without a per-batch scratch.
        isv[RL_PPO_A_UNNORM_ABS_MAX_INDEX]   = a_norm_max  * sigma_used;
        isv[RL_PPO_A_UNNORM_ABS_MEAN_INDEX]  = a_norm_mean * sigma_used;
        isv[RL_PPO_A_SIGMA_USED_INDEX]       = sigma_used;
        isv[RL_PPO_RATIO_DEV_ABS_MEAN_INDEX] = s_rd_sum[0] / denom;
        isv[RL_PPO_RATIO_CLIP_RATE_INDEX]    = s_rc_sum[0] / denom;
        isv[RL_PPO_L_SURROGATE_INDEX]        = s_su_sum[0] / denom;
    }
}
