/**
 * Standalone C51 distributional RL loss kernel for Branching DQN.
 *
 * Phase 2 Task 6: Extract C51 loss from dqn_forward_loss_kernel into a
 * standalone kernel that operates on batch-level tensors (pre-computed forward
 * pass outputs), eliminating the 46KB/thread register spill that reduces the
 * fused kernel to 1.56% occupancy.
 *
 * Design:
 *   - Input: value_logits + advantage_logits per branch from both online,
 *     target, and online_next networks (pre-computed by cuBLAS forward pass).
 *   - Grid: (batch_size, 1, 1)  — one block per sample.
 *   - Block: (256, 1, 1)        — 256 threads cooperate on NUM_ATOMS=51 work.
 *   - Shared memory: ~2.8 KB per block (vs 46 KB/thread in fused kernel).
 *
 * Does NOT modify dqn_forward_loss_kernel — that kernel remains as fallback.
 *
 * Requires common_device_functions.cuh prepended via NVRTC source concatenation
 * (same mechanism as dqn_training_kernel.cu). The following constants must be
 * injected via NVRTC #define before this source is compiled:
 *   NUM_ATOMS, V_MIN, V_MAX, BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE
 */

/* ── Compile-time defaults (overridden by NVRTC injection) ───────────── */
#ifndef MAX_PER_SAMPLE_CE
#define MAX_PER_SAMPLE_CE 50.0f  /* Layer 1: 10x converged loss — breaks PER feedback loop */
#endif
#ifndef LABEL_SMOOTHING_EPS
#define LABEL_SMOOTHING_EPS 0.01f /* Layer 2: 1% uniform mix — prevents target collapse */
#endif
#ifndef NUM_ATOMS
#define NUM_ATOMS 51
#endif
#ifndef BRANCH_0_SIZE
#define BRANCH_0_SIZE 9
#endif
#ifndef BRANCH_1_SIZE
#define BRANCH_1_SIZE 3
#endif
#ifndef BRANCH_2_SIZE
#define BRANCH_2_SIZE 3
#endif
#ifndef V_MIN
#define V_MIN (-25.0f)
#endif
#ifndef V_MAX
#define V_MAX (25.0f)
#endif
#define DELTA_Z ((V_MAX - V_MIN) / (float)(NUM_ATOMS - 1))

/* Total atoms per branch output */
#define BRANCH_0_ATOMS (BRANCH_0_SIZE * NUM_ATOMS)
#define BRANCH_1_ATOMS (BRANCH_1_SIZE * NUM_ATOMS)
#define BRANCH_2_ATOMS (BRANCH_2_SIZE * NUM_ATOMS)
/* Largest branch output size (floats) — used to size shmem regions */
#define MAX_BRANCH_ATOMS (MAX_BRANCH_SIZE * NUM_ATOMS)
#define MAX_BRANCH_SIZE 9  /* max(BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE) */
#define NUM_BRANCHES 3
/* Maximum warp count for BLOCK_THREADS = 256 */
#define NUM_WARPS 8  /* 256 / 32 */

/* ── Block-level thread count (must match launch config) ─────────────── */
/* 256 threads: 8 warps, enough to saturate the issue queue for small NUM_ATOMS.
 * At NUM_ATOMS=51, each thread handles at most 1 atom in most passes. */
#define BLOCK_THREADS 256

/* ── Shared memory layout (all regions in floats) ───────────────────────
 *
 * Region              Size (floats)   Purpose
 * ─────────────────────────────────────────────────────────────────
 * shmem_support       NUM_ATOMS=51    Precomputed z_j = V_MIN + j*DELTA_Z
 * shmem_val           NUM_ATOMS=51    Value logits for current network/state
 * shmem_adv           MAX_BRANCH_ATOMS=255  Adv logits for all actions in branch
 * shmem_lp            NUM_ATOMS=51    log_softmax result (current/target action)
 * shmem_proj          NUM_ATOMS=51    Bellman-projected target distribution
 * shmem_current_lp    NUM_ATOMS=51    Saved current action log-probs (for CE)
 * shmem_blk_reduce    NUM_WARPS=8     Block-level reduction scratch
 * ─────────────────────────────────────────────────────────────────
 * Total floats: 51+51+255+51+51+51+8 = 518
 * Total bytes:  518 * 4 = 2072 bytes ≈ 2.0 KB per block
 * H100 SM: 228 KB / 2.0 KB = 114 concurrent blocks per SM
 *
 * Note: shmem_blk_reduce is NUM_WARPS floats and reused for all block
 * reductions (max-finding, sum-finding, cross-entropy). One reduction
 * at a time — never overlapping.
 */
