// 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 f32 TD errors
//    Reads f32 TD errors directly (no separate cast kernel).
//    Standard PER (Schaul et al. 2016): priority p_i = |td_i| + eps,
//    sampling weight P(i) ∝ p_i^alpha. Stores BOTH columns of the
//    (priority, priority^alpha) row pair so the contract documented at
//    `gpu_replay_buffer.rs::priorities_pa_slice` (`priorities_pa[i] =
//    priorities[i]^alpha`) holds. atomicMax tracks the running raw-
//    priority max for `per_insert_pa`'s max-priority broadcast.
//
//    NOTE on prior bug (Class B Finding #2, fixed 2026-05-08): this
//    kernel previously stored `priorities[idx] = |td|^alpha + eps` and
//    then `priorities_pa[idx] = priorities^alpha`, applying alpha twice.
//    Sum-tree sampler reads `priorities_pa`, so the effective sampling
//    exponent was alpha² = 0.36 instead of alpha = 0.6 — high-TD events
//    were sampled only ~2.5× more than low-TD instead of ~10× as PER
//    intends. Fix below stores raw priority in `priorities` and the
//    α-weighted sampling weight in `priorities_pa`, matching the
//    Rainbow / Schaul-2016 standard.
//
//    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,
    /* SP21 T2.2 Phase 5+6 (2026-05-11) — combined E6/E7 signal-driven
     * PER alpha scaling. ISV[525] = winner P&L concentration from
     * E6 (top-decile mean P&L / all-trade mean P&L); ISV[526] =
     * hindsight counterfactual-P&L magnitude from E7. Both
     * normalised host-side; this kernel composes them
     * multiplicatively into `alpha_boost` and applies as
     * `alpha_eff = alpha + boost_delta` where boost_delta is
     * scaled to [0, 0.4] (so alpha_eff stays in
     * [alpha_baseline, alpha_baseline + 0.4]).
     *
     * Cold-start sentinel: when BOTH ISV[525] AND ISV[526] are at
     * sentinel 0.0 (pre-first-val-pass), kernel short-circuits
     * `alpha_eff = alpha` — no-op. Per
     * `pearl_first_observation_bootstrap`. NULL-tolerant on
     * `isv_signals` for tests / smoke scaffolds. */
    const float* __restrict__ isv_signals)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

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

    /* Raw PER priority p_i = |td_i| + eps (Schaul et al. 2016). */
    float td = fabsf((td_errors_bf16[i]));
    float new_prio = td + epsilon;
    if (new_prio < epsilon) new_prio = epsilon;
    if (new_prio > 1e6f) new_prio = 1e6f;

    priorities[idx] = new_prio;

    /* atomicMax tracks raw-priority max — matches `per_insert_pa`'s
     * `priorities[idx] = effective` (also raw) and the post-flush
     * `max_priority` broadcast that re-seeds new inserts. */
    int ival = __float_as_int(new_prio);
    atomicMax((int*)batch_max, ival);

    /* SP21 Phase 5+6+7: compose alpha_boost from three enrichment signals:
     *   - E6 winner_concentration (ISV[525]): top-decile pnl / all-mean pnl
     *   - E7 hindsight_magnitude (ISV[526]): mean |counterfactual_pnl|
     *   - E8 curriculum_concentration (ISV[527]): 1 - entropy/log(N) ∈ [0,1]
     *
     * All signals at sentinel 0.0 → boost_delta = 0 → alpha_eff = alpha
     * (cold-start no-op). Otherwise additive delta bounded to keep
     * alpha_eff in a sane range:
     *   - winner_term: (concentration - 1.0) clipped to [0, 2.0],
     *     scaled by 0.1 → [0, 0.2]. Concentration > 1.0 means top-decile
     *     drives more profit per trade than average.
     *   - hindsight_term: magnitude clipped to [0, 0.02] (typical
     *     counterfactual P&L scale ≈ 1-2%), scaled by 10.0 → [0, 0.2].
     *     Higher = more missed opportunities = need more PER focus.
     *   - curriculum_term: concentration ∈ [0, 1] (already normalised),
     *     scaled by 0.1 → [0, 0.1]. Higher = uneven mastery across
     *     curriculum segments = more focus on hard transitions.
     *   - Total delta clipped at 0.5 so alpha_eff ≤ alpha + 0.5. */
    float winner_conc   = (isv_signals != NULL) ? isv_signals[525] : 0.0f;
    float hindsight_mag = (isv_signals != NULL) ? isv_signals[526] : 0.0f;
    float curric_conc   = (isv_signals != NULL) ? isv_signals[527] : 0.0f;
    float boost_delta = 0.0f;
    if (winner_conc > 1e-6f || hindsight_mag > 1e-6f || curric_conc > 1e-6f) {
        float winner_term     = fmaxf(0.0f, fminf(2.0f, winner_conc - 1.0f)) * 0.1f;
        float hindsight_term  = fminf(0.02f, hindsight_mag) * 10.0f;
        float curriculum_term = fmaxf(0.0f, fminf(1.0f, curric_conc)) * 0.1f;
        boost_delta = fminf(0.5f, winner_term + hindsight_term + curriculum_term);
    }
    float alpha_eff = alpha + boost_delta;

    /* Sampling weight = priority^alpha_eff. Sum-tree (per_prefix_scan +
     * per_sample) walks this buffer; IS-weight (is_weights_f32) reads
     * the same column for normalisation. */
    float pa = powf(new_prio, alpha_eff);
    priorities_pa[idx] = pa;
}

