/**
 * 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.
 *
 * Rewards and dones are f32 throughout.
 *
 * 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.
 *
 * Rewards are f32.
 *
 * Launch config: grid=(ceil(n/256), 1, 1), block=(256, 1, 1).
 * ====================================================================== */

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;
    }
}

/**
 * TD(λ) truncated return kernel (v8 D1).
 *
 * Computes G^λ = (1-λ) Σ_{n=1}^{N-1} λ^{n-1} G^(n) + λ^{N-1} G^(N)
 * where G^(n) = r_t + γr_{t+1} + ... + γ^{n-1}r_{t+n-1} + γ^n Q(s_{t+n})
 *
 * λ=0: reduces to 1-step TD. λ=1: Monte Carlo (N-step, no intermediate bootstrap).
 * Strict generalization of nstep_accumulate_kernel.
 *
 * Requires bootstrapped Q(s') values. When not available, pass raw_rewards
 * as q_next for a self-bootstrap approximation.
 *
 * Launch: grid=(ceil(N*L/256)), block=(256).
 */
extern "C" __global__ void td_lambda_kernel(
    const float* __restrict__ raw_rewards,
    const float* __restrict__ raw_dones,
    const float* __restrict__ q_next,
    float*       __restrict__ out_rewards,
    float*       __restrict__ out_dones,
    float gamma,
    float lambda,
    int max_trace,
    int L,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= N * L) return;

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

    float G_lambda = 0.0f;
    float weight_sum = 0.0f;
    float lambda_pow = 1.0f;
    float any_done = 0.0f;

    for (int n = 1; n <= max_trace; n++) {
        float G_n = 0.0f;
        float gamma_pow = 1.0f;
        int valid = 1;

        for (int k = 0; k < n && (t + k) < L; k++) {
            int off = base + t + k;
            float d_k = raw_dones[off];
            if (d_k > 0.5f && k > 0) { valid = 0; any_done = 1.0f; break; }
            G_n += gamma_pow * raw_rewards[off];
            gamma_pow *= gamma;
            if (d_k > 0.5f) { any_done = 1.0f; break; }
        }

        if (!valid || (t + n) >= L) {
            G_lambda += lambda_pow * G_n;
            weight_sum += lambda_pow;
            break;
        }

        G_n += gamma_pow * q_next[base + t + n - 1];

        float w = (n < max_trace) ? (1.0f - lambda) * lambda_pow : lambda_pow;
        G_lambda += w * G_n;
        weight_sum += w;
        lambda_pow *= lambda;
    }

    int out_idx = base + t;
    out_rewards[out_idx] = (weight_sum > 1e-8f) ? G_lambda / weight_sum : raw_rewards[out_idx];
    out_dones[out_idx] = any_done;
}
