// rl_action_kernel.cu — GPU-resident Thompson sampler over the C51
// distributional Q-head output (Phase R4 of the integrated RL trainer
// rebuild; see
// docs/superpowers/plans/2026-05-23-integrated-rl-trainer-rebuild.md).
//
// Replaces the host-side Thompson loop the flawed Phase F shipped in
// step_with_lobsim (which violated `feedback_cpu_is_read_only` by
// DtoH-copying q_logits + sampling on CPU + HtoD-copying actions back
// to device).
//
// Per `pearl_thompson_for_distributional_action_selection`: Thompson
// is the canonical rollout selector under C51. For each batch + each
// action, sample ONE atom from that action's categorical (softmax
// over atom logits) and use the sampled atom's support value as the
// "sampled return"; pick the action with the highest sampled return.
// The companion `argmax_expected_q.cu` kernel does the Bellman-target
// argmax (over expected Q, not sampled).
//
// PRNG: per-batch xorshift32 state in `prng_state_d` (allocated +
// host-seeded once at trainer init per
// `pearl_scoped_init_seed_for_reproducibility`). To avoid the
// inter-thread race that a shared per-batch state would create, each
// per-action thread XORs the action index (via golden-ratio constant)
// into a thread-local copy of the per-batch state, advances locally
// for the atom sample, and only thread 0 writes the advanced per-batch
// state back. Reproducibility holds: identical (per-batch seed,
// blockDim, action grid, q_logits) → identical actions.
//
// Per `feedback_no_atomicadd`: thread 0 writes to `actions[b]` after
// __syncthreads. Per `pearl_no_host_branches_in_captured_graph`: no
// host branches; the only branches inside the kernel are device-side
// data-dependent (CDF walk, argmax).

#include <stdint.h>

#define N_ACTIONS    11
#define Q_N_ATOMS   21

__device__ static uint32_t xorshift32(uint32_t* state) {
    uint32_t x = *state;
    x ^= x << 13;
    x ^= x >> 17;
    x ^= x << 5;
    *state = x;
    return x;
}

// One block per batch (grid_dim.x = b_size). N_ACTIONS threads per
// block; each thread handles one action's atom sample, writes its
// sampled return to shared mem, then thread 0 argmaxes + writes
// actions[b] + advances the per-batch PRNG state.
//
// Inputs:
//   q_logits      [b_size, N_ACTIONS, Q_N_ATOMS] row-major
//   atom_supports [Q_N_ATOMS] e.g. linspace(Q_V_MIN, Q_V_MAX, Q_N_ATOMS)
//   prng_state    [b_size]  per-batch xorshift32 state (mutated in place)
// Outputs:
//   actions       [b_size]  Thompson-sampled action index per batch
extern "C" __global__ void rl_action_kernel(
    const float*    __restrict__ q_logits,
    const float*    __restrict__ atom_supports,
    uint32_t*       __restrict__ prng_state,
    int*            __restrict__ actions,
    int                          b_size
) {
    const int b = blockIdx.x;
    const int a = threadIdx.x;
    if (b >= b_size) return;
    if (a >= N_ACTIONS) return;

    __shared__ float sampled_returns[N_ACTIONS];

    // Derive thread-local PRNG state from the per-batch seed.
    // Golden-ratio constant decorrelates per-thread streams.
    const uint32_t base = prng_state[b];
    uint32_t local_state = base ^ ((uint32_t)a * 0x9E3779B1u);
    // A few warmup advances dispel correlations from the cheap mixing.
    #pragma unroll
    for (int w = 0; w < 4; ++w) xorshift32(&local_state);

    // Softmax over this action's atoms (numerically stabilised).
    const int row_off = (b * N_ACTIONS + a) * Q_N_ATOMS;
    float max_l = -INFINITY;
    #pragma unroll
    for (int i = 0; i < Q_N_ATOMS; ++i) {
        const float l = q_logits[row_off + i];
        if (l > max_l) max_l = l;
    }
    float sum_exp = 0.0f;
    #pragma unroll
    for (int i = 0; i < Q_N_ATOMS; ++i) {
        sum_exp += expf(q_logits[row_off + i] - max_l);
    }
    // Guard against numerical degeneracy (all-NaN or all-inf logits).
    if (!isfinite(sum_exp) || sum_exp <= 0.0f) sum_exp = 1.0f;

    // Categorical CDF walk against u ∈ [0, 1). Walk-to-end fallback
    // catches the case where cumulative-sum rounding edges below 1.0.
    const uint32_t r = xorshift32(&local_state);
    const float u = (float)r / 4294967296.0f;
    float cum = 0.0f;
    int chosen = Q_N_ATOMS - 1;
    for (int i = 0; i < Q_N_ATOMS; ++i) {
        cum += expf(q_logits[row_off + i] - max_l) / sum_exp;
        if (cum >= u) {
            chosen = i;
            break;
        }
    }

    sampled_returns[a] = atom_supports[chosen];
    __syncthreads();

    // Thread 0: argmax over per-action sampled returns; write action;
    // advance per-batch PRNG state for next call.
    if (a == 0) {
        int   best_a = 0;
        float best_v = sampled_returns[0];
        #pragma unroll
        for (int i = 1; i < N_ACTIONS; ++i) {
            if (sampled_returns[i] > best_v) {
                best_v = sampled_returns[i];
                best_a = i;
            }
        }
        actions[b] = best_a;

        uint32_t adv = base;
        xorshift32(&adv);
        prng_state[b] = adv;
    }
}
