/* =====================================================================
 * Q-bias correction kernel — ml-alpha RL training stability
 *
 * Tracks mean(Q_predicted - actual_return) as an EMA, writes a clamped
 * correction term into the ISV bus for downstream consumers (Bellman
 * target offset, diagnostic logging).
 *
 * Grid=(1), Block=(min(B, 256)). Tree-reduce for batch mean.
 * No atomicAdd — block tree-reduce in shared memory only.
 * No #include beyond cuda_runtime.h — pre-compiled cubin, no nvrtc.
 *
 * ISV slot layout:
 *   561 — Q_BIAS_EMA: exponential moving average of mean(Q - R)
 *   562 — Q_BIAS_CORRECTION: clamped correction term ∈ [-10, 10]
 *   563 — Q_BIAS_ALPHA: EMA smoothing rate (read-only, host-seeded)
 *
 * Pearl A bootstrap: on first call isv[Q_BIAS_EMA_INDEX] == 0.0 (sentinel);
 * thread 0 replaces directly with mean_diff rather than blending, per
 * pearl_first_observation_bootstrap.md.
 * ===================================================================== */

#include <cuda_runtime.h>

#define Q_BIAS_EMA_INDEX        561
#define Q_BIAS_CORRECTION_INDEX 562
#define Q_BIAS_ALPHA_INDEX      563

extern "C" __global__ void rl_q_bias_correction(
    const float* __restrict__ q_predicted,   /* [B] — max Q at sampled action  */
    const float* __restrict__ actual_return,  /* [B] — n-step return from PER   */
    float* __restrict__ isv,
    int b_size
) {
    __shared__ float sdata[256];

    int tid = threadIdx.x;
    int block_size = blockDim.x;

    /* ── Phase 1: each thread accumulates its slice of the batch ── */
    float local_sum = 0.0f;
    for (int i = tid; i < b_size; i += block_size) {
        local_sum += q_predicted[i] - actual_return[i];
    }
    sdata[tid] = local_sum;
    __syncthreads();

    /* ── Phase 2: block tree-reduce in shared memory ── */
    for (int stride = block_size / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sdata[tid] += sdata[tid + stride];
        }
        __syncthreads();
    }

    /* ── Phase 3: thread 0 — EMA update + correction emit ── */
    if (tid == 0) {
        float mean_diff = sdata[0] / fmaxf((float)b_size, 1.0f);

        float alpha = isv[Q_BIAS_ALPHA_INDEX];
        /* Defensive: if alpha not yet seeded (sentinel 0), use 0.01. */
        if (alpha < 1e-9f) {
            alpha = 0.01f;
        }

        float old_ema = isv[Q_BIAS_EMA_INDEX];

        /* Pearl A first-observation bootstrap: sentinel 0 → replace
         * directly; avoids contaminating the EMA with a blend towards
         * an uninformative zero initial value. */
        float bias_ema;
        if (old_ema > -1e-9f && old_ema < 1e-9f) {
            /* Sentinel — first observation. */
            bias_ema = mean_diff;
        } else {
            bias_ema = (1.0f - alpha) * old_ema + alpha * mean_diff;
        }

        isv[Q_BIAS_EMA_INDEX] = bias_ema;

        /* Correction: bilateral clamp per pearl_symmetric_clamp_audit. */
        float correction = fmaxf(-10.0f, fminf(bias_ema, 10.0f));
        isv[Q_BIAS_CORRECTION_INDEX] = correction;
    }
}
