// rl_sample_noise.cu — Device-side factored noise generation for NoisyNet.
//
// Replaces host-side ChaCha8 → f(x) transform → mapped-pinned upload
// so NoisyLinear can sit inside a CUDA Graph (no host RNG in captured
// region).
//
// Single block, max(in_dim, out_dim) threads. Each thread generates up
// to two noise values: one for eps_in (if tid < in_dim) and one for
// eps_out (if tid < out_dim). The factored noise transform
// f(x) = sign(x) × sqrt(|x|) is applied in-kernel.
//
// Per feedback_no_atomicadd: no atomics — per-thread writes only.
// Per feedback_cpu_is_read_only: pure device-side sampling.

#include <stdint.h>
#include <math.h>

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

// f(x) = sign(x) × sqrt(|x|)  —  Fortunato et al. 2017 factored noise.
__device__ static float factored_noise_f(float x) {
    return copysignf(sqrtf(fabsf(x)), x);
}

// Convert xorshift32 output to a float in [-1, +1).
__device__ static float rand_signed(uint32_t* state) {
    const uint32_t r = xorshift32(state);
    // Map to [-1, +1) via (r / 2^31) - 1.
    return (float)(int32_t)(r) * (1.0f / 2147483648.0f);
}

// Grid=(1, 1, 1), Block=(max(in_dim, out_dim), 1, 1).
//
// Inputs:
//   prng_state  [1]        persistent xorshift32 seed (mutated)
// Outputs:
//   eps_in      [in_dim]   f(rand) noise for input dimension
//   eps_out     [out_dim]  f(rand) noise for output dimension
extern "C" __global__ void rl_sample_noise(
    uint32_t* __restrict__ prng_state,
    float*    __restrict__ eps_in,
    float*    __restrict__ eps_out,
    int                    in_dim,
    int                    out_dim
) {
    const int tid = threadIdx.x;
    const int max_dim = (in_dim > out_dim) ? in_dim : out_dim;
    if (tid >= max_dim) return;

    // Self-seed on first call (alloc_zeros → 0).
    uint32_t seed = prng_state[0];
    if (seed == 0u) seed = 0x4E01u * 2654435761u + 0xCAFEu;

    // Derive thread-local PRNG from the global seed mixed with tid.
    uint32_t local_state = seed ^ (uint32_t)(tid * 2654435761u);

    #pragma unroll
    for (int w = 0; w < 4; ++w) xorshift32(&local_state);

    // Generate eps_in if tid < in_dim.
    if (tid < in_dim) {
        const float raw = rand_signed(&local_state);
        eps_in[tid] = factored_noise_f(raw);
    }

    // Generate eps_out if tid < out_dim (independent draw).
    if (tid < out_dim) {
        const float raw = rand_signed(&local_state);
        eps_out[tid] = factored_noise_f(raw);
    }

    // Thread 0 advances the persistent global seed. Uses `seed`
    // (includes self-seed) not `prng_state[0]` (may be zero on first call).
    if (tid == 0) {
        uint32_t adv = seed;
        #pragma unroll
        for (int w = 0; w < 8; ++w) xorshift32(&adv);
        prng_state[0] = adv;
    }
}
