// Standalone test-only kernel for Phase 0 verification.
// Implements the SAME math the Phase 2 production kernel will use,
// so Phase 0 tests verify the math AND Phase 2 Test 2.B verifies the
// production kernel matches this wrapper bit-for-bit (catches drift).
//
// All output buffers are device pointers into mapped pinned memory
// (cuMemHostAlloc with DEVICEMAP|PORTABLE). Kernels finish each write
// with __threadfence_system() so the host can read via read_volatile
// without an explicit memcpy_dtoh — see gpu_training_guard.rs MappedBuffer
// for the production reference of the same pattern.

#include <cuda_runtime.h>
#include <math_constants.h>

#ifndef N_IQN_QUANTILES
#define N_IQN_QUANTILES 5
#endif

// Inverse-CDF sample from C51 atom distribution.
// p[0..n_atoms] sums to ~1.0; atom_values[0..n_atoms] are the atom positions.
__device__ __forceinline__ float sample_c51_inverse_cdf(
    const float* __restrict__ p,
    const float* __restrict__ atom_values,
    int n_atoms,
    float u
) {
    float cum = 0.0f;
    for (int a = 0; a < n_atoms; a++) {
        cum += p[a];
        if (u < cum) return atom_values[a];
    }
    return atom_values[n_atoms - 1];  // numerical fallback
}

// Sample from IQN by uniform-τ interpolation over the 5 fixed quantiles.
// q[0..5] are the values at τ in {0.05, 0.25, 0.5, 0.75, 0.95}.
__device__ __forceinline__ float sample_iqn_quantile_interp(
    const float* __restrict__ q,
    float u
) {
    // Quantile boundaries: 0.05, 0.25, 0.50, 0.75, 0.95
    const float taus[N_IQN_QUANTILES] = {0.05f, 0.25f, 0.50f, 0.75f, 0.95f};

    if (u <= taus[0]) return q[0];
    if (u >= taus[N_IQN_QUANTILES - 1]) return q[N_IQN_QUANTILES - 1];

    for (int i = 0; i < N_IQN_QUANTILES - 1; i++) {
        if (u <= taus[i + 1]) {
            float w = (u - taus[i]) / (taus[i + 1] - taus[i]);
            return q[i] * (1.0f - w) + q[i + 1] * w;
        }
    }
    return q[N_IQN_QUANTILES - 1];
}

// Compute E[Q] from C51 atom probabilities + values.
__device__ __forceinline__ float compute_e_c51(
    const float* __restrict__ p,
    const float* __restrict__ atom_values,
    int n_atoms
) {
    float e = 0.0f;
    for (int a = 0; a < n_atoms; a++) {
        e += p[a] * atom_values[a];
    }
    return e;
}

// Compute E[Q] from IQN quantiles (mean of fixed quantiles).
__device__ __forceinline__ float compute_e_iqn(const float* __restrict__ q) {
    float s = 0.0f;
    for (int i = 0; i < N_IQN_QUANTILES; i++) s += q[i];
    return s / (float)N_IQN_QUANTILES;
}

// Linear-congruential RNG (deterministic from seed).
// Used only for tests; production uses Philox.
__device__ __forceinline__ float test_rand(unsigned int* state) {
    *state = (*state) * 1664525u + 1013904223u;
    return (float)(*state >> 8) / (float)(1u << 24);
}

// Test kernel: batched Thompson direction sampling (training mode).
//
// One thread per seed; thread `i` runs the full Thompson pick using
// `rng_state = base_seed + i` and writes `out_dir_idx[i]`. The output
// pointer is into mapped pinned memory; __threadfence_system() before
// thread exit makes the write visible to the host without memcpy_dtoh.
extern "C" __global__ void thompson_direction_test_batched(
    const float* __restrict__ c51_probs,        // [b0_size, n_atoms]
    const float* __restrict__ atom_values,      // [n_atoms]
    const float* __restrict__ iqn_quantiles,    // [b0_size, N_IQN_QUANTILES]
    int b0_size,
    int n_atoms,
    unsigned int base_seed,
    int n_seeds,
    int* __restrict__ out_dir_idx               // [n_seeds] — pinned-mapped
) {
    int seed_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (seed_idx >= n_seeds) return;

    unsigned int rng_state = base_seed + (unsigned int)seed_idx;
    float best_sample = -CUDART_INF_F;
    int best_d = 0;

    for (int d = 0; d < b0_size; d++) {
        float u_c51 = test_rand(&rng_state);
        float s_c51 = sample_c51_inverse_cdf(
            c51_probs + d * n_atoms, atom_values, n_atoms, u_c51
        );
        float u_iqn = test_rand(&rng_state);
        float s_iqn = sample_iqn_quantile_interp(
            iqn_quantiles + d * N_IQN_QUANTILES, u_iqn
        );
        // /2 dropped: argmax(a+b) == argmax((a+b)/2)
        float q_sample = s_c51 + s_iqn;
        if (q_sample > best_sample) {
            best_sample = q_sample;
            best_d = d;
        }
    }

    out_dir_idx[seed_idx] = best_d;
    __threadfence_system();   // make write PCIe-visible to host pointer
}

