// crates/ml/src/cuda_pipeline/c51_aux_dw_kernel.cu
//
// SP22 H6 Phase 3 α Step 8 → SP22 H6 vNext Phase D Step 8 — backward dW
// kernel for w_aux_to_q_dir [b0_size=4, K=3 outcomes] = 12 weights.
//
// Computes gradient of C51 loss w.r.t. the per-(action, outcome)
// atom-shift weights:
//   dW[a*K + k] =
//     Σ_b indicator(a_d_b == a) × isw_b × inv_batch × (SP_b / Δz_b)
//                                × state[121 + k]_b
//   + Σ_b indicator(a*_b == a) × isw_b × inv_batch × (-γ_eff_b × (1-done_b))
//                                × (SP_b / Δz_b) × next_state[121 + k]_b
//
// Phase D shift from K=2 / Phase 3:
//   - Phase 3 used state[121] only (scalar; K=2 head's recentered p_up)
//   - Phase D uses state[121+k] for k ∈ {0=Profit, 1=Stop, 2=Timeout}
//     (K=3 head's softmax probs from state slots [121..124))
//   - W is indexed (a, k) instead of just (a), so 12 weights total.
//
// where:
//   - SP_b = projection log-diff sum (saved by c51_loss_kernel forward):
//     SP_b = Σ_n p_target_n × (current_lp[upper_n] - current_lp[lower_n])
//   - a_d_b = taken direction action (= a0, computed from actions[b])
//   - a*_b  = sampled target action (saved by c51_loss_kernel as
//     aux_target_a_dir[b])
//   - Δz_b = per_sample_support[b * 12 + 0 * 3 + 2] (dir branch dz)
//   - γ_eff_b = effective gamma for dir branch from gamma_buf[b] (γ^n_steps,
//     ISV gamma_dir read inside)
//   - state_121_b = batch_states[b * state_dim + AUX_DIR_PROB_INDEX]
//   - next_state_121_b = next_batch_states[b * state_dim + AUX_DIR_PROB_INDEX]
//
// Output: dw_aux_buf[b0_size × K = 12] f32 (zeroed before launch —
// overwrite semantics).
//
// Discipline
// ──────────
// - Per pearl_no_atomicadd: ONE BLOCK PER (action, outcome) PAIR
//   (grid=(12,1,1), block=(256,1,1)); each block tree-reduces across
//   batch samples via shmem. No atomicAdd.
// - Per feedback_no_htod_htoh_only_mapped_pinned.md: all reads from device /
//   mapped pinned buffers; no HtoD inside kernel.
// - Per pearl_no_host_branches_in_captured_graph: launch happens OUTSIDE the
//   captured training graph (mirrors the existing post-c51_grad launches).
// - NULL-safety: when forward didn't populate scratch (aux_shift_active=false),
//   both aux_target_a_dir and aux_proj_logdiff_dir stay at alloc_zeros
//   sentinels; gradient stays at 0 → Adam step on W is a no-op.

#include <cuda_runtime.h>

#define BLOCK_THREADS 256
#define NUM_WARPS     (BLOCK_THREADS / 32)
#define MAX_DIR_ACTIONS 4

