// crates/ml/src/cuda_pipeline/popart_component_ema_kernel.cu
//
// SP11 B1b fix-up (2026-05-04): popart-component-specific magnitude EMA
// AND variance EMA (B1b smoke-recovery).
//
// Reads `popart_component_per_sample[N]` (an N-bar mapped-pinned f32
// buffer populated by `experience_env_step` at the `r_popart` write site)
// and computes:
//   scratch_out[0] = mean(|x|)   → ISV[POPART_COMPONENT_MAG_EMA_INDEX=360]
//   scratch_out[1] = var(|x|)    → ISV[REWARD_COMPONENT_VAR_EMA_BASE=361]
//
// The chained `apply_pearls_ad_kernel` launches on the same stream
// target ISV slots 360 (n_slots=1 mag) and 361 (n_slots=1 var) — see
// the two-singleton-chain in `launch_sp11_popart_component_ema`.
//
// Why both outputs in ONE kernel: a single-pass Welford via
// two-block-tree-reduces in the same kernel keeps the producer
// stream-ordered (no host sync between mag and var) and reuses the
// already-loaded |x| values across both reduces. The cost is a single
// extra `__syncthreads()` between the mean and variance passes.
//
// Why a separate kernel from `reward_component_ema_kernel`:
//   The pre-SP11 SP4 reward_component_ema kernel writes
//   `|reward_components_per_sample[+0]|` EMA to ISV[REWARD_POPART_EMA_INDEX=63].
//   That slot is overloaded — it is BOTH PopArt's normalisation-target
//   input (= total composed reward magnitude EMA) AND an attempted stand-
//   in for popart-component-specific magnitude. Pre-SP11 the two values
//   collapsed because reward composition was inline accumulation. B1b's
//   structural decomposition (per-component locals + Σ w_i × r_i)
//   surfaces the contamination — the SP11 mag-ratio canary reading slot
//   63 saw the inflated TOTAL reward magnitude as "popart's signal" and
//   redistributed the controller's weight toward popart (w_pop ≈ 2.0)
//   based on a contaminated ratio → 10× sharpe drop in B1b smoke. Spec
//   §4 amendment "Why slot 360" documents the resolution: a separate ISV
//   slot for popart-component-specific magnitude. Slot 63 is unchanged
//   (PopArt's input = total reward magnitude, pre-SP11 invariant
//   preserved); the SP11 mag-ratio canary reads slot 360 for the popart
//   axis (this kernel's output) and slots 64..68 for the other 5 axes.
//
// Why variance (B1b smoke-recovery, 2026-05-04, spec §4 amendment
// "Why z-score" lines 564-619): the linear mag-ratio formula
// `winner_weight = mag_ratio` amplified popart's intrinsic O(100)
// magnitude over the other 5 components' O(0.1-2). w_pop saturated
// toward MAX_WEIGHT, curiosity_b exploded, sharpe collapsed. Z-score
// makes ratios scale-invariant — every component contributes "shares
// of own scale". The variance EMA at slot 361 is the σ² denominator
// in the canary's z[0] = mag[0] / max(sqrt(var[0]), EPS_DIV).
//
// Single block (BLOCK_SIZE=256). Two block tree-reduces — no atomicAdd
// per `feedback_no_atomicadd.md`. Pure GPU compute per
// `feedback_no_cpu_compute_strict.md`.

#include <cuda_runtime.h>

#define BLOCK_SIZE 256

extern "C" __global__ void popart_component_ema_kernel(
    /* Per-bar popart-component-specific reward magnitude. Producer
     * `experience_env_step` writes the absolute popart-component value
     * (matching the `r_popart = capped_pnl;` voluntary-exit branch,
     * `0.0` everywhere else). Layout [n_samples] f32, mapped-pinned. */
    const float* __restrict__ popart_component_per_sample,
    /* Producer scratch buffer. The launcher passes a pointer offset to
     * `scratch[SCRATCH_SP11_POPART_COMPONENT_EMA]` so this kernel writes
     * to `scratch_out[0]` (mean) and `scratch_out[1]` (variance).
     * Chained `apply_pearls_ad_kernel` launches target ISV[360] (mag)
     * and ISV[361] (var) — see `launch_sp11_popart_component_ema`. */
    float* __restrict__ scratch_out,
    const int n_samples)
{
    /* Single-block reducer. */
    if (blockIdx.x != 0) return;
    extern __shared__ float sdata[];

    /* ── Pass 1: mean(|x|) ───────────────────────────────────────────
     * Strided accumulation over the per-bar buffer. Each thread folds
     * |popart_component_per_sample[i]| into its private accumulator,
     * then we tree-reduce in shared memory. */
    float acc = 0.0f;
    for (int i = threadIdx.x; i < n_samples; i += blockDim.x) {
        acc += fabsf(popart_component_per_sample[i]);
    }
    sdata[threadIdx.x] = acc;
    __syncthreads();

    /* Standard log2(BLOCK_SIZE) tree reduction — no atomicAdd. */
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            sdata[threadIdx.x] += sdata[threadIdx.x + s];
        }
        __syncthreads();
    }

    /* `fmaxf(1.0f, n_samples)` guards the impossible zero-bar edge case
     * (no producer ever writes; debug-only). All threads need `mean`
     * for the variance pass below — broadcast via shared mem. */
    __shared__ float mean_abs;
    if (threadIdx.x == 0) {
        mean_abs = sdata[0] / fmaxf(1.0f, (float)n_samples);
        scratch_out[0] = mean_abs;
    }
    __syncthreads();

    /* ── Pass 2: var(|x|) via Welford against the just-computed mean ──
     * var = (1/N) × Σ (|x_i| − mean)². Population variance (matches the
     * pre-SP11 `mean_abs` semantics — both are batch-scope statistics,
     * not sample-corrected estimators of an unknown population). */
    float v_acc = 0.0f;
    for (int i = threadIdx.x; i < n_samples; i += blockDim.x) {
        float d = fabsf(popart_component_per_sample[i]) - mean_abs;
        v_acc += d * d;
    }
    sdata[threadIdx.x] = v_acc;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s) {
            sdata[threadIdx.x] += sdata[threadIdx.x + s];
        }
        __syncthreads();
    }

    if (threadIdx.x == 0) {
        scratch_out[1] = sdata[0] / fmaxf(1.0f, (float)n_samples);
        __threadfence_system();
    }
}
