// PER (Prioritized Experience Replay) GPU kernels — decoupled lookback scan.
//
// Architecture: flat priority array + single-pass O(N) prefix sum
// (Merrill & Garland, 2016) + binary search sampling.
//
// Four kernels:
//   1. per_update_pa   — priority update from BF16 TD errors
//   2. per_insert_pa   — insert raw priorities at max_priority
//   3. per_prefix_scan — decoupled lookback inclusive prefix sum
//   4. per_sample      — proportional sampling via binary search

// ---------------------------------------------------------------------------
// 1. per_update_pa: priority leaf update from BF16 TD errors
//    Reads bf16 TD errors directly (no separate cast kernel).
//    Computes p = |td|^alpha + eps, writes raw priority, computes pa = p^alpha,
//    writes to priorities_pa. atomicMax for batch max.
//    Launch: grid=(ceil(batch_size/256)), block=(256).
// ---------------------------------------------------------------------------
extern "C" __global__ void per_update_pa(
    float* __restrict__ priorities_pa,
    float* __restrict__ priorities,
    float* __restrict__ batch_max,
    const unsigned int* __restrict__ indices,
    const float* __restrict__ td_errors_bf16,
    float alpha, float epsilon,
    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 td = fabsf((td_errors_bf16[i]));
    float new_prio = powf(td, alpha) + epsilon;
    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;
}

// ---------------------------------------------------------------------------
// 2. per_insert_pa: insert priorities at max_priority
//    Takes raw priority values, computes pa = priority^alpha, writes to
//    priorities_pa.
//    Launch: grid=(ceil(batch_size/256)), block=(256).
// ---------------------------------------------------------------------------
extern "C" __global__ void per_insert_pa(
    float* __restrict__ priorities_pa,
    const unsigned int* __restrict__ indices,
    const float* __restrict__ raw_priorities,
    float alpha,
    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 pa = powf(raw_priorities[i], alpha);
    priorities_pa[idx] = pa;
}

// ---------------------------------------------------------------------------
// 3. per_prefix_scan: decoupled lookback inclusive prefix sum
//    (Merrill & Garland, 2016). Single-pass O(N) scan with dynamic tile
//    assignment via global counter — guarantees forward progress even when
//    num_tiles >> resident blocks (no spin-wait deadlock).
//
//    Launch: grid=(min(num_tiles, MAX_RESIDENT_BLOCKS)), block=(256),
//            shared_mem=256*sizeof(float).
//    The global counter `tile_counter[0]` must be zeroed before launch.
// ---------------------------------------------------------------------------
#define PER_TILE_SIZE 256

#define PER_STATE_INVALID          0u
#define PER_STATE_AGGREGATE        1u
#define PER_STATE_PREFIX_AVAILABLE 2u

extern "C" __global__ void per_prefix_scan(
    const float* __restrict__ input,
    float* __restrict__ output,
    volatile unsigned int* __restrict__ tile_state,
    volatile float* __restrict__ tile_aggregate,
    volatile float* __restrict__ tile_prefix,
    unsigned int* __restrict__ tile_counter,
    const int* __restrict__ size_ptr,
    int _num_tiles_unused)
{
    // Read size from pinned device-mapped pointer — current value at replay time
    int size = *size_ptr;
    int num_tiles = (size + PER_TILE_SIZE - 1) / PER_TILE_SIZE;
    extern __shared__ float sdata[];
    int tid = threadIdx.x;

    // Dynamic tile assignment: each block grabs the next tile via atomicAdd.
    // This ensures tiles are processed in order by co-resident blocks,
    // preventing deadlock when num_tiles > resident blocks.
    __shared__ int s_tile_id;

    while (true) {
        // Thread 0 grabs next tile
        if (tid == 0) {
            s_tile_id = (int)atomicAdd(tile_counter, 1u);
        }
        __syncthreads();

        int tile_id = s_tile_id;
        if (tile_id >= num_tiles) return;

        int gid = tile_id * PER_TILE_SIZE + tid;

        // Load into shared memory
        float val = (gid < size) ? input[gid] : 0.0f;
        sdata[tid] = val;
        __syncthreads();

        // Intra-tile inclusive scan (Hillis-Steele)
        for (int offset = 1; offset < PER_TILE_SIZE; offset <<= 1) {
            float add = (tid >= offset) ? sdata[tid - offset] : 0.0f;
            __syncthreads();
            sdata[tid] += add;
            __syncthreads();
        }

        float tile_total = sdata[PER_TILE_SIZE - 1];
        // Save local scan result before thread 0 overwrites sdata
        float my_scan = sdata[tid];

        // Decoupled lookback: thread 0 publishes state and computes prefix
        __shared__ float s_exclusive_prefix;

        if (tid == 0) {
            if (tile_id == 0) {
                tile_aggregate[0] = tile_total;
                tile_prefix[0] = tile_total;
                __threadfence();
                tile_state[0] = PER_STATE_PREFIX_AVAILABLE;
                s_exclusive_prefix = 0.0f;
            } else {
                // Publish aggregate
                tile_aggregate[tile_id] = tile_total;
                __threadfence();
                tile_state[tile_id] = PER_STATE_AGGREGATE;

                // Lookback through preceding tiles
                float excl = 0.0f;
                int look = tile_id - 1;

                while (look >= 0) {
                    unsigned int state;
                    do {
                        state = tile_state[look];
                    } while (state == PER_STATE_INVALID);

                    if (state == PER_STATE_PREFIX_AVAILABLE) {
                        excl += tile_prefix[look];
                        break;
                    } else {
                        excl += tile_aggregate[look];
                        look--;
                    }
                }

                // Publish inclusive prefix
                tile_prefix[tile_id] = excl + tile_total;
                __threadfence();
                tile_state[tile_id] = PER_STATE_PREFIX_AVAILABLE;
                s_exclusive_prefix = excl;
            }
        }
        __syncthreads();

        // Write output: exclusive prefix + local inclusive scan
        if (gid < size) {
            output[gid] = s_exclusive_prefix + my_scan;
        }
        __syncthreads(); // Ensure all threads done before grabbing next tile
    }
}

