// CUDA kernels for GPU-resident replay buffer (pure cudarc, no Candle).
//
// Kernels:
//   1. scatter_insert_f32   — ring buffer insert for f32 arrays (rewards, dones, priorities)
//   2. scatter_insert_u32   — ring buffer insert for u32 arrays (actions)
//   2b. scatter_insert_i32  — ring buffer insert for i32 arrays (SP13 aux sign labels with -1 sentinel)
//   3. scatter_insert       — ring buffer insert for state matrices
//   4. gather_f32           — index_select for f32 arrays
//   5. gather_u32           — index_select for u32 arrays
//   6. gather_rows          — index_select rows from [cap, dim] matrix
//   7. pow_alpha_f32        — prios[i] = pow(abs(input[i]), alpha)
//   8. is_weights_f32       — importance sampling weight computation
//   9. priority_update_f32  — |td|^alpha + eps scatter-write + atomicMax
//  10. fill_f32             — fill array with scalar value
//  11. rand_scale_f32       — scale uniform [0,1) by total_sum

// ── 1. Scatter insert for f32 (1D) ─────────────────────────────────────────

extern "C" __global__
void scatter_insert_f32(
    float* __restrict__ dst,          // [capacity]
    const float* __restrict__ src,    // [batch_size]
    int start_idx,                     // write cursor position
    int capacity,                      // ring buffer capacity
    int batch_size                     // number of elements to write
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    int dst_idx = (start_idx + i) % capacity;
    dst[dst_idx] = src[i];
}

// ── 2. Scatter insert for u32 (1D) ─────────────────────────────────────────

extern "C" __global__
void scatter_insert_u32(
    unsigned int* __restrict__ dst,
    const unsigned int* __restrict__ src,
    int start_idx,
    int capacity,
    int batch_size
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    int dst_idx = (start_idx + i) % capacity;
    dst[dst_idx] = src[i];
}

// ── 2b. Scatter insert for i32 (1D, signed) ────────────────────────────────
//    Mirror of scatter_insert_u32 with i32 storage. SP13 Layer B uses this
//    for the 30-bar aux sign label where -1 is a "skip" sentinel that must
//    survive the ring-buffer round-trip as a signed value (gather_i32_scalar
//    at #18c is its consumer).

extern "C" __global__
void scatter_insert_i32(
    int* __restrict__ dst,
    const int* __restrict__ src,
    int start_idx,
    int capacity,
    int batch_size
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    int dst_idx = (start_idx + i) % capacity;
    dst[dst_idx] = src[i];
}

// ── 3. Scatter insert for f32 rows [batch, dim] → ring buffer [cap, dim] ──
//    2D-aware version of scatter_insert_f32 for state matrices.
//    Each thread writes one element: dst[(start + row) % cap, col] = src[row, col].

extern "C" __global__
void scatter_insert_f32_rows(
    float* __restrict__ dst,          // [capacity * state_dim]
    const float* __restrict__ src,    // [batch_size * state_dim]
    int start_idx,
    int capacity,
    int state_dim,
    int batch_size
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * state_dim;
    if (tid >= total) return;
    int row = tid / state_dim;
    int col = tid % state_dim;
    int dst_row = (start_idx + row) % capacity;
    dst[dst_row * state_dim + col] = src[row * state_dim + col];
}

// ── 4. Gather f32 by indices ────────────────────────────────────────────────

extern "C" __global__
void gather_f32(
    float* __restrict__ out,            // [batch_size]
    const float* __restrict__ src,      // [capacity]
    const long long* __restrict__ indices, // [batch_size]
    int batch_size
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    int idx = (int)indices[i];
    out[i] = src[idx];
}

// ── 5. Gather u32 by indices ────────────────────────────────────────────────

extern "C" __global__
void gather_u32(
    unsigned int* __restrict__ out,
    const unsigned int* __restrict__ src,
    const long long* __restrict__ indices,
    int batch_size,
    int capacity
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    int idx = (int)indices[i];
    if (idx < 0 || idx >= capacity) return;
    out[i] = src[idx];
}

// ── 6. Gather rows [cap, dim] → [batch, dim] by indices ───────────────

extern "C" __global__
void gather_rows(
    unsigned short* __restrict__ out,        // [batch_size * state_dim]
    const unsigned short* __restrict__ src,  // [capacity * state_dim]
    const long long* __restrict__ indices,   // [batch_size]
    int state_dim,
    int batch_size
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * state_dim;
    if (tid >= total) return;
    int row = tid / state_dim;
    int col = tid % state_dim;
    int src_row = (int)indices[row];
    out[row * state_dim + col] = src[src_row * state_dim + col];
}

