// crates/ml/src/cuda_pipeline/param_group_oracle_kernel.cu
//
// SP4 Layer A Task A7 — Pearl B fused per-param-group statistics oracle.
//
// Single kernel per param-group reads `(params, grads, adam_m, adam_v)` once
// and runs 4 sequential passes (5 for the trunk group):
//
//   Pass A: WEIGHT_BOUND[group] = p99(|params|)   via sp4_histogram_p99<256>
//   Pass B: ADAM_M_BOUND[group] = p99(|adam_m|)   via sp4_histogram_p99<256>
//   Pass C: ADAM_V_BOUND[group] = p99(|adam_v|)   via sp4_histogram_p99<256>
//   Pass D: WD_RATE[group]      = |Σ w·g| / max(Σ w², EPS_DIV)
//                                 side products: mean|g|, mean|w|
//   Pass E: L1_LAMBDA[trunk]    = (mean|g| / max(mean|w|, EPS_DIV))
//                                 × (log K − H) / log K
//                                 (Group 0 only; H = Shannon entropy of
//                                  per-feature gradient L2 norms)
//
// Pearl B 4× memory-bandwidth reduction vs four naive single-stat producer
// kernels: each of params/grads/m/v is read once across all 4-5 outputs
// rather than four times.
//
// Caller contract (same as the rest of the SP4 producer family):
//   - grid_dim  = (1, 1, 1)
//   - block_dim = (256, 1, 1)
//   - dynamic shared memory ≥ (256/32) × 256 × sizeof(int) = 8192 bytes
//     (matches `sp4_histogram_p99.cuh`'s per-warp tile contract; Pass D
//     reuses the same dynamic shmem region after Pass C finishes).
//
// The trunk-only Pass E uses statically allocated `__shared__ float
// s_g_feat[K_IN_MAX]` for per-feature gradient L2 norms; K_IN_MAX is set to
// 1024 to cover any plausible trunk input dimension (production trunk is
// 40 + OFI/TLOB widening, well under 1024).
//
// Cold-path launch — NOT in the captured graph for now. Task A10 may
// relocate select producers into the captured graph; this kernel mirrors
// the launch-shape that A5/A6 established so the migration is structural-
// only at that point.

#include "sp4_histogram_p99.cuh"

#define EPS_DIV 1.0e-8f

// Cap on Pass-E per-feature shared scratch. Production trunk K_in is much
// smaller (40 base features; widened with OFI/TLOB to <128). 1024 is a
// generous ceiling that still fits comfortably in a single block's static
// shmem budget alongside `s_bins[256]` and `s_step_max`.
#define K_IN_MAX 1024

// Helper: block-wide tree-reduce of `val` across the 256 threads via
// `s_buf` (already-allocated `__shared__ float[256]` storage). Returns the
// reduced value to thread 0; other threads' return value is undefined.
// Uses `__syncthreads()` between halves; no atomicAdd per
// `feedback_no_atomicadd`.
__device__ __forceinline__ float block_reduce_sum_256(float val, float* s_buf, int tid) {
    s_buf[tid] = val;
    __syncthreads();
    for (int s = 128; s > 0; s >>= 1) {
        if (tid < s) {
            s_buf[tid] += s_buf[tid + s];
        }
        __syncthreads();
    }
    return s_buf[0];
}

