// crates/ml-alpha/cuda/rl_popart_v_correct.cu
//
// PopArt V-head output correction kernel (van Hasselt et al., 2016).
//
// After `rl_popart_normalize` updates the running mean/sigma in ISV,
// the V-head's pre-existing predictions need an affine correction to
// account for the shift in normalization statistics. Without this,
// the value function has a discontinuity at every normalization update.
//
// Correction formula per element:
//   v_pred[b] = (sigma_old / sigma_new) * v_pred[b]
//             + (mean_old - mean_new) / sigma_new
//
// Grid=(ceil(B/32)), Block=(32).
//
// No shared memory, no reductions, no atomicAdd.
// Pre-compiled cubin — no nvrtc.

#include <cuda_runtime.h>

#define POPART_MEAN_INDEX                553
#define POPART_SIGMA_INDEX              555
#define POPART_SIGMA_OLD_INDEX          556
#define POPART_MEAN_OLD_INDEX           557
// Phase 3B-Y companion gate: when PopArt normalization is disabled
// (slot 793 = 0), V-correction MUST also be skipped. The running
// mean/sigma stats still update against raw reward magnitudes, so
// the correction formula would scale v_pred against an out-of-band
// scale that no longer matches the V regression target.
#define RL_POPART_NORMALIZE_ENABLED_INDEX 793

extern "C" __global__ void rl_popart_v_correct(
    float* __restrict__ v_pred,     // [B] IN/OUT
    const float* __restrict__ isv,
    int b_size
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= b_size) return;

    // Phase 3B-Y: skip correction when popart normalization is off.
    // Reward stream is in raw pnl-units; V is trained against the raw
    // returns directly; the correction would introduce a unit mismatch.
    if (isv[RL_POPART_NORMALIZE_ENABLED_INDEX] <= 0.5f) return;

    float sigma_old = isv[POPART_SIGMA_OLD_INDEX];
    float sigma_new = isv[POPART_SIGMA_INDEX];
    float mean_old  = isv[POPART_MEAN_OLD_INDEX];
    float mean_new  = isv[POPART_MEAN_INDEX];

    // Guard: if sigma_old is still at sentinel (pre-first-observation),
    // no correction is meaningful — the V-head has not been trained under
    // a previous normalization regime yet.
    if (sigma_old <= 1e-12f) return;

    // Guard: sigma_new is always >= sqrt(1e-6) from the normalize kernel,
    // but belt-and-suspenders.
    float inv_sigma_new = 1.0f / fmaxf(sigma_new, 1e-6f);

    v_pred[b] = (sigma_old * inv_sigma_new) * v_pred[b]
              + (mean_old - mean_new) * inv_sigma_new;
}