// ── 6a. Gather f32 rows [cap, dim] → [batch, dim] by indices (#30) ─────────

extern "C" __global__
void gather_f32_rows(
    float* __restrict__ out,         // [batch_size * state_dim]
    const float* __restrict__ src,   // [capacity * state_dim]
    const long long* __restrict__ indices,  // [batch_size]
    int state_dim,
    int batch_size
) {
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total = batch_size * state_dim;
    if (tid >= total) return;
    int row = tid / state_dim;
    int col = tid % state_dim;
    int src_row = (int)indices[row];
    out[row * state_dim + col] = src[src_row * state_dim + col];
}

// ── 6b. Gather scalars (1D) by indices ────────────────────────────────

extern "C" __global__
void gather(
    unsigned short* __restrict__ out,        // [batch_size]
    const unsigned short* __restrict__ src,  // [capacity]
    const long long* __restrict__ indices,   // [batch_size]
    int batch_size,
    int capacity                             // bounds check for index validity
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    int idx = (int)indices[i];
    /* Bounds check: corrupt PER segment tree can produce OOB indices.
     * Clamp to [0, capacity-1] instead of reading garbage (NaN bits). */
    if (idx < 0 || idx >= capacity) return;
    out[i] = src[idx];
}

// ── 7. prios_alpha: out[i] = pow(src[i], alpha) ────────────────────────────

extern "C" __global__
void pow_alpha_f32(
    float* __restrict__ out,
    const float* __restrict__ src,
    float alpha,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = powf(src[i], alpha);
}

// ── 8. IS weights: w_i = (N * sampled_prio[i] / total_sum)^(-beta) ─────────
//    Then normalize by max weight.
//    total_sum_buf is a GPU-resident scalar pointer (zero CPU readback).

extern "C" __global__
void is_weights_f32(
    float* __restrict__ weights,           // [batch_size] output (in-place)
    const float* __restrict__ sampled_prios, // [batch_size]
    const float* __restrict__ total_sum_buf, // [1] GPU-resident scalar
    float neg_beta,
    const int* __restrict__ n_buffer_ptr,   // pinned device-mapped — current buffer size at replay
    int batch_size
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    int n_buffer = *n_buffer_ptr;
    float ts = total_sum_buf[0];
    float prob = fmaxf((sampled_prios[i] * (float)n_buffer) / fmaxf(ts, 1e-8f), 1e-12f);
    float w = powf(prob, neg_beta);
    weights[i] = fminf(w, 60000.0f);
}

// Normalize weights by max (two-pass: first find max, then divide)
extern "C" __global__
void normalize_weights_f32(
    float* __restrict__ weights,
    const float* __restrict__ max_weight,  // [1] scalar
    int batch_size
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    float mw = max_weight[0];
    if (mw < 1e-8f) mw = 1e-8f;
    weights[i] /= mw;
}

// ── 9. Priority update: |td|^alpha + eps, scatter, atomicMax ────────────────

extern "C" __global__
void priority_update_f32(
    const float* __restrict__ td_errors,
    const unsigned int* __restrict__ indices,
    float* __restrict__ priorities,
    float* __restrict__ batch_max,
    float alpha,
    float epsilon,
    int batch_size
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    float td = fabsf(td_errors[i]);
    float new_prio = powf(td, alpha) + epsilon;

    // Clamp to [epsilon, 1e6]
    if (new_prio < epsilon) new_prio = epsilon;
    if (new_prio > 1e6f) new_prio = 1e6f;

    unsigned int idx = indices[i];
    priorities[idx] = new_prio;

    // atomicMax for batch max: IEEE 754 positive floats preserve int ordering
    int ival = __float_as_int(new_prio);
    atomicMax((int*)batch_max, ival);
}

// ── 10. Fill f32 array with scalar ──────────────────────────────────────────

extern "C" __global__
void fill_f32(
    float* __restrict__ out,
    float value,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = value;
}

// ── 11. Scale uniform [0,1) by total_sum scalar ────────────────────────────

extern "C" __global__
void rand_scale_f32(
    float* __restrict__ targets,        // [batch_size] in-place: targets[i] *= total_sum
    const float* __restrict__ total_sum, // [1]
    int batch_size
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    targets[i] *= total_sum[0];
}

// ── 12. Reduce max f32 ─────────────────────────────────────────────────────
//    Single-block parallel reduction to find max of array.