#define SHMEM_SUPPORT_OFF      0
#define SHMEM_VAL_OFF          (SHMEM_SUPPORT_OFF   + NUM_ATOMS)
#define SHMEM_ADV_OFF          (SHMEM_VAL_OFF        + NUM_ATOMS)
#define SHMEM_LP_OFF           (SHMEM_ADV_OFF        + MAX_BRANCH_SIZE * NUM_ATOMS)
#define SHMEM_PROJ_OFF         (SHMEM_LP_OFF         + NUM_ATOMS)
#define SHMEM_CURRENT_LP_OFF   (SHMEM_PROJ_OFF       + NUM_ATOMS)
#define SHMEM_REDUCE_OFF       (SHMEM_CURRENT_LP_OFF + NUM_ATOMS)
#define SHMEM_TOTAL_FLOATS     (SHMEM_REDUCE_OFF     + NUM_WARPS)

/* ── Device helpers ───────────────────────────────────────────────────── */

/**
 * Block-wide max reduction.
 * All BLOCK_THREADS threads contribute `val`; returns the block max.
 * Uses shmem_reduce[NUM_WARPS] as scratch (must be accessible by all threads).
 */
__device__ __forceinline__ float block_reduce_max(
    float val, float* __restrict__ shmem_reduce, int tid
) {
    unsigned int mask = 0xFFFFFFFF;
    int lane = tid & 31;
    int warp = tid >> 5;

    /* Warp-level reduction */
    for (int offset = 16; offset > 0; offset >>= 1)
        val = fmaxf(val, __shfl_xor_sync(mask, val, offset));

    if (lane == 0) shmem_reduce[warp] = val;
    __syncthreads();

    /* Thread 0's warp reduces across warp results */
    float result = -1e30f;
    if (tid < NUM_WARPS)
        result = shmem_reduce[tid];
    if (tid < 32) {
        for (int offset = NUM_WARPS >> 1; offset > 0; offset >>= 1)
            result = fmaxf(result, __shfl_xor_sync(mask, result, offset));
        if (tid == 0) shmem_reduce[0] = result;
    }
    __syncthreads();
    return shmem_reduce[0];
}

/**
 * Block-wide sum reduction.
 * All BLOCK_THREADS threads contribute `val`; returns the block sum.
 */
__device__ __forceinline__ float block_reduce_sum(
    float val, float* __restrict__ shmem_reduce, int tid
) {
    unsigned int mask = 0xFFFFFFFF;
    int lane = tid & 31;
    int warp = tid >> 5;

    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_xor_sync(mask, val, offset);

    if (lane == 0) shmem_reduce[warp] = val;
    __syncthreads();

    float result = 0.0f;
    if (tid < NUM_WARPS)
        result = shmem_reduce[tid];
    if (tid < 32) {
        for (int offset = NUM_WARPS >> 1; offset > 0; offset >>= 1)
            result += __shfl_xor_sync(mask, result, offset);
        if (tid == 0) shmem_reduce[0] = result;
    }
    __syncthreads();
    return shmem_reduce[0];
}

/**
 * Block-cooperative log_softmax over NUM_ATOMS elements.
 *
 * Reads from `logits[0..NUM_ATOMS]` in shared memory.
 * Writes result to `out[0..NUM_ATOMS]` in shared memory.
 * Both `logits` and `out` may alias the same region.
 * `shmem_reduce[NUM_WARPS]` is used as scratch.
 * Caller must __syncthreads() before reading `out`.
 */
