// rl_advantage_normalize.cu — Phase 4.5 per-batch advantage normalization (2026-05-30).
//
// Standard PPO practice (Schulman et al. 2017): normalize per-batch
// advantages before the PPO surrogate computation:
//
//   mean = (1/B) Σ_b advantage[b]
//   var  = (1/B) Σ_b (advantage[b] − mean)²
//   std  = sqrt(var + ε²)
//   advantage_norm[b] = (advantage[b] − mean) / std    (in-place)
//
// Why: PPO surrogate = ratio × A. Without normalization, |A| can vary
// 1000× across batches → l_pi magnitude unstable → Adam's per-parameter
// scaling still functional, but gradient direction confidence drops
// when advantage magnitudes are wildly inconsistent.
//
// In Phase 4.3, V_dq baseline produced advantages with ~100× more
// variance than Plan A v2's V_scalar baseline → l_pi grew to 1.6e9
// vs Plan A v2's 3.6e5. Adam internally normalizes, but the policy
// updates have higher variance per step → pnl trajectory chops.
//
// Per pearl_adaptive_not_tuned: this normalization is self-adaptive
// (uses observed per-batch statistics, no tuned hyperparameters).
// Per pearl_blend_formulas_must_have_permanent_floor: ε² floor on
// variance prevents div-by-zero when all advantages are identical.
//
// Block layout: grid=(1, 1, 1), block=(BLOCK_X=1024, 1, 1). Single
// block does parallel reduction for mean → variance → normalize.
// Suitable for batch sizes up to B=1024 (current production scale).

#define BLOCK_X 1024
#define ADVANTAGE_VAR_FLOOR 1e-6f

// Adaptive controller floors (spec 2026-05-30-adaptive-controller-floor-design):
// emit raw per-batch advantage variance to ISV[612] BEFORE the in-place
// normalize. `rl_rollout_steps_controller` consumes this as its driving
// signal — post-normalization variance is definitionally ~ε² floor under
// Phase 4.5 and useless as a controller input. Pre-norm variance is the
// real signal of "how much advantage spread the policy is producing".
#define RL_ADV_VAR_PRE_NORM_INDEX 612

extern "C" __global__ void rl_advantage_normalize(
    float* __restrict__ advantage,  // [B] in-place
    float* __restrict__ isv,        // ISV bus — only ADV_VAR_PRE_NORM written
    int    B
) {
    const int tid = threadIdx.x;
    if (tid >= BLOCK_X) return;

    // ── Pass 1: compute mean ──
    __shared__ float s_sum[BLOCK_X];
    float sum_partial = 0.0f;
    for (int b = tid; b < B; b += BLOCK_X) {
        sum_partial += advantage[b];
    }
    s_sum[tid] = sum_partial;
    __syncthreads();

    for (int stride = BLOCK_X / 2; stride > 0; stride >>= 1) {
        if (tid < stride) s_sum[tid] += s_sum[tid + stride];
        __syncthreads();
    }
    __shared__ float s_mean;
    if (tid == 0) s_mean = s_sum[0] / (float)B;
    __syncthreads();
    const float mean = s_mean;

    // ── Pass 2: compute variance ──
    __shared__ float s_var[BLOCK_X];
    float var_partial = 0.0f;
    for (int b = tid; b < B; b += BLOCK_X) {
        const float d = advantage[b] - mean;
        var_partial += d * d;
    }
    s_var[tid] = var_partial;
    __syncthreads();

    for (int stride = BLOCK_X / 2; stride > 0; stride >>= 1) {
        if (tid < stride) s_var[tid] += s_var[tid + stride];
        __syncthreads();
    }
    __shared__ float s_inv_std;
    if (tid == 0) {
        const float var = s_var[0] / (float)B;
        // Emit raw pre-normalization variance for rl_rollout_steps_controller.
        // Done BEFORE we add the ε² floor — the controller wants the true
        // signal, not the numerical-stability-adjusted divisor.
        isv[RL_ADV_VAR_PRE_NORM_INDEX] = var;
        s_inv_std = rsqrtf(var + ADVANTAGE_VAR_FLOOR);  // 1/sqrt(var + ε²)
    }
    __syncthreads();
    const float inv_std = s_inv_std;

    // ── Pass 3: normalize in-place ──
    for (int b = tid; b < B; b += BLOCK_X) {
        advantage[b] = (advantage[b] - mean) * inv_std;
    }
}