extern "C" __global__
void reduce_max_f32(
    const float* __restrict__ input,
    float* __restrict__ output,    // [1]
    int n
) {
    extern __shared__ float sdata[];

    int tid = threadIdx.x;
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    sdata[tid] = (i < n) ? input[i] : -1e30f;
    __syncthreads();

    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s && sdata[tid] < sdata[tid + s]) {
            sdata[tid] = sdata[tid + s];
        }
        __syncthreads();
    }

    if (tid == 0) {
        // atomicMax across blocks (for multi-block launches)
        int ival = __float_as_int(sdata[0]);
        atomicMax((int*)output, ival);
    }
}

// ── 13. Convert i64 indices to u32 ──────────────────────────────────────────

extern "C" __global__
void i64_to_u32(
    unsigned int* __restrict__ out,
    const long long* __restrict__ input,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = (unsigned int)input[i];
}

// ── 14. Cast F32 states to U16 for ring buffer storage ─────────────────────

extern "C" __global__
void f32_to_u16_cast(
    unsigned short* __restrict__ out,
    const float* __restrict__ input,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    /* Reinterpret f32 bit pattern to u16. */
    float val = (input[i]);
    out[i] = *(unsigned short*)&val;
}

// ── 14b. Fill f32 array from a GPU scalar pointer ───────────────────────────

extern "C" __global__
void fill_from_gpu_f32(
    float* __restrict__ out,
    const float* __restrict__ value_ptr,  // [1] scalar on GPU
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = value_ptr[0];
}

// ── 15. Cast U16 to F32 on GPU ───────────────────────────────────────
//    Inverse of f32_to_u16_cast. Converts u16 storage bits to f32.

extern "C" __global__
void u16_to_f32_cast(
    float* __restrict__ out,
    const unsigned short* __restrict__ input,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    // Upper 16 bits of F32, so shift left by 16
    unsigned int bits = ((unsigned int)input[i]) << 16;
    out[i] = __uint_as_float(bits);
}

// ── 16. Cast u32 to F32 on GPU ──────────────────────────────────────────────

extern "C" __global__
void u32_to_f32_cast(
    float* __restrict__ out,
    const unsigned int* __restrict__ input,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = (float)input[i];
}

// ── 17. GPU max of two single-element f32 scalars ───────────────────────────
//    out[0] = max(a[0], b[0])

extern "C" __global__
void max_of_two_f32(
    float* __restrict__ out,
    const float* __restrict__ a,
    const float* __restrict__ b
) {
    float va = a[0];
    float vb = b[0];
    out[0] = (va > vb) ? va : vb;
}

// ── 18a. Gather f32 rows with zero-padding to dst_stride ───────────────────
//    Replaces 2-step gather + pad_states: gathers rows from [cap, src_dim]
//    into [batch, dst_stride], zero-filling columns [src_dim, dst_stride).

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

// ── 18b. Gather scalar f32 by i64 indices ─────────────────────────────────
//    Direct gather for rewards, dones, IS-weights into trainer buffers.

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

// ── 18c. Gather scalar i32 by i64 indices ─────────────────────────────────
//    Direct gather for actions (u32/i32) into trainer buffers.

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

// ── 18. NaN/Inf check kernel ────────────────────────────────────────────────
//    out[i] = (isnan(v) || isinf(v)) ? 1.0f : 0.0f

extern "C" __global__
void nan_inf_check_f32(
    float* __restrict__ out,
    const float* __restrict__ input,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    float v = input[i];
    out[i] = (isnan(v) || isinf(v)) ? 1.0f : 0.0f;
}

// ── 19. Dead neuron (near-zero) check kernel ────────────────────────────────
//    out[i] = (fabsf(v) < threshold) ? 1.0f : 0.0f

extern "C" __global__
void dead_neuron_check_f32(
    float* __restrict__ out,
    const float* __restrict__ input,
    float threshold,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = (fabsf(input[i]) < threshold) ? 1.0f : 0.0f;
}

// ── 20. Cast f32-encoded indices to u32 ─────────────────────────────────────
//    out[i] = (unsigned int)input[i]

extern "C" __global__
void f32_idx_to_u32(
    unsigned int* __restrict__ out,
    const float* __restrict__ input,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = (unsigned int)input[i];
}