__device__ void block_log_softmax(
    const float* __restrict__ logits,  /* [NUM_ATOMS] shmem */
    float*       __restrict__ out,     /* [NUM_ATOMS] shmem output */
    float*       __restrict__ shmem_reduce, /* [NUM_WARPS] scratch */
    int tid
) {
    /* Pass 1: block-max */
    float local_max = -1e30f;
    for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS)
        local_max = fmaxf(local_max, logits[i]);
    float block_max = block_reduce_max(local_max, shmem_reduce, tid);

    /* Pass 2: sum(exp(x - max)) */
    float local_sum = 0.0f;
    for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS)
        local_sum += expf(logits[i] - block_max);
    float block_sum = block_reduce_sum(local_sum, shmem_reduce, tid);
    float log_sum_exp = logf(block_sum + 1e-10f);

    /* Pass 3: write log(softmax) = x - max - log_sum_exp */
    for (int i = tid; i < NUM_ATOMS; i += BLOCK_THREADS)
        out[i] = logits[i] - block_max - log_sum_exp;
    __syncthreads();
}

/**
 * Block-cooperative expected Q value: E[Q] = sum_j(exp(log_probs[j]) * z_j).
 *
 * `log_probs` and `support` are in shared memory.
 * Returns the scalar expected Q (identical on all threads after call).
 */
__device__ float block_expected_q(
    const float* __restrict__ log_probs,  /* [NUM_ATOMS] shmem */
    const float* __restrict__ support,    /* [NUM_ATOMS] shmem */
    float*       __restrict__ shmem_reduce, /* [NUM_WARPS] scratch */
    int tid
) {
    float local_sum = 0.0f;
    for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
        local_sum += expf(log_probs[j]) * support[j];
    return block_reduce_sum(local_sum, shmem_reduce, tid);
}

/**
 * Block-cooperative Bellman projection.
 *
 * Projects `target_probs[NUM_ATOMS]` (probabilities, not log-probs) onto the
 * C51 support using T_z = clamp(r + gamma * z_j * (1 - done), V_MIN, V_MAX).
 * Result accumulated into `projected[NUM_ATOMS]` (must be zero-initialized).
 *
 * With BLOCK_THREADS=256 and NUM_ATOMS=51, each thread handles at most 1 source
 * atom. The per-thread `local_proj[NUM_ATOMS]` register array therefore has at
 * most 2 non-zero entries (lower and upper neighbors for that one atom).
 * The warp butterfly + block accumulate then merges them.
 */
__device__ void block_bellman_project(
    const float* __restrict__ target_probs, /* [NUM_ATOMS] shmem probabilities */
    float reward, float done, float gamma,
    float*       __restrict__ projected,     /* [NUM_ATOMS] shmem, zero-initialized */
    const float* __restrict__ support,       /* [NUM_ATOMS] shmem */
    float*       __restrict__ shmem_reduce,  /* [NUM_WARPS] scratch */
    int tid
) {
    /* Phase 1: each thread accumulates its atoms' contributions in registers.
     * With BLOCK_THREADS=256 and NUM_ATOMS=51, each thread processes at most
     * one source atom (tid 0..50) or zero atoms (tid 51..255). */
    float local_proj[NUM_ATOMS];
    for (int k = 0; k < NUM_ATOMS; k++) local_proj[k] = 0.0f;

    for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
        float z_j  = support[j];
        float t_z  = reward + gamma * z_j * (1.0f - done);
        t_z = fminf(fmaxf(t_z, V_MIN), V_MAX);

        float b_pos = (t_z - V_MIN) / DELTA_Z;
        /* Keep b_pos away from exact upper boundary to prevent phantom
         * upper==lower when ceilf(N-1) == floorf(N-1) = N-1.
         * Without this, probability mass piles on boundary atoms. */
        b_pos = fminf(b_pos, (float)(NUM_ATOMS - 1) - 1e-6f);
        b_pos = fmaxf(b_pos, 1e-6f);
        int lower   = (int)floorf(b_pos);
        int upper   = (int)ceilf(b_pos);
        lower = max(min(lower, NUM_ATOMS - 1), 0);
        upper = max(min(upper, NUM_ATOMS - 1), 0);

        float frac = b_pos - floorf(b_pos);
        float p_j  = target_probs[j];
        local_proj[lower] += p_j * (1.0f - frac);
        if (upper != lower)
            local_proj[upper] += p_j * frac;
    }

    /* Phase 2: reduce each bin across the block.
     * For each atom k, sum local_proj[k] over all threads using the warp
     * butterfly + block-level collect pattern. */
    unsigned int mask = 0xFFFFFFFF;
    int lane = tid & 31;
    int warp = tid >> 5;

    for (int k = 0; k < NUM_ATOMS; k++) {
        /* Warp butterfly sum for atom k */
        float val = local_proj[k];
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(mask, val, offset);

        /* Collect warp results via shmem_reduce */
        if (lane == 0) shmem_reduce[warp] = val;
        __syncthreads();

        if (tid == 0) {
            float total = 0.0f;
            for (int w = 0; w < NUM_WARPS; w++)
                total += shmem_reduce[w];
            projected[k] = total;
        }
        __syncthreads();
    }
    /* projected[k] is now correct for all k — written by thread 0 only. */
}