// SP4 Task A7 fix-up #2 (2026-04-30): the kernel now accepts an array of
// sub-buffer pointers + counts so groups with non-contiguous storage
// (Curiosity's `[w1, b1, w2, b2]` layout) can be described as a single
// logical group. For groups 0-6 the launcher passes `n_sub=1` and the
// behaviour is identical to the original single-pointer signature. For
// group 7 (Curiosity) it passes `n_sub=4` and the union of the four
// sub-buffers is treated as one distribution for Pass A/B/C p99 and
// Pass D (`Σ w·g`, `Σ w²`, `Σ|g|`, `Σ|w|`). Pass E (trunk-only) operates
// on the contiguous trunk grad-w slice via the `g_idx == 0` arm and is
// guaranteed to see `n_sub=1` (Curiosity has `l1_lambda_scratch_idx==-1`).
extern "C" __global__ void param_group_oracle_update(
    const unsigned long long* __restrict__ params_ptrs,    /* [n_sub] device pointers */
    const unsigned long long* __restrict__ grads_ptrs,     /* [n_sub] device pointers */
    const unsigned long long* __restrict__ adam_m_ptrs,    /* [n_sub] device pointers */
    const unsigned long long* __restrict__ adam_v_ptrs,    /* [n_sub] device pointers */
    const int*                __restrict__ sub_counts,     /* [n_sub] per-sub-buffer element counts */
    int          n_sub,                /* number of sub-buffers (1 for groups 0-6, 4 for group 7) */
    int          total_count,          /* sum of sub_counts (used for p99 normalisation + WD inv_count) */
    int          k_in,                 /* trunk input dim (Pass E only; 0 for non-trunk) */
    int          h_dim,                /* trunk output dim of layer (Pass E only) */
    float*       __restrict__ scratch_buf,
    int          weight_scratch_idx,
    int          adam_m_scratch_idx,
    int          adam_v_scratch_idx,
    int          wd_rate_scratch_idx,
    int          l1_lambda_scratch_idx /* -1 if not group 0 */
) {
    if (blockIdx.x != 0) return;
    const int tid = threadIdx.x;

    // ── Pass A: WEIGHT_BOUND[group] = p99(|params|) ─────────────────────
    {
        float p99 = sp4_histogram_p99_multi<256>(params_ptrs, sub_counts, n_sub, total_count);
        if (tid == 0) {
            scratch_buf[weight_scratch_idx] = p99;
            __threadfence_system();   // make host-mapped write visible
        }
    }
    __syncthreads();

    // ── Pass B: ADAM_M_BOUND[group] = p99(|adam_m|) ─────────────────────
    {
        float p99 = sp4_histogram_p99_multi<256>(adam_m_ptrs, sub_counts, n_sub, total_count);
        if (tid == 0) {
            scratch_buf[adam_m_scratch_idx] = p99;
            __threadfence_system();
        }
    }
    __syncthreads();

    // ── Pass C: ADAM_V_BOUND[group] = p99(|adam_v|) ─────────────────────
    {
        float p99 = sp4_histogram_p99_multi<256>(adam_v_ptrs, sub_counts, n_sub, total_count);
        if (tid == 0) {
            scratch_buf[adam_v_scratch_idx] = p99;
            __threadfence_system();
        }
    }
    __syncthreads();

    // ── Pass D: WD_RATE[group] = |Σ w·g| / max(Σ w², EPS_DIV) ──────────
    // Side products: mean|g| = Σ|g|/total_count, mean|w| = Σ|w|/total_count
    // for Pass E.  Iterates `n_sub` sub-buffers in the per-thread accumulator
    // loop; for `n_sub=1` this is equivalent to the original single-buffer
    // implementation.
    //
    // Reuses the histogram dynamic shmem region (`s_warp_tiles`) as the
    // per-thread accumulator buffer; the Pass C `__syncthreads` above
    // ensures the previous histogram computation has fully drained before
    // we overwrite. We reinterpret `s_warp_tiles` (declared as `int*` in
    // `sp4_histogram_p99.cuh`) via `float*` cast: 8192 bytes = 2048 floats
    // ≥ 256 floats × 4 accumulators = 1024 floats. Safe.
    extern __shared__ int s_warp_tiles[];
    float* s_buf = reinterpret_cast<float*>(s_warp_tiles);

    float t_sum_wg     = 0.0f;
    float t_sum_w2     = 0.0f;
    float t_sum_abs_g  = 0.0f;
    float t_sum_abs_w  = 0.0f;

    for (int s = 0; s < n_sub; ++s) {
        const float* w_buf = reinterpret_cast<const float*>(params_ptrs[s]);
        const float* g_buf = reinterpret_cast<const float*>(grads_ptrs[s]);
        const int    cnt_s = sub_counts[s];
        for (int i = tid; i < cnt_s; i += 256) {
            float w = w_buf[i];
            float g = g_buf[i];
            t_sum_wg    += w * g;
            t_sum_w2    += w * w;
            t_sum_abs_g += fabsf(g);
            t_sum_abs_w += fabsf(w);
        }
    }

    // Four sequential block-reduces sharing the same 256-float shmem region.
    float sum_wg    = block_reduce_sum_256(t_sum_wg,    s_buf, tid);
    float sum_w2    = block_reduce_sum_256(t_sum_w2,    s_buf, tid);
    float sum_abs_g = block_reduce_sum_256(t_sum_abs_g, s_buf, tid);
    float sum_abs_w = block_reduce_sum_256(t_sum_abs_w, s_buf, tid);

    // Stash mean|g| / mean|w| in static shmem so Pass E (group 0) sees them
    // after any further `__syncthreads`. We use `s_bins[]` (also int) by
    // reinterpret-cast — same trick as `sp4_histogram_p99.cuh`'s Pass 1
    // does for max-reduce. With 256 ints available we only need 2 slots.
    __shared__ float s_mean_abs_g;
    __shared__ float s_mean_abs_w;
    __shared__ float s_wd_rate;
    if (tid == 0) {
        const float inv_count = (total_count > 0) ? (1.0f / (float)total_count) : 0.0f;
        s_mean_abs_g = sum_abs_g * inv_count;
        s_mean_abs_w = sum_abs_w * inv_count;
        const float denom = fmaxf(sum_w2, EPS_DIV);
        s_wd_rate = fabsf(sum_wg) / denom;
        scratch_buf[wd_rate_scratch_idx] = s_wd_rate;
        __threadfence_system();
    }
    __syncthreads();

    // ── Pass E (group 0 trunk only): L1_LAMBDA_TRUNK ────────────────────
    // Per-feature gradient L2 norm `||g_w[:,i]||` for i ∈ [0, k_in).
    // Trunk weight matrix layout per `compute_param_sizes`: tensor 0 is
    // `W_h_s1 [k_in × h_dim]` row-major, so feature `i` occupies the
    // `i`-th column → strided indices `[i, k_in + i, 2*k_in + i, ...,
    // (h_dim-1)*k_in + i]` in row-major. We compute the per-feature L2
    // norm by squaring + summing across `h_dim` then sqrt.
    //
    // Pass E reads the contiguous trunk grad-w sub-buffer at
    // `grads_ptrs[0]` directly — group 0 is contiguous so `n_sub == 1` and
    // the `[h_dim × k_in]` weight-matrix layout sits in sub_counts[0].
    // For multi-sub-buffer groups (Curiosity) `l1_lambda_scratch_idx==-1`
    // and Pass E is gated off entirely.
    //
    // Feature normalization → discrete distribution; Shannon entropy H
    // → entropy_deficit = (log K − H) / log K ∈ [0, 1] where K = k_in
    // (1 = perfectly concentrated on one feature; 0 = uniform). The L1
    // λ for the trunk is then
    //
    //     L1_LAMBDA_TRUNK = (mean|g| / max(mean|w|, EPS_DIV)) × deficit
    //
    // — magnitude scale × concentration. Single thread (thread 0) does the
    // normalize + entropy computation; per-feature L2 squares are produced
    // by the block in parallel.
    if (l1_lambda_scratch_idx >= 0 && k_in > 0 && h_dim > 0 && k_in <= K_IN_MAX) {
        __shared__ float s_g_feat[K_IN_MAX];

        // Zero the per-feature scratch (cooperative across the block).
        for (int i = tid; i < k_in; i += 256) {
            s_g_feat[i] = 0.0f;
        }
        __syncthreads();

        // Read the contiguous trunk grads sub-buffer. Pass E is only
        // dispatched for group 0 (DqnTrunk) which has `n_sub==1`, so
        // `grads_ptrs[0]` is the canonical trunk grad-w pointer.
        const float* grads_e = reinterpret_cast<const float*>(grads_ptrs[0]);

        // Each thread accumulates squared gradients for a strided slice
        // of the [k_in × h_dim] grad_w matrix (= grads[0..k_in*h_dim)).
        // Per-feature accumulation is naturally race-free across threads
        // that hit different (i, h) cells; collisions on the same `i`
        // across threads are resolved by reducing through shared memory
        // — but to avoid atomicAdd, we instead serialize per-i: threads
        // collaborate by reduction across `h ∈ [0, h_dim)` for each `i`
        // in turn. That's k_in × (h_dim parallel reduce) = k_in × ~h_dim
        // ops; for k_in ≤ 1024 and h_dim ≤ ~1024 this is fine for a cold-
        // path producer.
        const int kh = k_in * h_dim;

        // Cooperative thread-strided accumulation into s_g_feat[i] using
        // sequential synchronisation per feature. Avoids atomicAdd while
        // keeping the cost linear in k_in × h_dim.
        for (int i = 0; i < k_in; ++i) {
            float partial = 0.0f;
            for (int h = tid; h < h_dim; h += 256) {
                int idx = h * k_in + i;
                if (idx < kh) {
                    float gv = grads_e[idx];
                    partial += gv * gv;
                }
            }
            float sq_sum = block_reduce_sum_256(partial, s_buf, tid);
            if (tid == 0) {
                s_g_feat[i] = sqrtf(sq_sum);
            }
            __syncthreads();
        }

        // Single-thread normalization + Shannon entropy.
        if (tid == 0) {
            float total = 0.0f;
            for (int i = 0; i < k_in; ++i) total += s_g_feat[i];

            float deficit = 0.0f;
            if (total > EPS_DIV && k_in > 1) {
                const float log_k = logf((float)k_in);
                float entropy = 0.0f;
                for (int i = 0; i < k_in; ++i) {
                    float p = s_g_feat[i] / total;
                    if (p > EPS_DIV) {
                        entropy -= p * logf(p);
                    }
                }
                deficit = (log_k - entropy) / log_k;
                if (deficit < 0.0f) deficit = 0.0f;
                if (deficit > 1.0f) deficit = 1.0f;
            }

            const float mean_abs_g = s_mean_abs_g;
            const float mean_abs_w = s_mean_abs_w;
            const float scale = mean_abs_g / fmaxf(mean_abs_w, EPS_DIV);
            const float l1_lambda = scale * deficit;

            scratch_buf[l1_lambda_scratch_idx] = l1_lambda;
            __threadfence_system();
        }
    }
}
