/**
 * Graph utility kernels — GPU-native gather + step counter increment.
 *
 * These kernels replace host-side operations on the training hot path,
 * enabling a true single-graph architecture with zero CPU involvement
 * per step (single cuGraphExecLaunch).
 *
 * Kernels:
 *   gather_f32_rows_padded   — gather rows from replay buffer + zero-pad
 *                              to 128-byte-aligned dst_stride (replaces
 *                              2-step gather + pad_states host pipeline)
 *   gather_f32_scalar        — gather scalar f32 values (rewards, dones,
 *                              IS-weights) from replay buffer indices
 *   gather_i32_scalar        — gather scalar i32 values (actions) from
 *                              replay buffer indices
 *   increment_step_counters  — single-thread kernel that increments all
 *                              per-step Adam/scheduler counters on GPU
 *                              and computes cosine-annealed tau
 *
 * No #include of common_device_functions.cuh is needed here — the build
 * function in build.rs prepends it via string concatenation before
 * passing the combined source to nvcc.
 */

/* ══════════════════════════════════════════════════════════════════════
 * gather_f32_rows_padded
 *
 * Gathers rows from a flat row-major src buffer into dst, zero-padding
 * columns [src_dim, dst_stride) so the output is 128-byte aligned per row.
 *
 * Launch config: grid=(ceil(batch_size * dst_stride / 256)), block=(256)
 *
 * Parameters:
 *   dst        — output buffer [batch_size * dst_stride]
 *   src        — source buffer [capacity * src_dim]
 *   indices    — i64 indices into src rows [batch_size]
 *   src_dim    — number of valid f32 elements per src row (unpadded)
 *   dst_stride — padded row width in f32 elements (dst_dim >= src_dim,
 *                aligned to 128 bytes = 32 f32 elements)
 *   batch_size — number of rows to gather
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void gather_f32_rows_padded(
    float* __restrict__ dst,
    const float* __restrict__ src,
    const long long* __restrict__ indices,
    int src_dim, int dst_stride, int batch_size)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * dst_stride;
    if (tid >= total) return;
    int row = tid / dst_stride;
    int col = tid % dst_stride;
    long long src_row = indices[row];
    dst[tid] = (col < src_dim) ? src[src_row * src_dim + col] : 0.0f;
}

/* ══════════════════════════════════════════════════════════════════════
 * gather_f32_scalar
 *
 * Gathers scalar f32 values (rewards, dones, IS-weights) from the replay
 * buffer into a flat output array using the sampled indices.
 *
 * Launch config: grid=(ceil(batch_size / 256)), block=(256)
 *
 * Parameters:
 *   dst        — output buffer [batch_size]
 *   src        — source buffer [capacity] (replay buffer scalar field)
 *   indices    — i64 indices [batch_size]
 *   batch_size — number of elements to gather
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void gather_f32_scalar(
    float* __restrict__ dst,
    const float* __restrict__ src,
    const long long* __restrict__ indices,
    int batch_size)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= batch_size) return;
    dst[tid] = src[indices[tid]];
}

/* ══════════════════════════════════════════════════════════════════════
 * gather_i32_scalar
 *
 * Gathers scalar i32 values (packed actions) from the replay buffer into
 * a flat output array using the sampled indices.
 *
 * Launch config: grid=(ceil(batch_size / 256)), block=(256)
 *
 * Parameters:
 *   dst        — output buffer [batch_size]
 *   src        — source buffer [capacity] (replay buffer action field)
 *   indices    — i64 indices [batch_size]
 *   batch_size — number of elements to gather
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void gather_i32_scalar(
    int* __restrict__ dst,
    const int* __restrict__ src,
    const long long* __restrict__ indices,
    int batch_size)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= batch_size) return;
    dst[tid] = src[indices[tid]];
}

/* ══════════════════════════════════════════════════════════════════════
 * increment_step_counters
 *
 * Single-thread kernel that atomically increments all per-step counters
 * and recomputes the cosine-annealed IQN tau value — all on GPU, zero
 * CPU synchronisation required.
 *
 * All counter pointers must be device-accessible (device alloc or pinned
 * host-mapped memory).  tau_out is written on every call.
 *
 * Launch config: grid=(1), block=(1)
 *
 * Parameters:
 *   adam_t, sel_t, denoise_t, iql_t, iql_low_t, iqn_t, attn_t, rng_step
 *              — per-scheduler step counters (int*, device-mapped)
 *   tau_out    — cosine-annealed tau written here (float*, device-mapped)
 *   tau_init   — initial tau value (start of annealing)
 *   tau_final  — final tau value (end of annealing)
 *   anneal_steps — total number of steps over which to anneal
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void increment_step_counters(
    int* adam_t, int* sel_t, int* denoise_t, int* iql_t,
    int* iql_low_t, int* iqn_t, int* attn_t, int* rng_step,
    float* tau_out, float tau_init, float tau_final, int anneal_steps)
{
    /* adam_t is the master step — its post-increment value drives the
     * cosine schedule so that step 1 produces the first non-initial tau. */
    // Single thread — plain writes, no atomicAdd needed
    int step = (*adam_t += 1);
    *sel_t     += 1;
    *denoise_t += 1;
    *iql_t     += 1;
    *iql_low_t += 1;
    if (iqn_t)    *iqn_t     += 1;
    if (attn_t)   *attn_t    += 1;
    if (rng_step) *rng_step  += 1;

    float progress = fminf((float)step / (float)anneal_steps, 1.0f);
    float cosine   = 0.5f * (1.0f + cosf(progress * 3.14159265f));
    *tau_out = tau_final + (tau_init - tau_final) * cosine;
}

/* ══════════════════════════════════════════════════════════════════════
 * copy_f32 — Graph-safe device-to-device copy.
 *
 * Replaces memcpy_dtod_async which is NOT captured by CUDA Graph.
 * Any DtoD copy inside a graph-captured code path must use this kernel
 * instead of cuMemcpyDtoDAsync.
 *
 * Launch config: grid=(ceil(N / 256)), block=(256)
 *
 * Parameters:
 *   src — source buffer [N] (f32, or any 4-byte type)
 *   dst — destination buffer [N] (f32, or any 4-byte type)
 *   N   — number of 4-byte elements to copy
 * ══════════════════════════════════════════════════════════════════════ */
extern "C" __global__ void copy_f32(
    const float* __restrict__ src,
    float* __restrict__ dst,
    int N
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < N) dst[i] = src[i];
}
