/**
 * C51 distributional RL loss gradient kernel — FULLY DETERMINISTIC.
 *
 * Zero atomicAdd: one thread per (b, j), loops over 4 branches.
 * d_value accumulates in register, d_adv written directly (unique slot per thread).
 *
 * dL/d_combined[b,d,j] = is_weights[b] * (exp(current_lp[b,d,j]) - projected[b,d,j])
 * d_value[b,j] = sum_d dL/d_combined[b,d,j]
 * d_adv[b,d,a,j] = dL/d_combined[b,d,j] * (delta(a,a_d) - 1/A_d)
 *
 * Launch config: grid=(ceil(batch_size*num_atoms/256), 1, 1), block=(256, 1, 1).
 */

extern "C" __global__ void c51_grad_kernel(
    const float* __restrict__ current_lp,      // [B, 4, NA]
    const float* __restrict__ projected,        // [B, 4, NA]
    const float* __restrict__ is_weights,                // [B] f32 (bf16 overflows to Inf)
    const int*   __restrict__ actions,          // [B] factored
    float* __restrict__ d_value_logits,         // [B, NA] f32
    float* __restrict__ d_adv_logits,           // [B, (B0+B1+B2+B3)*NA] f32
    int batch_size,
    int num_atoms,
    int b0_size, int b1_size, int b2_size, int b3_size,
    int total_branch_atoms,
    float entropy_coeff,
    const float* __restrict__ branch_scales,    /* [B, 4] per-sample per-branch gradient scale */
    const float* __restrict__ per_sample_support, /* [B, 4, 3] stride-12 per-sample per-branch [v_min, v_max, delta_z] (Phase 2d) */
    const float* __restrict__ atom_positions,    /* [4, num_atoms] adaptive positions. NULL = linear. */
    const float* __restrict__ q_mean_ema_ptr,   /* [1] pinned device-mapped Q-mean EMA */
    /* Task 2.X adaptive magnitude fix + Task 2.Y adaptive direction fix —
     * ISV signal bus [ISV_DIM=23]. Task 2.X slots [13..16] drive an adaptive
     * bin weight applied to the magnitude branch (d==1); Task 2.Y slots
     * [17..21] drive an adaptive bin weight applied to the direction branch
     * (d==0). Slot 22 is SHARPE_EMA, consumed on-GPU by isv_signal_update to
     * drive slot 12 (learning_health). NULL → identity (safe: reduces to
     * pre-fix behaviour). */
    const float* __restrict__ isv_signals,
    /* Backward-standardization fix (2026-04-23): the forward kernel
     * `c51_loss_batched` applies per-atom advantage standardization
     * `centered /= (a_std + 1e-6)` for the magnitude branch (d==1).
     * The backward gradient must replicate the same normalisation factor
     * or the gradient is not the partial derivative of the (scaled) loss.
     * We approximate `a_std` as constant w.r.t. advantage values (same
     * simplification used in batch-norm with `track_running_stats=False`),
     * so the chain rule reduces to `d_adv[a,j] *= 1/(a_std[b,j] + 1e-6)`.
     * `on_adv_logits_b1` layout matches the forward: f32 stride
     * `[B, b1_size, num_atoms]` = `[B, b1_size * num_atoms]` per sample.
     * NULL → standardization not applied (legacy behaviour). */
    const float* __restrict__ on_adv_logits_b1)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    int total_elems = batch_size * num_atoms;
    if (tid >= total_elems) return;

    int j = tid % num_atoms;
    int b = tid / num_atoms;

    float isw = fminf(is_weights[b], 10.0f);
    float inv_batch = 1.0f / (float)batch_size;

    /* Factored action decode (shared across all 4 branches) */
    int factored = actions[b];
    int max_action = b0_size * b1_size * b2_size * b3_size;
    if (factored < 0 || factored >= max_action) factored = 0;
    int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size };
    int branch_actions[4];
    branch_actions[0] = factored / (b1_size * b2_size * b3_size);
    branch_actions[1] = (factored / (b2_size * b3_size)) % b1_size;
    branch_actions[2] = (factored / b3_size) % b2_size;
    branch_actions[3] = factored % b3_size;

    /* Accumulate d_value in register across 4 branches — no atomicAdd */
    float d_val_sum = 0.0f;
    int branch_base = 0;

    for (int d = 0; d < 4; d++) {
        int lp_idx = b * 4 * num_atoms + d * num_atoms + j;
        float lp   = (float)current_lp[lp_idx];
        float proj = (float)projected[lp_idx];

        float d_combined = inv_batch * isw * (expf(lp) - proj);

        /* Entropy regularization — must use inv_batch for consistent scaling */
        if (entropy_coeff > 0.0f) {
            float lp_clamped = fmaxf(lp, -10.0f);
            d_combined += inv_batch * entropy_coeff * (1.0f + lp_clamped);
        }

        /* Task 2.X "make Full useful": scale magnitude-branch gradient by
         * the same signal-driven bin weight applied to branch_ce in
         * c51_loss_batched. Consistency is required — the forward CE and
         * the backward gradient must share the same scalar multiplier so
         * the update step follows the true gradient of the (scaled) loss.
         *
         * Composite signal — see c51_loss_kernel.cu::get_magnitude_bin_weight
         * for the full rationale. Component A reads ISV slots [13..16]
         * (per-mag Q-mean EMAs + |Q| ref) and fires when the taken bin's
         * Q-mean lags max. Component B reads ISV slot [12] (learning_health)
         * and fires when training has not stabilised. collapse_frac =
         * min(1, frac_bin + frac_health). Bounded bin_weight ∈ [1, 2]. */
        if (d == 1 && isv_signals != NULL) {
            int a1 = branch_actions[1];
            if (b1_size > 0 && a1 >= 0 && a1 < b1_size && a1 < 3) {
                /* Component A: bin-specific bias-vs-argmax */
                float q_q = isv_signals[13];
                float q_h = isv_signals[14];
                float q_f = isv_signals[15];
                float q_abs_ref = isv_signals[16];
                float means[3] = { q_q, q_h, q_f };
                float max_mean = means[0];
                for (int k = 1; k < 3; k++)
                    max_mean = fmaxf(max_mean, means[k]);
                float frac_bin = 0.0f;
                if (q_abs_ref > 1e-6f) {
                    float bias_gap = fmaxf(0.0f, max_mean - means[a1]);
                    frac_bin = fminf(1.0f, bias_gap / fmaxf(q_abs_ref, 1e-6f));
                }
                /* Component B: meta-health stress signal */
                float health = fminf(1.0f, fmaxf(0.0f, isv_signals[12]));
                float frac_health = 1.0f - health;
                /* Composite */
                float collapse_frac = fminf(1.0f, frac_bin + frac_health);
                if (collapse_frac > 0.0f) {
                    float mag_bias_signal = (float)(a1 + 1) / (float)b1_size;
                    float bin_weight = 1.0f + collapse_frac * mag_bias_signal;
                    d_combined *= bin_weight;
                }
            }
        }

        /* Task 2.Y "make direction branch useful at eval": scale
         * direction-branch gradient by the signal-driven bin weight applied
         * to branch_ce in c51_loss_batched (d==0). Consistency is required —
         * forward CE and backward gradient must share the same scalar
         * multiplier so the update step follows the true gradient of the
         * (scaled) loss.
         *
         * Mirror of magnitude block above. See c51_loss_kernel.cu::
         * get_direction_bin_weight for rationale. Component A reads ISV
         * slots [17..21] (per-direction Q-mean EMAs + |Q| ref) and fires
         * when the taken direction's Q-mean lags max. Component B shares
         * ISV slot [12] (learning_health) with the magnitude mechanism.
         * dir_bias_signal[k] = {1.0, 0.5, 1.0, 0.0} for Short/Hold/Long/Flat:
         * Flat always returns bin_weight=1.0 (no amplification of the bin
         * the mechanism is meant to correct AGAINST). Bounded ∈ [1, 2]. */
        if (d == 0 && isv_signals != NULL) {
            int a0 = branch_actions[0];
            if (b0_size > 0 && a0 >= 0 && a0 < b0_size && a0 < 4) {
                /* Component A: bin-specific bias-vs-argmax (direction) */
                float q_short = isv_signals[17];
                float q_hold  = isv_signals[18];
                float q_long  = isv_signals[19];
                float q_flat  = isv_signals[20];
                float q_dir_abs_ref = isv_signals[21];
                float means[4] = { q_short, q_hold, q_long, q_flat };
                float max_mean = means[0];
                for (int k = 1; k < 4; k++)
                    max_mean = fmaxf(max_mean, means[k]);
                float frac_bin = 0.0f;
                if (q_dir_abs_ref > 1e-6f) {
                    float bias_gap = fmaxf(0.0f, max_mean - means[a0]);
                    frac_bin = fminf(1.0f, bias_gap / fmaxf(q_dir_abs_ref, 1e-6f));
                }
                /* Component B: shared meta-health stress signal */
                float health = fminf(1.0f, fmaxf(0.0f, isv_signals[12]));
                float frac_health = 1.0f - health;
                /* Composite */
                float collapse_frac = fminf(1.0f, frac_bin + frac_health);
                if (collapse_frac > 0.0f) {
                    /* Architectural shape — NOT a tuning knob. Same
                     * rationale as get_direction_bin_weight in c51_loss_kernel. */
                    float dir_bias_signal = 0.0f;
                    if (a0 == 0) dir_bias_signal = 1.0f;       /* Short */
                    else if (a0 == 1) dir_bias_signal = 0.5f;  /* Hold */
                    else if (a0 == 2) dir_bias_signal = 1.0f;  /* Long */
                    else if (a0 == 3) dir_bias_signal = 0.0f;  /* Flat */
                    if (dir_bias_signal > 0.0f) {
                        /* Cross-branch compression boost — mirror of the
                         * c51_loss_kernel get_direction_bin_weight formulation.
                         * When direction |Q|-scale is much tighter than
                         * magnitude |Q|-scale, amplify more aggressively.
                         * ISV-driven; self-disables when spreads equalise. */
                        float q_abs_ref_mag = isv_signals[16];
                        float compression = 0.0f;
                        if (q_abs_ref_mag > 1e-6f) {
                            compression = fmaxf(0.0f,
                                1.0f - q_dir_abs_ref / fmaxf(q_abs_ref_mag, 1e-6f));
                            compression = fminf(1.0f, compression);
                        }
                        float compression_boost = 1.0f + compression; /* ∈ [1,2] */
                        float bin_weight = 1.0f
                            + collapse_frac * dir_bias_signal * compression_boost;
                        d_combined *= bin_weight;
                    }
                }
            }
        }

        d_val_sum += d_combined;

        float branch_scale = branch_scales[b * 4 + d];

        int A_d = branch_sizes[d];
        float inv_A = 1.0f / (float)A_d;
        int a_d = branch_actions[d];

        /* d_adv: each (b,d,a,j) slot is written by exactly ONE thread — plain write */

        /* Q-gap floor gradient: pushes taken action's distribution toward
         * higher atoms and non-taken toward lower. Creates Q-value spread
         * even when the C51 Bellman has converged (zero CE gradient).
         * Scale: proportional to atom position (normalized to [-1, +1]).
         * Weight: inv_batch * 0.01 (small relative to main gradient).
         * The gradient is orthogonal to Bellman — it depends on atom
         * position, not on the target match. Active on ALL samples,
         * providing perpetual differentiation pressure. */
        /* z_norm: use adaptive atom positions if available, else linear grid.
         * Phase 2d: per_sample_support is [B, 4, 3] stride-12; index with
         * (b, d, {0,1,2}) so each branch's normalisation uses its own
         * v_min/v_max/delta_z. */
        float z_norm;
        float delta_z;
        int support_base = b * 12 + d * 3;
        if (atom_positions != NULL) {
            float z_val = atom_positions[(long long)d * num_atoms + j];
            float z_min = per_sample_support[support_base + 0];
            float z_max = per_sample_support[support_base + 1];
            z_norm = 2.0f * (z_val - z_min) / fmaxf(z_max - z_min, 1e-6f) - 1.0f;
            /* Approximate delta_z from adjacent atom spacing */
            delta_z = (j < num_atoms - 1)
                ? atom_positions[(long long)d * num_atoms + j + 1] - z_val
                : z_val - atom_positions[(long long)d * num_atoms + j - 1];
        } else {
            z_norm = 2.0f * (float)j / fmaxf((float)(num_atoms - 1), 1.0f) - 1.0f;
            delta_z = per_sample_support[support_base + 2];
        }
        float spread_scale = inv_batch * delta_z;

        /* Asymmetric spread: challenger (adjacent to taken) gets pushed UP.
         * Creates a two-horse race instead of one-horse monopoly.
         * Taken → UP (full), challenger → UP (half), rest → DOWN.
         * Zero-sum: rest absorbs -(1.0 + 0.5) / (A_d - 2) to prevent Q-inflation.
         * For A_d=3: taken=+1.0, challenger=+0.5, other=-1.5. */
        int challenger = (a_d + 1) % A_d;
        int n_rest = A_d - 2;  /* actions that are neither taken nor challenger */
        float rest_weight = (n_rest > 0) ? (1.0f + 0.5f) / (float)n_rest : 0.0f;

        /* Backward-standardization fix (2026-04-23): for d==1 (magnitude)
         * the forward applies `centered /= (a_std + 1e-6f)` per (b, j).
         * Recompute the same a_std here so the chain-rule factor lands
         * on every d==1 grad write. `on_adv_logits_b1` layout is
         * `[B, b1_size, num_atoms]`; we read the 3 advantage values at
         * this (b, *, j) slice, compute mean + variance, then sqrt.
         * Approximates a_std as constant w.r.t. advantage — the simpler
         * chain rule used in practice (skips da_std/dadv terms). Guarded
         * on `on_adv_logits_b1 != NULL` so the kernel is backward-
         * compatible with smoke-test launchers that don't wire it. */
        float inv_a_std = 1.0f;
        if (d == 1 && on_adv_logits_b1 != NULL) {
            int A1 = A_d;
            /* Only compute when b1_size looks sane (avoids OOB on
             * degenerate configurations). */
            if (A1 > 0) {
                float sum = 0.0f;
                for (int a = 0; a < A1; a++) {
                    sum += on_adv_logits_b1[(long long)b * A1 * num_atoms + a * num_atoms + j];
                }
                float mean = sum / (float)A1;
                float sq_sum = 0.0f;
                for (int a = 0; a < A1; a++) {
                    float v = on_adv_logits_b1[(long long)b * A1 * num_atoms + a * num_atoms + j] - mean;
                    sq_sum += v * v;
                }
                float a_std = sqrtf(sq_sum / (float)A1 + 1e-12f);
                /* SP14 A.1: lift floor 1e-6 → 1e-3 to cap inv_a_std at 1000
                 * (was ~1e6 when magnitude advantage logits collapse near-
                 * uniform — Smoke A captured var_q≈9e-10 producing 1109
                 * GRAD_CLIP_OUTLIER events with C51 grad reaching 9.5e6).
                 * Numerical-stability anchor (Invariant 1: prevent division-
                 * by-near-zero amplification), not an ISV-driven behavioural
                 * bound; same class of fix as the existing 1e-12f a_std floor
                 * above. */
                inv_a_std = 1.0f / (a_std + 1e-3f);
            }
        }

        for (int a = 0; a < A_d; a++) {
            /* SP17 annotation (Commit E, 2026-05-08): this `dueling_grad` IS
             * the SP17 mean-zero Jacobian — `J[a, a_d] = δ(a, a_d) − 1/n_d`
             * (where `inv_A = 1.0 / A_d`). For the action `a` that was taken
             * in the forward pass, the per-action centered logit is
             * `A_taken − mean_a A`; the chain rule gradient w.r.t. logit at
             * action a' is `∂(A_taken − mean_a A)/∂A[a']` = δ(a', taken) −
             * 1/n_d. This site has carried the correct mean-zero Jacobian
             * since 56373f094 — no code-path change in this commit.
             *
             * SP17: this IS the mean-zero Jacobian — Δ(a,a_d) - 1/n_d. */
            float dueling_grad = (a == a_d) ? (1.0f - inv_A) : (-inv_A);
            float grad_val = branch_scale * d_combined * dueling_grad;
            if (d == 1) grad_val *= inv_a_std;   /* SP17: per-d=1 mag std normalization (orthogonal) */

            float spread_grad;
            if (a == a_d) {
                spread_grad = spread_scale * z_norm;
            } else if (a == challenger) {
                spread_grad = spread_scale * z_norm * 0.5f;
            } else {
                spread_grad = -spread_scale * z_norm * rest_weight;
            }
            grad_val += branch_scale * spread_grad;

            int adv_idx = branch_base + b * (A_d * num_atoms) + a * num_atoms + j;
            d_adv_logits[adv_idx] = grad_val;
        }
        branch_base += batch_size * A_d * num_atoms;
    }

    /* Q-mean drift: handled adaptively by E1 enrichment (Q-value reality check).
     * Hardcoded penalty removed — E1 computes bias correction from actual eval
     * performance and adapts to changing reward scales automatically. */
    (void)q_mean_ema_ptr; /* parameter kept for ABI compatibility */

    /* Single deterministic write — no atomicAdd */
    d_value_logits[b * num_atoms + j] = d_val_sum;
}
