/* q_quantile_reduce — GPU-driven per-branch P5/P95 EMA update (Plan 2 Task 1, spec §4.C.1).
 *
 * For each of the 4 action branches (direction=0, magnitude=1, order=2, urgency=3),
 * computes the per-sample max-Q-over-actions, sorts those N values in shared memory
 * via bitonic sort or quickselect, picks the 5th and 95th percentile values, and
 * applies an EMA to ISV[Q_P05_*] and ISV[Q_P95_*].
 *
 * Cold-path per-epoch kernel. Launched with grid=(4,1,1), block=(1,1,1) — single
 * thread per block, one block per branch.  N samples can reach up to the full batch
 * size; for N <= MAX_SHMEM_SORT_SIZE the sort runs in shared memory; larger N is
 * truncated to a representative subsample (still epoch-cold-path accurate).
 *
 * No atomicAdd anywhere — per-sample array + single-thread reduce only.
 * (feedback_no_atomicadd.md).
 *
 * ISV slot mapping (Plan 2 Task 1 + Plan 2 Task 3 D.2 shifted layout):
 *   Q_P05_DIR=50, Q_P05_MAG=51, Q_P05_ORD=52, Q_P05_URG=53
 *   Q_P95_DIR=54, Q_P95_MAG=55, Q_P95_ORD=56, Q_P95_URG=57
 * Shifted +3 from original layout by D.2 per-branch gamma insertion at [43..47).
 */

/* Maximum samples that fit in shared memory for the sort path.
 * Two arrays of this size × sizeof(float) must fit within 48KB (0xc000) shmem.
 * 3072 × 4 × 2 = 24576 bytes — fits with plenty of margin.
 * For N > 3072, the first 3072 samples form a representative subsample. */
#define MAX_SHMEM_SORT_SIZE 3072

/* ── Bitonic sort — in-place ascending sort on arr[0..len).
 * len must be a power of 2. Single-thread; safe on cold path. ──────────── */
__device__ static void bitonic_sort_block(float* arr, int len)
{
    for (int k = 2; k <= len; k <<= 1) {
        for (int j = k >> 1; j >= 1; j >>= 1) {
            for (int i = 0; i < len; i++) {
                int ixj = i ^ j;
                if (ixj > i) {
                    bool ascending = ((i & k) == 0);
                    if (( ascending && arr[i] > arr[ixj]) ||
                        (!ascending && arr[i] < arr[ixj])) {
                        float tmp  = arr[i];
                        arr[i]     = arr[ixj];
                        arr[ixj]   = tmp;
                    }
                }
            }
        }
    }
}

/* ── Quickselect: find the k-th smallest value (0-indexed) in arr[0..n).
 * Modifies arr in place. Single-thread; safe on cold path. ─────────────── */
__device__ static float nth_element_float(float* arr, int n, int k)
{
    int lo = 0, hi = n - 1;
    while (lo < hi) {
        float pivot = arr[(lo + hi) >> 1];
        int l = lo, r = hi;
        while (l <= r) {
            while (arr[l] < pivot) l++;
            while (arr[r] > pivot) r--;
            if (l <= r) {
                float tmp = arr[l];
                arr[l]    = arr[r];
                arr[r]    = tmp;
                l++; r--;
            }
        }
        if (k <= r)      hi = r;
        else if (k >= l) lo = l;
        else             break;
    }
    return arr[k];
}

/* ── Main kernel ─────────────────────────────────────────────────────────── */
extern "C" __global__ void q_quantile_reduce(
    const float* __restrict__ q_values,          /* [n_samples, total_actions] row-major */
    int          n_samples,
    int          total_actions,
    const int*   __restrict__ branch_offsets,    /* [4] — action index of each branch start */
    const int*   __restrict__ branch_sizes,      /* [4] — number of actions per branch */
    const float* __restrict__ isv,               /* read: current Q_P05/Q_P95 EMA seeds */
    float*       __restrict__ isv_out,           /* write: updated Q_P05/Q_P95 EMA values */
    float        ema_alpha
)
{
    /* Single-thread per block; blockIdx.x selects the branch. */
    if (threadIdx.x != 0) return;

    const int branch = (int)blockIdx.x;           /* 0=dir, 1=mag, 2=ord, 3=urg */

    /* ISV slot indices for this branch (Plan 2 Task 1 + D.2 shifted layout).
     * Q_P05 base = 50, Q_P95 base = 54. Shifted +3 by D.2 per-branch gamma. */
    const int p05_slot = 50 + branch;
    const int p95_slot = 54 + branch;

    const int b_off  = branch_offsets[branch];
    const int b_size = branch_sizes[branch];
    const int n      = n_samples;

    /* Guard: degenerate configuration — skip without modifying ISV. */
    if (b_size <= 0 || b_off < 0 || n <= 0 ||
        total_actions <= 0 || b_off + b_size > total_actions) {
        return;
    }

    /* ── Phase 1: collect per-sample max-Q over this branch's actions.
     * Stored in svals[0..fill_n).  fill_n = min(n, MAX_SHMEM_SORT_SIZE).
     * When n > MAX_SHMEM_SORT_SIZE the first fill_n samples are used; this is a
     * representative subsample for the epoch-boundary percentile estimate.     */
    __shared__ float svals[MAX_SHMEM_SORT_SIZE];
    __shared__ float svals_copy[MAX_SHMEM_SORT_SIZE];

    const int fill_n = (n <= MAX_SHMEM_SORT_SIZE) ? n : MAX_SHMEM_SORT_SIZE;

    for (int s = 0; s < fill_n; s++) {
        const float* row = q_values + (long long)s * total_actions + b_off;
        float mx = row[0];
        for (int a = 1; a < b_size; a++) {
            float v = row[a];
            if (v > mx) mx = v;
        }
        svals[s] = mx;
    }

    /* ── Phase 2: percentile selection.
     * p05_idx = floor(fill_n * 0.05), p95_idx = floor(fill_n * 0.95), clamped. */
    int p05_idx = (int)((float)fill_n * 0.05f);
    int p95_idx = (int)((float)fill_n * 0.95f);
    if (p05_idx < 0)       p05_idx = 0;
    if (p05_idx >= fill_n) p05_idx = fill_n - 1;
    if (p95_idx < 0)       p95_idx = 0;
    if (p95_idx >= fill_n) p95_idx = fill_n - 1;

    float p05_val, p95_val;

    /* Bitonic sort path: exact, in-place, for power-of-2 fill_n. */
    bool is_pow2 = (fill_n > 1) && ((fill_n & (fill_n - 1)) == 0);

    if (is_pow2) {
        bitonic_sort_block(svals, fill_n);
        p05_val = svals[p05_idx];
        p95_val = svals[p95_idx];
    } else {
        /* Quickselect path (two independent passes; use svals_copy for the second
         * so svals is not clobbered between calls). */
        for (int i = 0; i < fill_n; i++) svals_copy[i] = svals[i];

        p05_val = nth_element_float(svals,      fill_n, p05_idx);
        p95_val = nth_element_float(svals_copy, fill_n, p95_idx);
    }

    /* ── Phase 3: EMA update into ISV[p05_slot] and ISV[p95_slot].
     * Read both previous values before any write (handles isv == isv_out aliasing). */
    float alpha = fminf(fmaxf(ema_alpha, 0.01f), 0.5f);

    float prev_p05 = isv[p05_slot];
    float prev_p95 = isv[p95_slot];

    isv_out[p05_slot] = (1.0f - alpha) * prev_p05 + alpha * p05_val;
    isv_out[p95_slot] = (1.0f - alpha) * prev_p95 + alpha * p95_val;
}
