/**
 * HER Episode Boundary Kernels
 *
 * GPU-native episode boundary tracking to enable HER Future and Final
 * strategies. All three kernels operate on a flat [N * L] buffer where
 * episode `ep` occupies indices [ep*L, (ep+1)*L).
 *
 * Zero CPU in the hot path — episode IDs filled by GPU kernel.
 */

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 1: fill_episode_ids_kernel
 *
 * Fills episode_ids[i] = i / L for each transition index i.
 * Direct arithmetic — no linear scan.
 *
 * Launch config: grid=(ceil(total/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void fill_episode_ids_kernel(
    int* __restrict__ episode_ids,  /* [N * L] output */
    int L,                          /* timesteps per episode */
    int total                       /* N * L */
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total) return;
    episode_ids[i] = i / L;
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 2: her_sample_future_donors
 *
 * HER Future strategy: for each source transition, sample a random
 * donor transition LATER in the same episode.
 *
 * Episode ep occupies [ep*L, (ep+1)*L) in the flat buffer.
 * Samples uniformly from [src+1, ep_end) using per-thread LCG.
 * Fallback to src itself when the source is the last step of its episode.
 *
 * Launch config: grid=(ceil(her_batch_size/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void her_sample_future_donors(
    const int* __restrict__ episode_ids,    /* [capacity] */
    const int* __restrict__ source_indices, /* [her_batch_size] */
    int* __restrict__ donor_indices,        /* [her_batch_size] output */
    unsigned int* __restrict__ rng_states,  /* [her_batch_size] per-thread LCG state */
    int L,                                  /* timesteps per episode */
    int capacity,                           /* N * L, total buffer size */
    int her_batch_size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= her_batch_size) return;

    int src = source_indices[idx];
    int ep  = episode_ids[src];

    /* Episode occupies [ep * L, (ep+1) * L) in the flat buffer */
    int ep_end = ep * L + L;
    if (ep_end > capacity) ep_end = capacity;

    /* Sample from [src+1, ep_end) */
    int range = ep_end - src - 1;
    if (range <= 0) {
        donor_indices[idx] = src; /* fallback: self (last step of episode) */
        return;
    }

    /* LCG — same pattern as action selection kernel */
    unsigned int rng = rng_states[idx];
    rng = rng * 1664525u + 1013904223u;
    rng_states[idx] = rng;

    /* Map [0, 2^24) uniformly onto [0, range) */
    int offset = (int)((float)(rng >> 8) / 16777216.0f * (float)range);
    if (offset >= range) offset = range - 1;

    donor_indices[idx] = src + 1 + offset;
}

/* ══════════════════════════════════════════════════════════════════════
 * KERNEL 3: her_find_episode_end
 *
 * HER Final strategy: for each source transition, find the last
 * transition index of the same episode.
 *
 * Episode ep occupies [ep*L, (ep+1)*L) so the last index is
 * (ep+1)*L - 1, clamped to capacity-1.
 *
 * Launch config: grid=(ceil(her_batch_size/256), 1, 1), block=(256, 1, 1).
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void her_find_episode_end(
    const int* __restrict__ episode_ids,    /* [capacity] */
    const int* __restrict__ source_indices, /* [her_batch_size] */
    int* __restrict__ end_indices,          /* [her_batch_size] output: last idx in episode */
    int L,                                  /* timesteps per episode */
    int capacity,                           /* N * L, total buffer size */
    int her_batch_size
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= her_batch_size) return;

    int src = source_indices[idx];
    int ep  = episode_ids[src];

    /* Last index of this episode */
    int ep_end = (ep + 1) * L - 1;
    if (ep_end >= capacity) ep_end = capacity - 1;

    end_indices[idx] = ep_end;
}
