/*
 * attention_focus_ema_update — Plan 4 Task 5 Mode A (E.5 attention-focus
 * interpretability). Fully GPU-driven per spec §4.C.6.
 *
 * Single kernel, 3-block multi-reduction:
 *   Block 0 reduces the magnitude-branch VSN-weight slice of the params buffer
 *           → mean |w| → EMA-updates ISV[VSN_MAG_EMA_INDEX].
 *   Block 1 reduces the direction-branch VSN-weight slice
 *           → mean |w| → EMA-updates ISV[VSN_DIR_EMA_INDEX].
 *   Block 2 reduces the Mamba2 enriched-hidden buffer
 *           → mean |state_transition| → EMA-updates ISV[MAMBA2_RETENTION_EMA_INDEX].
 *
 * Each block does a smem tree-reduction (256 threads). Thread 0 of each
 * block reads the previous EMA value, computes the adaptive α, and writes
 * the updated EMA back to its ISV slot. ISV is pinned device-mapped, so
 * the host sees the new values without an explicit DtoH.
 *
 * Adaptive α convention matches Plan 3 Task 3 / Task 4 / Task 7:
 *     α = α_base × (1 + 0.5 × |clamp(ISV[sharpe], -2, 2)|)
 *
 * Cold-path: launched once per epoch boundary alongside the other Plan 3
 * ISV producers. No DtoH, no host scalar computation, no atomicAdd.
 */

#include <cuda_runtime.h>

extern "C" __global__ void attention_focus_ema_update(
    const float* __restrict__ params_ptr,           /* DQN params buffer */
    int          mag_vsn_start_idx,                 /* slice for mag VSN */
    int          mag_vsn_len,
    int          dir_vsn_start_idx,                 /* slice for dir VSN */
    int          dir_vsn_len,
    const float* __restrict__ mamba2_h_enriched,    /* [B*SH2] live buffer */
    int          mamba2_n_elements,
    float*       __restrict__ isv_out,              /* pinned device-mapped */
    int          vsn_mag_ema_idx,
    int          vsn_dir_ema_idx,
    int          mamba2_retention_ema_idx,
    int          sharpe_ema_idx,
    float        alpha_base
) {
    /* Pick the slice and target ISV slot for this block. */
    const float* slice_ptr;
    int          slice_len;
    int          target_idx;
    switch (blockIdx.x) {
        case 0:
            slice_ptr  = (mag_vsn_len > 0) ? params_ptr + mag_vsn_start_idx : NULL;
            slice_len  = mag_vsn_len;
            target_idx = vsn_mag_ema_idx;
            break;
        case 1:
            slice_ptr  = (dir_vsn_len > 0) ? params_ptr + dir_vsn_start_idx : NULL;
            slice_len  = dir_vsn_len;
            target_idx = vsn_dir_ema_idx;
            break;
        case 2:
            slice_ptr  = (mamba2_n_elements > 0) ? mamba2_h_enriched : NULL;
            slice_len  = mamba2_n_elements;
            target_idx = mamba2_retention_ema_idx;
            break;
        default:
            return;
    }

    /* Smem tree reduction: each thread accumulates a stride, then reduce. */
    extern __shared__ float smem[];
    const int tid   = threadIdx.x;
    const int block = blockDim.x;

    float local_sum = 0.0f;
    if (slice_ptr != NULL) {
        for (int i = tid; i < slice_len; i += block) {
            local_sum += fabsf(slice_ptr[i]);
        }
    }
    smem[tid] = local_sum;
    __syncthreads();

    for (int s = block / 2; s > 0; s >>= 1) {
        if (tid < s) smem[tid] += smem[tid + s];
        __syncthreads();
    }

    if (tid == 0) {
        const float mean = (slice_len > 0) ? (smem[0] / (float)slice_len) : 0.0f;

        const float sharpe = fmaxf(-2.0f, fminf(2.0f, isv_out[sharpe_ema_idx]));
        const float alpha  = alpha_base * (1.0f + 0.5f * fabsf(sharpe));
        const float prev   = isv_out[target_idx];
        isv_out[target_idx] = (1.0f - alpha) * prev + alpha * mean;
    }
}
