/**
 * N-step return accumulation kernel.
 *
 * Converts 1-step transitions (s_t, a_t, r_t, s_{t+1}, done_t) into
 * n-step transitions (s_t, a_t, R_n, s_{t+n}, done_n) for Rainbow DQN.
 *
 * R_n = r_t + gamma*r_{t+1} + ... + gamma^{n-1}*r_{t+n-1}
 * done_n = 1 if any step in [t..t+n-1] was terminal
 *
 * Reads from raw_* (copy of original 1-step data), writes to out_*
 * (overwritten in-place). Double-buffering eliminates race conditions.
 *
 * Launch config: grid=(ceil(N*L/256), 1, 1), block=(256, 1, 1).
 */

extern "C" __global__ void nstep_accumulate_kernel(
    const float* __restrict__ raw_rewards,  /* [N * L] original 1-step rewards */
    const float* __restrict__ raw_dones,    /* [N * L] original 1-step dones */
    float* __restrict__ out_rewards,        /* [N * L] overwritten with R_n */
    float* __restrict__ out_dones,          /* [N * L] overwritten with done_n */
    float gamma,
    int n_steps,
    int L,          /* timesteps per episode */
    int N           /* number of episodes */
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * L) return;

    int ep = idx / L;
    int t  = idx % L;

    float R_n      = 0.0f;
    float gamma_pow = 1.0f;
    float any_done = 0.0f;

    for (int i = 0; i < n_steps; i++) {
        int step = t + i;
        if (step >= L) break;

        int off = ep * L + step;
        float r_i = raw_rewards[off];
        float d_i = raw_dones[off];

        R_n += gamma_pow * r_i;
        gamma_pow *= gamma;

        if (d_i > 0.5f) {
            any_done = 1.0f;
            break;  /* Episode terminated — stop accumulating */
        }
    }

    int base = ep * L + t;
    out_rewards[base] = R_n;
    out_dones[base]   = any_done;
}

/* ══════════════════════════════════════════════════════════════════════
 * REWARD NORMALIZATION KERNEL
 *
 * reward[i] = (reward[i] - mean) * inv_std
 *
 * Normalizes rewards in-place using pre-computed batch statistics.
 * This puts dense shaping (0.01x) and sparse trade-completion (±2.0)
 * rewards on the same scale for C51 distributional learning.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

/* regime_scale_td_errors — authoritative version in dqn_utility_kernels.cu.
 * Removed duplicate from this file. The trainer uses the dqn_utility version
 * which has STATE_DIM as a compile-time constant for tensor core alignment. */

extern "C" __global__ void reward_normalize_kernel(
    float* __restrict__ rewards,
    float mean,
    float inv_std,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n) rewards[i] = (rewards[i] - mean) * inv_std;
}