// ---------------------------------------------------------------------------
// Philox single-round helper (counter-based PRNG)
// ---------------------------------------------------------------------------
__device__ __forceinline__
unsigned int philox_single(unsigned int counter, unsigned int key) {
    unsigned int hi, lo;
    lo = counter * 0xD2511F53u;
    hi = __umulhi(counter, 0xD2511F53u);
    for (int r = 0; r < 10; r++) {
        unsigned int t = hi ^ key;
        hi = lo * 0xCD9E8D57u;
        lo = __umulhi(lo, 0xCD9E8D57u);
        lo ^= t;
        key += 0x9E3779B9u;
    }
    return lo ^ hi;
}

// ---------------------------------------------------------------------------
// 4. per_sample: deterministic stratified proportional sampling
//    Divides [0, total_sum) into batch_size equal segments and picks
//    the midpoint of each segment. Binary searches prefix_sum to find
//    the corresponding transition. Zero randomness — fully deterministic.
//    Launch: grid=(ceil(batch_size/256)), block=(256).
// ---------------------------------------------------------------------------
extern "C" __global__ void per_sample(
    const float* __restrict__ prefix_sum,
    const float* __restrict__ priorities_pa,
    long long* __restrict__ out_indices,
    float* __restrict__ out_priorities,
    float* __restrict__ total_sum_out,
    const int* __restrict__ rng_step_ptr,
    const int* __restrict__ size_ptr,
    int batch_size)
{
    // Read from pinned device-mapped pointers — current values at graph replay
    int size = *size_ptr;
    unsigned int seed = (unsigned int)(*rng_step_ptr);
    int i = blockIdx.x * blockDim.x + threadIdx.x;

    // Thread 0 writes total sum
    if (i == 0) {
        total_sum_out[0] = (size > 0) ? prefix_sum[size - 1] : 0.0f;
    }

    if (i >= batch_size) return;

    float total_sum = (size > 0) ? prefix_sum[size - 1] : 1.0f;

    // Deterministic stratified sampling: midpoint of segment i
    // segment_i = [i * total_sum / B, (i+1) * total_sum / B)
    // threshold = midpoint = (i + 0.5) * total_sum / B
    float segment_width = total_sum / (float)batch_size;
    float threshold = ((float)i + 0.5f) * segment_width;

    // Binary search: find smallest index where prefix_sum[index] > threshold
    int lo = 0;
    int hi = size - 1;
    while (lo < hi) {
        int mid = lo + (hi - lo) / 2;
        if (prefix_sum[mid] <= threshold) {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }

    // Clamp to valid range
    if (lo >= size) lo = size - 1;
    if (lo < 0) lo = 0;

    out_indices[i] = (long long)lo;
    out_priorities[i] = priorities_pa[lo];
}
