// GPU segment tree kernels for O(log n) PER (Prioritized Experience Replay).
//
// Replaces the O(n) prefix sum + searchsorted pipeline with a binary segment
// tree stored as a flat array: tree[0] unused, root at tree[1], leaves at
// tree[capacity..2*capacity]. Internal nodes store sums of children.
//
// Three kernels:
//   1. seg_tree_update    — priority update: compute (|td|^alpha + eps),
//                           write to priorities[], compute priority^alpha,
//                           write to tree leaf, propagate sums to root.
//   2. seg_tree_insert    — insert: take raw priorities (max_priority fill),
//                           compute priority^alpha, write to tree leaf,
//                           propagate sums to root.
//   3. seg_tree_sample    — proportional sampling: parallel root-to-leaf
//                           traversal with Philox RNG. Output i64 indices.

// ---------------------------------------------------------------------------
// 1. seg_tree_update: priority update + tree propagation
//    Called from update_priorities_gpu(). One thread per updated index.
//    Replaces: priority_update_f32 + pow_alpha + prefix_sum rebuild.
// ---------------------------------------------------------------------------
extern "C" __global__ void seg_tree_update(
    float* __restrict__ tree,               // [2 * capacity] segment tree
    float* __restrict__ priorities,          // [capacity] raw priority buffer
    float* __restrict__ batch_max,           // [1] atomicMax accumulator
    const unsigned int* __restrict__ indices, // [batch_size] buffer indices
    const float* __restrict__ td_errors,     // [batch_size] TD errors
    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;

    // Compute priority = |td|^alpha + eps (same as priority_update_f32)
    float td = fabsf(td_errors[i]);
    float new_prio = powf(td, alpha) + epsilon;
    if (new_prio < epsilon) new_prio = epsilon;
    if (new_prio > 1e6f) new_prio = 1e6f;

    // Write to priorities buffer (for max_priority tracking)
    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);

    // Compute priority^alpha for tree leaf (same as pow_alpha_f32)
    float pa = powf(new_prio, alpha);

    // Write leaf and propagate sums to root
    int leaf = capacity + (int)idx;
    tree[leaf] = pa;

    int node = leaf >> 1;
    while (node >= 1) {
        tree[node] = tree[2 * node] + tree[2 * node + 1];
        node >>= 1;
    }
}

// ---------------------------------------------------------------------------
// 2. seg_tree_insert: insert raw priorities into tree
//    Called from insert_batch(). Takes already-computed priority values
//    (e.g., max_priority fill), computes priority^alpha, writes to tree.
// ---------------------------------------------------------------------------
extern "C" __global__ void seg_tree_insert(
    float* __restrict__ tree,               // [2 * capacity] segment tree
    const unsigned int* __restrict__ indices, // [batch_size] buffer indices
    const float* __restrict__ priorities,    // [batch_size] raw priority values
    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;

    // Compute priority^alpha for tree leaf
    float pa = powf(priorities[i], alpha);

    // Write leaf and propagate sums to root
    int leaf = capacity + (int)idx;
    tree[leaf] = pa;

    int node = leaf >> 1;
    while (node >= 1) {
        tree[node] = tree[2 * node] + tree[2 * node + 1];
        node >>= 1;
    }
}

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

// ---------------------------------------------------------------------------
// 3. seg_tree_sample: proportional PER sampling via tree traversal
//    Each thread generates a Philox random threshold in [0, tree[1]) and
//    traverses root to leaf. Outputs i64 indices for gather kernel compat.
//    Replaces: pow_alpha + prefix_sum + threshold_gen + searchsorted.
// ---------------------------------------------------------------------------
extern "C" __global__ void seg_tree_sample(
    const float* __restrict__ tree,          // [2 * capacity] segment tree
    long long* __restrict__ out_indices,      // [batch_size] output indices (i64)
    unsigned int seed,
    int capacity, int batch_size,
    int buffer_size)                          // actual number of valid entries
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    // Philox RNG: generate uniform in [0, total_sum)
    unsigned int bits = philox_single((unsigned int)i, seed);
    float u = (float)(bits >> 8) * (1.0f / 16777216.0f);
    float total_sum = tree[1];  // root = total priority sum
    float threshold = u * total_sum;

    // Tree traversal: root -> leaf
    int node = 1;
    while (node < capacity) {
        float left_sum = tree[2 * node];
        if (threshold <= left_sum) {
            node = 2 * node;
        } else {
            threshold -= left_sum;
            node = 2 * node + 1;
        }
    }

    // node is now a leaf: index = node - capacity
    int idx = node - capacity;
    // Clamp to valid range [0, buffer_size - 1]
    if (idx >= buffer_size) idx = buffer_size - 1;
    if (idx < 0) idx = 0;

    out_indices[i] = (long long)idx;
}

// ---------------------------------------------------------------------------
// 4. seg_tree_gather_prios: gather priority^alpha values from tree leaves
//    for IS weight computation. Reads tree[capacity + idx] for each sampled
//    index. Replaces: gather from pa_buf.
// ---------------------------------------------------------------------------
extern "C" __global__ void seg_tree_gather_prios(
    float* __restrict__ out_prios,           // [batch_size] output sampled priorities
    const float* __restrict__ tree,          // [2 * capacity] segment tree
    const long long* __restrict__ indices,   // [batch_size] sampled indices
    int capacity, int batch_size)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    int idx = (int)indices[i];
    out_prios[i] = tree[capacity + idx];
}