// Test kernel: argmax of (E_C51 + E_IQN) — eval mode.
// Single-threaded (argmax is deterministic over b0_size==4 directions).
// Output pointer is into mapped pinned memory; __threadfence_system()
// before exit ensures the host read sees the write.
extern "C" __global__ void argmax_eq_test(
    const float* __restrict__ c51_probs,
    const float* __restrict__ atom_values,
    const float* __restrict__ iqn_quantiles,
    int b0_size,
    int n_atoms,
    int* out_dir_idx
) {
    if (threadIdx.x != 0 || blockIdx.x != 0) return;

    float best_q = -CUDART_INF_F;
    int best_d = 0;

    for (int d = 0; d < b0_size; d++) {
        float e_c51 = compute_e_c51(c51_probs + d * n_atoms, atom_values, n_atoms);
        float e_iqn = compute_e_iqn(iqn_quantiles + d * N_IQN_QUANTILES);
        float e_joint = e_c51 + e_iqn;
        if (e_joint > best_q) { best_q = e_joint; best_d = d; }
    }

    *out_dir_idx = best_d;
    __threadfence_system();
}

// Diagnostic kernel: compute σ from C51 atoms (used by Test 0.F).
// out_sigma is mapped pinned memory; __threadfence_system() before thread exit.
extern "C" __global__ void compute_sigma_c51_test(
    const float* __restrict__ c51_probs,
    const float* __restrict__ atom_values,
    int b0_size,
    int n_atoms,
    float* out_sigma   // [b0_size]
) {
    int d = blockIdx.x * blockDim.x + threadIdx.x;
    if (d >= b0_size) return;

    const float* p = c51_probs + d * n_atoms;
    float mean = compute_e_c51(p, atom_values, n_atoms);
    float var = 0.0f;
    for (int a = 0; a < n_atoms; a++) {
        float diff = atom_values[a] - mean;
        var += p[a] * diff * diff;
    }
    out_sigma[d] = sqrtf(fmaxf(var, 0.0f));
    __threadfence_system();
}

// Diagnostic kernel: compute σ from IQN quantiles (used by Test 0.F).
// out_sigma is mapped pinned memory; __threadfence_system() before thread exit.
extern "C" __global__ void compute_sigma_iqn_test(
    const float* __restrict__ iqn_quantiles,
    int b0_size,
    float* out_sigma   // [b0_size]
) {
    int d = blockIdx.x * blockDim.x + threadIdx.x;
    if (d >= b0_size) return;

    const float* q = iqn_quantiles + d * N_IQN_QUANTILES;
    // IQR / 1.349 ≈ σ for Gaussian-like distributions
    float iqr = q[3] - q[1];  // Q_τ=0.75 - Q_τ=0.25
    out_sigma[d] = iqr / 1.349f;
    __threadfence_system();
}

// ─── Plan C Phase 2 Task 2 amendment (2026-04-29) ─────────────────────
//
// The kernels above (thompson_direction_test_batched, argmax_eq_test) are
// the Plan A Phase 0 entry points: they consume post-softmax `c51_probs`,
// a single linear `atom_values [n_atoms]`, AND `iqn_quantiles` (joint
// C51+IQN sample). They remain in place because Plan A tests 0.B-0.F (in
// `distributional_q_tests.rs`) still verify the underlying inverse-CDF +
// uniform-tau primitives via these wrappers — none of that math changed.
//
// The kernels below match the AMENDED production API (single-distribution
// C51 with adaptive per-direction support; no IQN inference at rollout):
//
//   - direction_thompson_v2_test: Thompson sample per direction using
//     `b_logits` + `per_sample_support` + optional `atom_positions`.
//     Single-distribution; no IQN term. Bit-equivalent to the production
//     direction-branch Thompson loop in
//     `experience_kernels.cu::experience_action_select`.
//
//   - argmax_eq_v2_test: argmax over E[Q] computed from softmax(logits)
//     and the same adaptive support. Bit-equivalent to the production
//     eval-mode argmax.
//
// Test 2.B (Plan C Task 7) launches the production kernel with the same
// inputs as `direction_thompson_v2_test` and asserts both produce
// identical `dir_idx` output across many seeds.
//
// RNG note: `direction_thompson_v2_test` uses the same LCG (`test_rand`)
// as `thompson_direction_test_batched` for parity with Plan A wrappers,
// NOT Philox. Test 2.B (Task 7) bridges to Philox at execution time by
// laying out `u` draws produced via Philox and passing them through this
// kernel's deterministic LCG path is NOT bit-equivalent to production —
// Test 2.B will need to either (a) drive both kernels from the same
// uniform array (preferred — strips RNG implementation choice from the
// equivalence test), or (b) invoke a separate Philox-driven variant. (a)
// is the simpler integration for Task 7.