/* Graph-safe GPU reduction: scaled-mean of u32 actions → i32[1].
 *
 * Writes (sum(actions) * 1000) / n  into mean_scaled_out[0]. Launched as a
 * single 256-thread block; deterministic tree reduction in shared memory.
 * Replaces the prior dtoh-then-sum-on-host path, which was not graph-safe:
 * memcpy_dtoh inside a CUDA Graph capture region invalidates the stream
 * (L40S smoke surfaced this on Fold 1 re-capture). Keep n small — called
 * with batch-sized windows (≤ max_batch_size), one block handles the lot.
 */
extern "C" __global__ void actions_sum_scale_reduce_u32(
    const unsigned int* __restrict__ actions,
    int* __restrict__ mean_scaled_out,   /* i32[1]: (sum * 1000) / n */
    int n
) {
    extern __shared__ long long sdata_ll[];
    int tid = threadIdx.x;

    long long local = 0;
    for (int i = tid; i < n; i += blockDim.x) {
        local += (long long)actions[i];
    }
    sdata_ll[tid] = local;
    __syncthreads();

    /* Deterministic tree reduction — no atomics, no warp shuffles with
     * data-dependent lane ordering. */
    for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
        if (tid < offset) {
            sdata_ll[tid] += sdata_ll[tid + offset];
        }
        __syncthreads();
    }

    if (tid == 0) {
        long long denom = n > 0 ? (long long)n : 1LL;
        long long mean_scaled = (sdata_ll[0] * 1000LL) / denom;
        mean_scaled_out[0] = (int)mean_scaled;
    }
}

/* C1/P1: Health-weighted PER priority — boosts priorities of diverse-action
 * experiences during collapse (health < 0.8).
 *
 *   new_prio = clamp((|td_error[i]| + epsilon)^alpha
 *                    × (1 + 2×(1−health) × |action[i] − mean_action|),
 *                    epsilon, 1e6)
 *   priorities[indices[i]]    = new_prio
 *   priorities_pa[indices[i]] = new_prio^alpha
 *   atomicMax(batch_max, new_prio)  (IEEE 754 int trick)
 *
 * `mean_action_scaled_ptr` points to a device-side i32 set by the companion
 * kernel `actions_sum_scale_reduce_u32` (same batch, scaled mean × 1000).
 * Device-pointer passthrough keeps the update graph-safe — no host readback.
 *
 * Drop-in replacement for per_update_pa when health < 0.8.
 */
extern "C" __global__ void pow_alpha_diverse_f32(
    float* __restrict__ priorities_pa,
    float* __restrict__ priorities,
    float* __restrict__ batch_max,
    const unsigned int* __restrict__ indices,
    const float* __restrict__ td_errors,
    const int* __restrict__ actions,
    const int* __restrict__ mean_action_scaled_ptr,   /* device i32[1] */
    float alpha,
    float epsilon,
    float health,
    int capacity,
    int batch_size
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    unsigned int idx = indices[i];
    if (idx >= (unsigned int)capacity) return;

    float base = powf(fabsf(td_errors[i]) + epsilon, alpha);

    /* Plain load (not __ldg): the backing buffer is device-mapped pinned host
     * memory (cuMemHostAlloc + cuMemHostGetDevicePointer). __ldg routes through
     * the read-only cache which is not guaranteed coherent for mapped host
     * memory — stale reads manifested as a 13× drop in final q_gap on the
     * multi-trial smoke. Plain load respects stream serialization between
     * actions_sum_scale_reduce_u32 (writer) and this kernel (reader). */
    float mean_action = (float)(mean_action_scaled_ptr[0]) * 0.001f;
    float action_diff = fabsf((float)actions[i] - mean_action);
    float diversity_mult = 1.0f + 2.0f * (1.0f - health) * action_diff;

    float new_prio = base * diversity_mult;
    if (new_prio < epsilon) new_prio = epsilon;
    if (new_prio > 1e6f) new_prio = 1e6f;

    priorities[idx] = new_prio;

    int ival = __float_as_int(new_prio);
    atomicMax((int*)batch_max, ival);

    float pa = powf(new_prio, alpha);
    priorities_pa[idx] = pa;
}

// ── 21. Regime weight combination kernel ────────────────────────────────────
//    out[i] = trending[i] * t_scale + ranging[i] * r_scale + volatile[i] * v_scale

extern "C" __global__
void regime_weight_combine_f32(
    float* __restrict__ out,
    const float* __restrict__ trending,
    const float* __restrict__ ranging,
    const float* __restrict__ volatile_mask,
    float t_scale,
    float r_scale,
    float v_scale,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    out[i] = trending[i] * t_scale + ranging[i] * r_scale + volatile_mask[i] * v_scale;
}
