/*
 * target_drift_ema_update — Plan 4 follow-up to T5: GPU-driven RMS drift
 * between online and target params per branch (mag + dir), EMA-updated
 * into ISV slots. Replaces the legacy `per_branch_target_drift()` host
 * helper which used DtoH + CPU loop (anti-pattern per
 * pearl_cold_path_no_exception_to_gpu_drives.md).
 *
 * 2-block multi-reduction:
 *   Block 0 reduces magnitude-branch param slice → RMS(target - online)
 *           → EMA-updates ISV[TARGET_DRIFT_MAG_EMA_INDEX].
 *   Block 1 reduces direction-branch param slice → RMS(target - online)
 *           → EMA-updates ISV[TARGET_DRIFT_DIR_EMA_INDEX].
 *
 * Each block: 256-thread smem tree-reduce of (target[i] - online[i])²
 * over the slice, sqrt(mean) at thread 0, EMA into ISV slot. Pinned
 * device-mapped → host sees updated values without explicit DtoH.
 *
 * Adaptive α matches Plan 3 Task 3 / Task 7 / P4 Task 5 convention:
 *     α = α_base × (1 + 0.5 × |clamp(ISV[sharpe], -2, 2)|)
 *
 * Cold-path: launched once per epoch alongside attention_focus_ema.
 */

#include <cuda_runtime.h>

extern "C" __global__ void target_drift_ema_update(
    const float* __restrict__ online_params,        /* params_buf */
    const float* __restrict__ target_params,        /* target_params_buf */
    int          mag_branch_start_idx,              /* mag branch tensor 12..16 */
    int          mag_branch_len,
    int          dir_branch_start_idx,              /* dir branch tensor 8..12 */
    int          dir_branch_len,
    float*       __restrict__ isv_out,              /* pinned device-mapped */
    int          target_drift_mag_ema_idx,
    int          target_drift_dir_ema_idx,
    int          sharpe_ema_idx,
    float        alpha_base
) {
    /* Pick branch slice + ISV target for this block. */
    int          slice_start;
    int          slice_len;
    int          target_idx;
    switch (blockIdx.x) {
        case 0:
            slice_start = mag_branch_start_idx;
            slice_len   = mag_branch_len;
            target_idx  = target_drift_mag_ema_idx;
            break;
        case 1:
            slice_start = dir_branch_start_idx;
            slice_len   = dir_branch_len;
            target_idx  = target_drift_dir_ema_idx;
            break;
        default:
            return;
    }

    extern __shared__ float smem[];
    const int tid   = threadIdx.x;
    const int block = blockDim.x;

    /* Sum of squared differences over the slice (strided over threads). */
    float local_sumsq = 0.0f;
    if (slice_len > 0) {
        for (int i = tid; i < slice_len; i += block) {
            const float d = target_params[slice_start + i] - online_params[slice_start + i];
            local_sumsq += d * d;
        }
    }
    smem[tid] = local_sumsq;
    __syncthreads();

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

    if (tid == 0) {
        const float rms = (slice_len > 0) ? sqrtf(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 * rms;
    }
}