// Numerically stable softmax over C51 logits. Matches the production
// `softmax_c51_inline` in `experience_kernels.cu`.
__device__ __forceinline__ void softmax_c51_test(
    const float* __restrict__ logits,
    int n_atoms,
    float* probs_out
) {
    float max_logit = logits[0];
    for (int a = 1; a < n_atoms; a++) {
        if (logits[a] > max_logit) max_logit = logits[a];
    }
    float sum = 0.0f;
    for (int a = 0; a < n_atoms; a++) {
        probs_out[a] = expf(logits[a] - max_logit);
        sum += probs_out[a];
    }
    float inv_sum = (sum > 0.0f) ? (1.0f / sum) : (1.0f / (float)n_atoms);
    for (int a = 0; a < n_atoms; a++) {
        probs_out[a] *= inv_sum;
    }
}

// Compute atom values for one (sample, direction). Matches the production
// `compute_atom_values_inline` in `experience_kernels.cu`.
__device__ __forceinline__ void compute_atom_values_test(
    const float* __restrict__ per_sample_support,
    const float* __restrict__ atom_positions,
    int i,
    int d,
    int b0_size,
    int n_atoms,
    float* atom_values_out
) {
    long long support_base = (long long)i * b0_size * 3 + (long long)d * 3;
    float v_min   = per_sample_support[support_base + 0];
    float v_max   = per_sample_support[support_base + 1];
    float delta_z = per_sample_support[support_base + 2];
    if (atom_positions != NULL) {
        long long pos_base = (long long)d * n_atoms;
        for (int a = 0; a < n_atoms; a++) {
            atom_values_out[a] = atom_positions[pos_base + a];
        }
    } else {
        for (int a = 0; a < n_atoms; a++) {
            atom_values_out[a] = v_min + (float)a * delta_z;
        }
    }
    (void)v_max;
}

#ifndef THOMPSON_TEST_MAX_ATOMS
#define THOMPSON_TEST_MAX_ATOMS 128
#endif

// Test kernel: single-distribution C51 Thompson direction sampling
// (training mode), aligned with the amended production API.
//
// Input layout (single-sample, since this is a standalone test wrapper):
//   - b_logits           [b0_size, n_atoms]  raw C51 logits
//   - per_sample_support [b0_size, 3]        — single-sample slice; the
//                                              caller passes a 1×4×3 tile
//                                              corresponding to sample i=0.
//   - atom_positions     [b0_size, n_atoms]  may be NULL (linear support).
//
// One thread per seed; `out_dir_idx[i]` records the Thompson pick for
// `rng_state = base_seed + i`.
extern "C" __global__ void direction_thompson_v2_test(
    const float* __restrict__ b_logits,
    const float* __restrict__ per_sample_support,
    const float* __restrict__ atom_positions,
    int b0_size,
    int n_atoms,
    unsigned int base_seed,
    int n_seeds,
    int* __restrict__ out_dir_idx               // [n_seeds] — pinned-mapped
) {
    int seed_idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (seed_idx >= n_seeds) return;

    unsigned int rng_state = base_seed + (unsigned int)seed_idx;
    float best_sample = -CUDART_INF_F;
    int best_d = 0;

    float probs_d[THOMPSON_TEST_MAX_ATOMS];
    float atom_vals_d[THOMPSON_TEST_MAX_ATOMS];

    for (int d = 0; d < b0_size; d++) {
        softmax_c51_test(b_logits + d * n_atoms, n_atoms, probs_d);
        // Sample index 0 — the standalone wrapper feeds a single tile.
        compute_atom_values_test(per_sample_support, atom_positions,
                                 /*i=*/0, d, b0_size, n_atoms, atom_vals_d);
        float u = test_rand(&rng_state);
        float q_sample = sample_c51_inverse_cdf(probs_d, atom_vals_d, n_atoms, u);
        if (q_sample > best_sample) {
            best_sample = q_sample;
            best_d = d;
        }
    }

    out_dir_idx[seed_idx] = best_d;
    __threadfence_system();
}

// Test kernel: argmax of E[Q] (eval mode), aligned with the amended
// production API. Single-threaded.
extern "C" __global__ void argmax_eq_v2_test(
    const float* __restrict__ b_logits,
    const float* __restrict__ per_sample_support,
    const float* __restrict__ atom_positions,
    int b0_size,
    int n_atoms,
    int* out_dir_idx
) {
    if (threadIdx.x != 0 || blockIdx.x != 0) return;

    float best_q = -CUDART_INF_F;
    int best_d = 0;

    float probs_d[THOMPSON_TEST_MAX_ATOMS];
    float atom_vals_d[THOMPSON_TEST_MAX_ATOMS];

    for (int d = 0; d < b0_size; d++) {
        softmax_c51_test(b_logits + d * n_atoms, n_atoms, probs_d);
        compute_atom_values_test(per_sample_support, atom_positions,
                                 /*i=*/0, d, b0_size, n_atoms, atom_vals_d);
        float e = compute_e_c51(probs_d, atom_vals_d, n_atoms);
        if (e > best_q) { best_q = e; best_d = d; }
    }

    *out_dir_idx = best_d;
    __threadfence_system();
}
