// crates/ml/src/cuda_pipeline/q_skew_kurtosis_kernel.cu
//
// SP5 Task A5 auxiliary: per-branch Q-distribution skew + kurtosis reduction.
// Reads save_q_online[B × 13] row-major; computes per-branch standardized
// skewness and excess kurtosis. Used by Pearl 5 to adapt the IQN τ schedule.
//
// Single-block, 4-thread (one per branch: Direction(4) / Magnitude(3) /
// Order(3) / Urgency(3)). Each thread reduces over its branch's
// (batch × n_actions) slice using a numerically-stable two-pass approach:
//   Pass 1: mean
//   Pass 2: m2 (variance), m3 (third central moment), m4 (fourth)
//   skew    = m3 / m2^1.5            (with EPS_DIV denom; clamped to [-3, +3])
//   ex_kurt = m4 / m2^2 - 3.0        (excess kurtosis; clamped to [-3, +30])
//
// Writes (skew[4], ex_kurt[4]) — 8 floats — to scratch_out at provided bases.
// __threadfence_system() after writes.
//
// No atomicAdd (feedback_no_atomicadd) — one thread per branch, independent
// scratch writes.

#include <cuda_runtime.h>

extern "C" __global__ void q_skew_kurtosis_update(
    const float* __restrict__ save_q_online_ptr,  // [B × 13] row-major
    int   batch_size,
    const int*   __restrict__ action_counts,       // [4] = {4, 3, 3, 3}
    const int*   __restrict__ action_offsets,      // [4] = {0, 4, 7, 10}
    float* __restrict__ scratch_out,
    int   skew_idx_base,                           // skew[4] at [base..base+4)
    int   ex_kurt_idx_base                         // excess kurtosis[4] at [base..base+4)
) {
    int b = threadIdx.x;
    if (blockIdx.x != 0 || b >= 4) return;

    int n_actions = action_counts[b];
    int act_off   = action_offsets[b];
    int total     = batch_size * n_actions;
    if (total == 0) return;

    // Pass 1: mean
    float sum = 0.0f;
    for (int i = 0; i < batch_size; ++i) {
        for (int a = 0; a < n_actions; ++a) {
            sum += save_q_online_ptr[i * 13 + act_off + a];
        }
    }
    float mean = sum / (float)total;

    // Pass 2: central moments m2, m3, m4
    float m2 = 0.0f, m3 = 0.0f, m4 = 0.0f;
    for (int i = 0; i < batch_size; ++i) {
        for (int a = 0; a < n_actions; ++a) {
            float d = save_q_online_ptr[i * 13 + act_off + a] - mean;
            float d2 = d * d;
            m2 += d2;
            m3 += d2 * d;
            m4 += d2 * d2;
        }
    }
    m2 /= (float)total;
    m3 /= (float)total;
    m4 /= (float)total;

    // EPS_DIV = 1e-12 is an Invariant 1 anchor — numerical stability for
    // m2^1.5 / m2^2 denominators. Prevents 0/0 when Q-distribution is constant.
    const float EPS_DIV = 1e-12f;
    float m2_15 = m2 * sqrtf(m2) + EPS_DIV;
    float m2_2  = m2 * m2 + EPS_DIV;

    float skew    = m3 / m2_15;
    float ex_kurt = m4 / m2_2 - 3.0f;

    // Structural envelopes (Invariant 1 anchors): clamp to keep downstream τ
    // shifts within [0.01, 0.99] and prevent runaway under degenerate input.
    skew    = fmaxf(-3.0f, fminf(3.0f,  skew));
    ex_kurt = fmaxf(-3.0f, fminf(30.0f, ex_kurt));

    scratch_out[skew_idx_base + b]    = skew;
    scratch_out[ex_kurt_idx_base + b] = ex_kurt;
    __threadfence_system();
}