// ---------------------------------------------------------------------------
// 2. per_insert_pa: insert priorities at max_priority with optional
//    SP15 Phase 3.5.5.b recovery oversample boost (per-transition
//    weighted via Phase 1.3.b-followup-B).
//
//    Takes raw priority values, applies the recovery multiplier
//        effective = raw * (1 + RECOVERY_OVERSAMPLE_WEIGHT × dd_trajectory_per_env[env_id])
//    and writes BOTH `priorities[idx] = effective` and
//    `priorities_pa[idx] = effective^alpha` so subsequent prefix-sum
//    sampling sees the boosted weight. The `priorities` write subsumes
//    the broadcast `scatter_insert_f32` previously used to seed
//    `priorities[idx] = max_priority` — per_insert_pa now owns both
//    columns of the (priority, priority^alpha) row in one launch.
//
//    Phase 1.3.b-followup-B (2026-05-07) — per-(env, t) trajectory
//    lookup. Wave 4.3 / Phase 3.5.5.b read `ISV[DD_TRAJECTORY_DECREASING
//    _INDEX=439]` ONCE per insert call and applied the SAME boost factor
//    to ALL transitions in the batch. That was statistically biased: in
//    a multi-env training step `n_envs * lookback * 2` transitions are
//    inserted in one batch, but the scalar ISV[439] reflected only the
//    LAST env's trajectory write, so any env's recovery context could
//    determine the boost for ALL inserted transitions regardless of
//    their actual env. The fix is per-(env, t) lookup: each transition
//    `j` maps to its own env via `env_id = (j % (n_envs * lookback)) /
//    lookback` (the experience-collector batch layout is
//    `[on_policy_main: N*L | counterfactual: N*L]` — see
//    `experience_kernels.cu::out_off = i * L + t` and
//    `cf_off = out_off + N * L`); the kernel reads
//    `dd_trajectory_per_env[env_id]` for that env's recovery flag.
//
//    `isv_dev_ptr` and `dd_trajectory_per_env_dev_ptr` may BOTH be
//    nullptr (passed from a setup that hasn't wired the SP15 ISV bus
//    yet — e.g. unit tests that exercise PER in isolation, or smoke
//    tests). When EITHER is null, the boost factor is 1.0 (no recovery
//    oversample), and behavior is bit-identical to the pre-Phase-3.5.5.b
//    kernel modulo the additional `priorities[idx]` write. The two
//    pointers degrade together (the per-env tile is meaningless without
//    the ISV[440] weight, and vice versa), so the null check is a
//    single combined gate.
//
//    Slot reads (literal indices to keep the kernel self-contained;
//    canonical names live in
//    `crates/ml/src/cuda_pipeline/sp15_isv_slots.rs` and the trainer
//    constructor pins those names to these indices via compile-time
//    assert_eq! regression checks). Mirrors the
//    `dd_trajectory_kernel.cu` literal-index convention:
//      ISV[440] RECOVERY_OVERSAMPLE_WEIGHT     (read; global scalar —
//                                              the oversample weight is
//                                              not per-env)
//      dd_trajectory_per_env[env_id]           (read; per-env tile)
//
//    `n_envs` and `lookback` are launch-time scalar args (fixed across
//    the run); `n_envs * lookback * 2` MUST equal `batch_size` when the
//    pointers are non-null (the experience-collector contract). When
//    `n_envs` or `lookback` is 0 the kernel falls back to `boost_factor
//    = 1.0` — matches the same-degraded-as-null contract for the
//    non-production callers (gpu_per_integration_test, smoke tests).
//
//    Launch: grid=(ceil(batch_size/256)), block=(256).
// ---------------------------------------------------------------------------
extern "C" __global__ void per_insert_pa(
    float* __restrict__ priorities_pa,
    float* __restrict__ priorities,                  /* SP15 3.5.5.b: also written here */
    const unsigned int* __restrict__ indices,
    const float* __restrict__ raw_priorities,
    const float* __restrict__ isv,                   /* nullable; SP15 ISV bus */
    const float* __restrict__ dd_trajectory_per_env, /* nullable; [n_envs] per-env trajectory */
    int n_envs,                                       /* SP15 1.3.b-followup-B: per-env dim */
    int lookback,                                     /* SP15 1.3.b-followup-B: per-env time dim */
    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;

    // SP15 Phase 3.5.5.b recovery oversample boost (per-transition
    // weighted per Phase 1.3.b-followup-B). Both pointers null OR
    // n_envs/lookback ≤ 0 → no boost (boost_factor=1.0), preserving
    // callers that don't wire the SP15 production path (oracle tests,
    // smoke tests, `gpu_per_integration_test`).
    //
    // env_id = (i % (n_envs * lookback)) / lookback maps the insert
    // batch index to its source env (the experience-collector batch
    // layout is `[on_policy_main: N*L | counterfactual: N*L]` — both
    // halves use `slot_it = j_main % (N*L)` semantics within their
    // half, so the modulo handles both halves uniformly; both halves
    // share the same `env_id = slot_it / L` mapping because both
    // `out_off = i*L+t` and `cf_off = N*L + i*L+t` are
    // env-major-then-time-stride). The mapping is independent of
    // whether the transition is the on-policy main or the counterfactual
    // sibling — they share the same env's DD trajectory because the
    // env's portfolio state machine is unitary across the cf-sibling
    // pair.
    float boost_factor = 1.0f;
    if (isv != nullptr && dd_trajectory_per_env != nullptr
        && n_envs > 0 && lookback > 0) {
        const int slot_it = i % (n_envs * lookback);
        const int env_id  = slot_it / lookback;
        // env_id ∈ [0, n_envs) by construction of the modulo above.
        const float decreasing = dd_trajectory_per_env[env_id];
        const float weight     = isv[/*RECOVERY_OVERSAMPLE_WEIGHT_INDEX*/ 440];
        boost_factor = 1.0f + weight * decreasing;
    }

    /* SP21 T2.2 Phase 7.5 (2026-05-11) — per-segment curriculum
     * boost on insert priority. E8's per-segment weights live at
     * ISV[528..536) (CURRICULUM_WEIGHT_{0..8}_INDEX), normalised so
     * sum = 1 across the 8 segments. Each insert gets tagged with
     * `seg_id = i % 8` (round-robin within the batch). Effective
     * priority gets multiplied by `N_SEGMENTS × weight[seg_id]`
     * — the `× N_SEGMENTS = 8` factor makes uniform weights
     * (= 1/8 each) a no-op (= 1.0), so the controller only
     * REDISTRIBUTES sampling priority across segments without
     * changing the overall priority scale.
     *
     * Cold-start short-circuit: when `isv[528]` is at sentinel 0.0
     * (pre-first-emit), curriculum_boost = 1.0 (uniform no-op).
     * Per `pearl_first_observation_bootstrap`. NULL-tolerant on
     * `isv` for tests / smoke scaffolds.
     *
     * Slot indices written as literals (mirroring the canonical
     * `dd_trajectory_kernel.cu` convention); names live in
     * `crates/ml/src/cuda_pipeline/sp21_isv_slots.rs` and the
     * trainer constructor pins those names to these indices via
     * the `curriculum_weight_index` accessor used in the producer
     * (host-side write path). */
    float curriculum_boost = 1.0f;
    if (isv != nullptr) {
        const float w0 = isv[/*CURRICULUM_WEIGHT_0_INDEX*/ 528];
        if (w0 > 1e-6f) {
            const int seg_id = i % 8;
            const float w = isv[528 + seg_id];
            /* × 8 makes uniform (1/8 per segment) = 1.0 no-op. */
            curriculum_boost = w * 8.0f;
            /* Guard against pathological zero-weight segments: a
             * single zero weight would zero the insert priority,
             * which would prevent that segment from ever being
             * sampled (sticky exclusion). Floor at 0.1× to keep
             * all segments at least nominally sampleable. */
            if (curriculum_boost < 0.1f) curriculum_boost = 0.1f;
        }
    }

    float effective = raw_priorities[i] * boost_factor * curriculum_boost;

    priorities[idx] = effective;
    priorities_pa[idx] = powf(effective, alpha);
}

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