/* ══════════════════════════════════════════════════════════════════════
 * MAIN KERNEL: c51_loss_batched
 *
 * Takes pre-computed forward pass outputs (value logits + per-branch
 * advantage logits) from three network/state combinations:
 *   - online network on current states  (for current Q-distribution)
 *   - target network on next states     (for target Q-distribution)
 *   - online network on next states     (Double DQN action selection)
 *
 * Grid  = (batch_size, 1, 1)
 * Block = (BLOCK_THREADS, 1, 1) = (256, 1, 1)
 *
 * For each sample b:
 *   Per branch d:
 *     1. Dueling: Q[a,j] = V[j] + A[a,j] - mean_a(A[*,j])
 *     2. log_softmax(Q[a_d, :]) → current action log-probs
 *     3. Double DQN: for each action a, dueling+expected_Q → argmax
 *     4. Dueling for target network + log_softmax for best_next action
 *     5. Bellman projection of target probs onto support
 *     6. Cross-entropy: -sum(projected * current_lp)
 *   Average CE over branches, apply IS weight.
 * ══════════════════════════════════════════════════════════════════════ */

extern "C" __global__ void c51_loss_batched(
    /* ── Online network outputs on current STATES ─────────────────── */
    const float* __restrict__ on_value_logits,    /* [B, NUM_ATOMS] */
    const float* __restrict__ on_adv_logits_b0,   /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
    const float* __restrict__ on_adv_logits_b1,   /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
    const float* __restrict__ on_adv_logits_b2,   /* [B, BRANCH_2_SIZE * NUM_ATOMS] */

    /* ── Target network outputs on NEXT_STATES ────────────────────── */
    const float* __restrict__ tg_value_logits,    /* [B, NUM_ATOMS] */
    const float* __restrict__ tg_adv_logits_b0,   /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
    const float* __restrict__ tg_adv_logits_b1,   /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
    const float* __restrict__ tg_adv_logits_b2,   /* [B, BRANCH_2_SIZE * NUM_ATOMS] */

    /* ── Online network outputs on NEXT_STATES (Double DQN selector) */
    const float* __restrict__ on_next_value_logits,   /* [B, NUM_ATOMS] */
    const float* __restrict__ on_next_adv_logits_b0,  /* [B, BRANCH_0_SIZE * NUM_ATOMS] */
    const float* __restrict__ on_next_adv_logits_b1,  /* [B, BRANCH_1_SIZE * NUM_ATOMS] */
    const float* __restrict__ on_next_adv_logits_b2,  /* [B, BRANCH_2_SIZE * NUM_ATOMS] */

    /* ── Batch data ───────────────────────────────────────────────── */
    const int*   __restrict__ actions,            /* [B] factored action indices 0-44 */
    const float* __restrict__ rewards,            /* [B] */
    const float* __restrict__ dones,              /* [B] */
    const float* __restrict__ is_weights,         /* [B] PER importance-sampling weights */

    /* ── Outputs ──────────────────────────────────────────────────── */
    float* __restrict__ per_sample_loss,          /* [B] IS-weighted loss per sample */
    float* __restrict__ td_errors,               /* [B] unweighted, for PER priority update */
    float* __restrict__ total_loss,              /* [1] batch mean loss (atomicAdd) */

    /* ── Saved tensors for backward pass ─────────────────────────── */
    float* __restrict__ save_current_lp,          /* [B, NUM_BRANCHES, NUM_ATOMS] */
    float* __restrict__ save_projected,           /* [B, NUM_BRANCHES, NUM_ATOMS] */

    /* ── Config ───────────────────────────────────────────────────── */
    float gamma,
    int batch_size
) {
    extern __shared__ float shmem[];

    /* Named pointers into the flat shared memory region */
    float* shmem_support    = shmem + SHMEM_SUPPORT_OFF;    /* [NUM_ATOMS] z_j */
    float* shmem_val        = shmem + SHMEM_VAL_OFF;        /* [NUM_ATOMS] value logits */
    float* shmem_adv        = shmem + SHMEM_ADV_OFF;        /* [MAX_BRANCH_SIZE*NUM_ATOMS] adv logits */
    float* shmem_lp         = shmem + SHMEM_LP_OFF;         /* [NUM_ATOMS] log-prob result */
    float* shmem_proj       = shmem + SHMEM_PROJ_OFF;       /* [NUM_ATOMS] Bellman projection */
    float* shmem_current_lp = shmem + SHMEM_CURRENT_LP_OFF; /* [NUM_ATOMS] taken action lp */
    float* shmem_reduce     = shmem + SHMEM_REDUCE_OFF;     /* [NUM_WARPS] reduction scratch */

    int sample_id = blockIdx.x;
    int tid = threadIdx.x;

    if (sample_id >= batch_size) return;

    /* ── Precompute C51 support atoms once per block ─────────────── */
    for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
        shmem_support[j] = V_MIN + (float)j * DELTA_Z;
    __syncthreads();

    /* ── Decode factored action into per-branch indices ──────────── */
    int factored_action = actions[sample_id];
    /* Clamp to valid range — guards against uninitialized replay buffer entries */
    int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE;
    if (factored_action < 0 || factored_action >= max_action)
        factored_action = 0;
    int a0 = factored_action / (BRANCH_1_SIZE * BRANCH_2_SIZE);
    int a1 = (factored_action / BRANCH_2_SIZE) % BRANCH_1_SIZE;
    int a2 = factored_action % BRANCH_2_SIZE;

    /* Per-branch metadata (compile-time sizes, runtime-selected values) */
    const int branch_action[NUM_BRANCHES] = { a0, a1, a2 };
    const int branch_sizes[NUM_BRANCHES]  = { BRANCH_0_SIZE, BRANCH_1_SIZE, BRANCH_2_SIZE };

    /* Per-sample row pointers into each branch's advantage logit tensors */
    const float* on_adv_row[NUM_BRANCHES] = {
        on_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
        on_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
        on_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
    };
    const float* tg_adv_row[NUM_BRANCHES] = {
        tg_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
        tg_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
        tg_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
    };
    const float* on_next_adv_row[NUM_BRANCHES] = {
        on_next_adv_logits_b0 + (long long)sample_id * BRANCH_0_ATOMS,
        on_next_adv_logits_b1 + (long long)sample_id * BRANCH_1_ATOMS,
        on_next_adv_logits_b2 + (long long)sample_id * BRANCH_2_ATOMS
    };

    /* Value logit row pointers (same NUM_ATOMS for all branches) */
    const float* on_val_row      = on_value_logits      + (long long)sample_id * NUM_ATOMS;
    const float* tg_val_row      = tg_value_logits      + (long long)sample_id * NUM_ATOMS;
    const float* on_next_val_row = on_next_value_logits + (long long)sample_id * NUM_ATOMS;

    float reward    = rewards[sample_id];
    float done      = dones[sample_id];
    float is_weight = is_weights[sample_id];

    float total_ce = 0.0f;

    /* ═══════════════════════════════════════════════════════════════
     * Process each branch sequentially to limit register pressure.
     * ═══════════════════════════════════════════════════════════════ */

    for (int d = 0; d < NUM_BRANCHES; d++) {
        int n_d    = branch_sizes[d];
        int a_d    = branch_action[d];
        int n_atoms = n_d * NUM_ATOMS;  /* total adv logit elements for this branch */

        /* Global memory offset for save tensors: [sample, branch, atom] */
        long long save_off = ((long long)sample_id * NUM_BRANCHES + d) * NUM_ATOMS;

        /* ── Helper: compute dueling logits for a single action into shmem_lp
         *   shmem_val must already hold value logits [NUM_ATOMS]
         *   shmem_adv must already hold all adv logits [n_d * NUM_ATOMS]
         *   Result: dueling(a_sel)[j] = V[j] + A[a_sel,j] - mean_a(A[*,j])
         *           written into shmem_lp[j] */
        /* (Inlined below to avoid device-function overhead for this hot path) */

        /* ═══════════════════════════════════════════════════════════
         * STEP a: Current log-probs for the taken action a_d
         *         (online network, current states)
         * ═══════════════════════════════════════════════════════════ */

        /* Load value logits */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            shmem_val[j] = on_val_row[j];
        __syncthreads();

        /* Load advantage logits for all n_d actions */
        for (int i = tid; i < n_atoms; i += BLOCK_THREADS)
            shmem_adv[i] = on_adv_row[d][i];
        __syncthreads();

        /* Dueling for action a_d: Q[j] = V[j] + A[a_d,j] - mean_a(A[*,j]) */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
            float a_mean = 0.0f;
            for (int a = 0; a < n_d; a++)
                a_mean += shmem_adv[a * NUM_ATOMS + j];
            a_mean /= (float)n_d;
            shmem_lp[j] = shmem_val[j] + shmem_adv[a_d * NUM_ATOMS + j] - a_mean;
        }
        __syncthreads();

        /* log_softmax → result back into shmem_lp */
        block_log_softmax(shmem_lp, shmem_lp, shmem_reduce, tid);
        /* shmem_lp[0..NUM_ATOMS] = log_probs for taken action a_d */

        /* Copy into shmem_current_lp for use in cross-entropy (step e) */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            shmem_current_lp[j] = shmem_lp[j];

        /* Save for backward kernel */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            save_current_lp[save_off + j] = shmem_lp[j];
        __syncthreads();

        /* ═══════════════════════════════════════════════════════════
         * STEP b: Double DQN — argmax action on online_next network
         *
         * For each action a (0..n_d): compute dueling logits, log_softmax,
         * then expected Q. Select action with highest expected Q.
         * ═══════════════════════════════════════════════════════════ */

        /* Load online_next value and advantage logits */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            shmem_val[j] = on_next_val_row[j];
        __syncthreads();

        for (int i = tid; i < n_atoms; i += BLOCK_THREADS)
            shmem_adv[i] = on_next_adv_row[d][i];
        __syncthreads();

        /* Compute average advantage across actions for dueling (same for all a) */
        /* shmem_proj temporarily holds per-atom mean advantage — cleared next step */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
            float a_mean = 0.0f;
            for (int a = 0; a < n_d; a++)
                a_mean += shmem_adv[a * NUM_ATOMS + j];
            a_mean /= (float)n_d;
            shmem_proj[j] = a_mean;  /* reuse proj region for mean advantage */
        }
        __syncthreads();

        /* Now compute expected Q per action and find argmax.
         * We store per-action expected Q in registers on thread 0 (via shared election).
         * The MAX_BRANCH_SIZE actions fit in a small register array on thread 0. */
        __shared__ float   eq_per_action[MAX_BRANCH_SIZE];
        __shared__ int     best_next_a_shared;

        for (int a = 0; a < n_d; a++) {
            /* Dueling logits for action a: V[j] + A[a,j] - mean_j(A[*,j]) */
            for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
                shmem_lp[j] = shmem_val[j] + shmem_adv[a * NUM_ATOMS + j] - shmem_proj[j];
            __syncthreads();

            /* log_softmax in-place on shmem_lp */
            block_log_softmax(shmem_lp, shmem_lp, shmem_reduce, tid);

            /* Expected Q for action a */
            float eq = block_expected_q(shmem_lp, shmem_support, shmem_reduce, tid);
            if (tid == 0) eq_per_action[a] = eq;
            __syncthreads();
        }

        /* Thread 0 finds argmax across actions */
        if (tid == 0) {
            int best = 0;
            for (int a = 1; a < n_d; a++)
                if (eq_per_action[a] > eq_per_action[best]) best = a;
            best_next_a_shared = best;
        }
        __syncthreads();
        int best_next_a = best_next_a_shared;

        /* ═══════════════════════════════════════════════════════════
         * STEP c: Target distribution for best_next action
         *         (target network, next states)
         * ═══════════════════════════════════════════════════════════ */

        /* Load target value and advantage logits */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            shmem_val[j] = tg_val_row[j];
        __syncthreads();

        for (int i = tid; i < n_atoms; i += BLOCK_THREADS)
            shmem_adv[i] = tg_adv_row[d][i];
        __syncthreads();

        /* Per-atom mean advantage for target network */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS) {
            float a_mean = 0.0f;
            for (int a = 0; a < n_d; a++)
                a_mean += shmem_adv[a * NUM_ATOMS + j];
            a_mean /= (float)n_d;
            shmem_lp[j] = shmem_val[j] + shmem_adv[best_next_a * NUM_ATOMS + j] - a_mean;
        }
        __syncthreads();

        /* log_softmax → log-probs for target's best_next action */
        block_log_softmax(shmem_lp, shmem_lp, shmem_reduce, tid);

        /* Convert to probabilities for Bellman projection */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            shmem_proj[j] = expf(shmem_lp[j]);
        __syncthreads();

        /* ═══════════════════════════════════════════════════════════
         * STEP d: Bellman projection
         *         T_z_j = clamp(r + gamma * z_j * (1-done), V_MIN, V_MAX)
         *         Project shmem_proj (target probs) onto shmem_proj (output).
         *
         * The projection writes back to shmem_proj in-place (the input
         * probabilities are read from registers before the output is written).
         * ═══════════════════════════════════════════════════════════ */

        /* Zero-initialize projected output — block_bellman_project accumulates */
        /* We use shmem_lp as the output projected array (shmem_proj still holds probs) */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            shmem_lp[j] = 0.0f;
        __syncthreads();

        block_bellman_project(shmem_proj, reward, done, gamma,
                              shmem_lp, shmem_support, shmem_reduce, tid);
        __syncthreads();

        /* Layer 2: Label smoothing — mix ε uniform into projected target.
         * Prevents any atom from having zero probability → no log(0) in CE. */
        {
            const float ls_uniform = LABEL_SMOOTHING_EPS / (float)NUM_ATOMS;
            const float ls_scale = 1.0f - LABEL_SMOOTHING_EPS;
            for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
                shmem_lp[j] = ls_scale * shmem_lp[j] + ls_uniform;
            __syncthreads();
        }

        /* shmem_lp[0..NUM_ATOMS] = smoothed projected target distribution */

        /* Save projected for backward kernel */
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            save_projected[save_off + j] = shmem_lp[j];
        __syncthreads();

        /* ═══════════════════════════════════════════════════════════
         * STEP e: Cross-entropy loss for this branch
         *         CE = -sum_j(projected[j] * current_lp[j])
         * ═══════════════════════════════════════════════════════════ */

        float local_ce = 0.0f;
        for (int j = tid; j < NUM_ATOMS; j += BLOCK_THREADS)
            local_ce -= shmem_lp[j] * shmem_current_lp[j];

        float branch_ce = block_reduce_sum(local_ce, shmem_reduce, tid);
        /* branch_ce is now the same on all threads */
        total_ce += branch_ce;

    }  /* end branch loop (d = 0, 1, 2) */

    /* ── Aggregate: average over branches, apply IS weight ───────── */
    float avg_ce = total_ce / (float)NUM_BRANCHES;

    /* ── Write outputs (thread 0 only to avoid races) ────────────── */
    if (tid == 0) {
        /* Layer 1: clamp per-sample CE to prevent PER feedback loop.
         * High loss → high PER priority → high IS weight → amplified loss → repeat.
         * Clamping at 50.0 (10x converged loss of ~5) breaks the loop.
         * The gradient itself is bounded (softmax - target ∈ [0,1]),
         * so clamping the LOSS only affects PER priority, not gradient direction. */
        float clamped_ce = fminf(avg_ce, MAX_PER_SAMPLE_CE);
        float weighted_loss = clamped_ce * is_weight;
        per_sample_loss[sample_id] = weighted_loss;
        td_errors[sample_id]       = clamped_ce;  /* PER sees clamped too */
        atomicAdd(total_loss, weighted_loss / (float)batch_size);
    }
}
