// 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)
//   3. scatter_insert_bf16  — ring buffer insert for bf16 state matrices
//   4. gather_f32           — index_select for f32 arrays
//   5. gather_u32           — index_select for u32 arrays
//   6. gather_bf16_rows     — index_select rows from [cap, dim] bf16 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];
}

// ── 3. Scatter insert for bf16 rows [batch, dim] → ring buffer [cap, dim] ──

extern "C" __global__
void scatter_insert_bf16(
    unsigned short* __restrict__ dst,     // [capacity * state_dim]
    const unsigned short* __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 i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    int idx = (int)indices[i];
    out[i] = src[idx];
}

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

extern "C" __global__
void gather_bf16_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];
}

// ── 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,
    int n_buffer,                           // number of experiences in buffer
    int batch_size
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;
    float ts = total_sum_buf[0];
    float prob = (sampled_prios[i] * (float)n_buffer) / ts;
    weights[i] = powf(prob, neg_beta);
}

// 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 BF16 for ring buffer storage ─────────────────────

extern "C" __global__
void f32_to_bf16_cast(
    unsigned short* __restrict__ out,
    const float* __restrict__ input,
    int n
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= n) return;
    // Truncate F32 to BF16: shift right by 16 bits
    unsigned int bits = __float_as_uint(input[i]);
    // Round to nearest even (RNE)
    unsigned int lsb = (bits >> 16) & 1;
    unsigned int rounding_bias = 0x7FFF + lsb;
    bits += rounding_bias;
    out[i] = (unsigned short)(bits >> 16);
}

// ── 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 BF16 (u16) to F32 on GPU ───────────────────────────────────────
//    Inverse of f32_to_bf16_cast. Converts bf16 storage (u16 bits) to f32.

extern "C" __global__
void bf16_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;
    // BF16 is the 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;
}

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

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