extern "C" __global__ void c51_aux_dw_kernel(
    /* Inputs from c51_loss_kernel forward scratch */
    const int*   __restrict__ aux_target_a_dir,    /* [B] sampled a* per sample */
    const float* __restrict__ aux_proj_logdiff_dir,/* [B] SP_b per sample */

    /* Inputs that the gradient formula needs */
    const int*   __restrict__ actions,             /* [B] factored action codes */
    const float* __restrict__ is_weights,          /* [B] PER importance weights (clamped to 10) */
    const float* __restrict__ dones,               /* [B] 0 or 1 */
    const float* __restrict__ gamma_buf,           /* [B] effective γ^n_steps */
    const float* __restrict__ per_sample_support,  /* [B, 4, 3] for Δz */
    const float* __restrict__ batch_states,        /* [B, state_dim] for state_121 */
    const float* __restrict__ next_batch_states,   /* [B, state_dim] for next_state_121 */

    /* Output */
    float* __restrict__ dw_aux,                    /* [b0_size × K = 12] gradient — written by tid==0 per block */

    /* Config */
    int batch_size,
    int b0_size,
    int b1_size,
    int b2_size,
    int b3_size,
    int aux_dir_prob_index,                        /* = 121 (= AUX_OUTCOME_PROFIT_INDEX) */
    int state_dim,                                 /* = 128 */
    /* SP22 H6 vNext Phase D (2026-05-14): K=3 outcome dimension (3 in
     * production = AUX_OUTCOME_K). Each block computes one (a, k) pair;
     * grid = (b0_size × K, 1, 1). Decoded as `a = blockIdx.x / K`,
     * `k = blockIdx.x % K`. The state slot is `aux_dir_prob_index + k`
     * (slots 121, 122, 123 for k = 0, 1, 2). */
    int aux_outcome_k
) {
    /* SP22 H6 vNext Phase D (2026-05-14): decode (a, k) from blockIdx.x.
     * Layout: grid block index = a * K + k. */
    int block_idx = blockIdx.x;
    int a = block_idx / aux_outcome_k;
    int k = block_idx - a * aux_outcome_k;
    int tid = threadIdx.x;
    if (a >= b0_size) return;
    if (a >= MAX_DIR_ACTIONS) return;
    if (k >= aux_outcome_k) return;

    /* State slot for this outcome class (k=0 → slot 121 = p_Profit,
     * k=1 → slot 122 = p_Stop, k=2 → slot 123 = p_Timeout). */
    int state_slot = aux_dir_prob_index + k;
    if (state_slot < 0 || state_slot >= state_dim) return;

    float inv_batch = 1.0f / (float)batch_size;

    /* Decode a_d (taken dir action = a0) inverse to c51_loss_kernel's decode:
     *   a0 = factored / (b1 * b2 * b3)
     * Note: b0_size=4 in production but read from kernel arg for ABI safety. */
    int divisor_a0 = b1_size * b2_size * b3_size;
    if (divisor_a0 <= 0) divisor_a0 = 1;

    /* Per-thread partial sum across samples assigned to this thread. */
    float local_sum = 0.0f;
    for (int b = tid; b < batch_size; b += BLOCK_THREADS) {
        int factored = actions[b];
        if (factored < 0) factored = 0;
        int a_d = factored / divisor_a0;
        if (a_d < 0 || a_d >= b0_size) a_d = 0;

        int a_star = aux_target_a_dir[b];
        if (a_star < 0 || a_star >= b0_size) a_star = 0;

        float sp = aux_proj_logdiff_dir[b];
        float isw = fminf(is_weights[b], 10.0f);
        float dz = per_sample_support[b * 12 + 0 * 3 + 2];
        /* DEFENSIVE NaN guards: any NaN/Inf in inputs → skip sample. */
        if (!isfinite(sp) || !isfinite(isw) || !isfinite(dz) || dz < 1e-7f) continue;
        float gamma_eff = gamma_buf[b];
        float done = dones[b];
        if (!isfinite(gamma_eff) || !isfinite(done)) continue;

        float dL_dDelta_online = sp / dz;
        float dL_dDelta_target = -gamma_eff * (1.0f - done) * dL_dDelta_online;

        if (a == a_d) {
            /* Phase D: state slot = aux_dir_prob_index + k (= 121, 122,
             * or 123 for k = 0, 1, 2 outcome classes). */
            float s_k = batch_states[(long long)b * state_dim + state_slot];
            if (!isfinite(s_k)) s_k = 0.0f;
            local_sum += inv_batch * isw * dL_dDelta_online * s_k;
        }
        if (a == a_star) {
            float ns_k = next_batch_states[(long long)b * state_dim + state_slot];
            if (!isfinite(ns_k)) ns_k = 0.0f;
            local_sum += inv_batch * isw * dL_dDelta_target * ns_k;
        }
    }

    /* Block tree-reduce (per pearl_no_atomicadd): warp shuffle + shmem. */
    __shared__ float shmem_warp[NUM_WARPS];
    unsigned int mask = 0xFFFFFFFF;
    int lane = tid & 31;
    int warp = tid >> 5;

    /* Warp-level reduce. */
    for (int offset = 16; offset > 0; offset >>= 1) {
        local_sum += __shfl_xor_sync(mask, local_sum, offset);
    }
    if (lane == 0) shmem_warp[warp] = local_sum;
    __syncthreads();

    /* Final reduce across warps (sequential in tid==0). */
    if (tid == 0) {
        float total = 0.0f;
        for (int w = 0; w < NUM_WARPS; w++) total += shmem_warp[w];
        /* Defensive NaN guard — set dW to 0 instead of propagating
         * NaN into W via Adam. Phase D: write to dw_aux[a * K + k]
         * instead of dw_aux[a]. */
        if (!isfinite(total)) total = 0.0f;
        dw_aux[a * aux_outcome_k + k] = total;
    }
}
