/* =====================================================================
 * rl_per_update_priority.cu — GPU-resident PER: update leaf priorities
 *
 * Grid=(1), Block=(b_size). After a training step, writes new priorities
 * (from absolute TD-errors) into the sum-tree leaf positions.
 *
 * No atomicAdd — max_priority reduction uses block-wide tree-reduce in
 * shared memory.
 *
 * ISV reads:
 *   isv[RL_PER_ALPHA_INDEX=405] — prioritization exponent α
 * ===================================================================== */

#define RL_PER_ALPHA_INDEX 405

extern "C" __global__ void rl_per_update_priority(
    const unsigned int* __restrict__ sample_indices, /* [B] */
    const float* __restrict__ td_errors,             /* [B] absolute values */
    float* __restrict__ priority_tree,               /* [2 * capacity] */
    float* __restrict__ max_priority,                /* [1] */
    const float* __restrict__ isv,
    int b_size,
    int capacity
)
{
    extern __shared__ float smax[];

    const int b = threadIdx.x;
    if (b >= b_size) return;

    /* ── Read alpha from ISV ──────────────────────────────────────────── */
    const float alpha = isv[RL_PER_ALPHA_INDEX];

    /* ── Compute priority = (|td_error| + ε)^α ───────────────────────── */
    const float td_abs = fabsf(td_errors[b]);
    const float priority = powf(td_abs + 1e-6f, alpha);

    /* ── Write priority to tree leaf ──────────────────────────────────── */
    const int leaf = (int)sample_indices[b];
    priority_tree[capacity + leaf] = priority;

    /* ── Block-wide tree-reduce to find max priority ──────────────────── */
    smax[b] = priority;
    __syncthreads();

    /* Power-of-two tree reduce (handles non-power-of-two b_size) */
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (b < stride && (b + stride) < b_size) {
            const float other = smax[b + stride];
            if (other > smax[b]) smax[b] = other;
        }
        __syncthreads();
    }

    /* Thread 0 writes the block max to global max_priority */
    if (b == 0) {
        const float old_max = max_priority[0];
        const float new_max = smax[0];
        max_priority[0] = (new_max > old_max) ? new_max : old_max;
    }
}
