/**
 * IQN (Implicit Quantile Network) Dual-Head Kernel
 *
 * Auxiliary quantile head that shares the DQN trunk (h_s2) but trains its own
 * weights independently. Provides quantile-based Q-values for risk-sensitive
 * action selection (CVaR) alongside the primary C51 distributional head.
 *
 * Architecture per sample:
 *   For each τ_i ∈ (0,1), i = 1..N:
 *     cos_feat  = [cos(π·1·τ), ..., cos(π·D·τ)]          [IQN_EMBED_DIM]
 *     embed     = ReLU(W_embed × cos_feat + b_embed)      [IQN_HIDDEN]
 *     combined  = h_s2 ⊙ embed                             [IQN_HIDDEN]
 *     per-branch: q_d = W_bd × combined + b_bd             [n_d] scalars
 *
 * Loss: Quantile Huber loss (Dabney et al., 2018b)
 *   For each predicted τ_i and target τ'_j:
 *     δ_ij = target_θ_j - predicted_θ_i
 *     ρ_τi(δ) = |τ_i - 𝟙{δ<0}| × L_κ(δ)
 *   L_iqn = (1/N) Σ_i (1/N') Σ_j ρ_τi(δ_ij)
 *
 * Launch config:
 *   Forward + loss:  grid=(batch_size, 1, 1), block=(32, 1, 1)
 *   Backward:        grid=(batch_size, 1, 1), block=(32, 1, 1)
 *   Grad norm:       grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1)
 *   Adam:            grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1)
 *
 * Compile-time defines (injected via NVRTC dim_overrides):
 *   IQN_HIDDEN        -- hidden dim = SHARED_H2 from DQN trunk (default 256)
 *   IQN_EMBED_DIM     -- cosine embedding dimension (default 64)
 *   IQN_NUM_QUANTILES -- τ samples per forward pass (default 32)
 *   IQN_KAPPA         -- Huber threshold for quantile loss (default 1.0)
 *   BRANCH_0_SIZE     -- exposure actions (default 5)
 *   BRANCH_1_SIZE     -- order type actions (default 3)
 *   BRANCH_2_SIZE     -- urgency actions (default 3)
 */

/* ── Compile-time defaults (overridden by NVRTC injection) ──────────── */
#ifndef IQN_HIDDEN
#define IQN_HIDDEN 256
#endif
#ifndef IQN_EMBED_DIM
#define IQN_EMBED_DIM 64
#endif
#ifndef IQN_NUM_QUANTILES
#define IQN_NUM_QUANTILES 32
#endif
#ifndef IQN_KAPPA
#define IQN_KAPPA 1.0f
#endif
#ifndef BRANCH_0_SIZE
#define BRANCH_0_SIZE 5
#endif
#ifndef BRANCH_1_SIZE
#define BRANCH_1_SIZE 3
#endif
#ifndef BRANCH_2_SIZE
#define BRANCH_2_SIZE 3
#endif

#define TOTAL_BRANCH_ACTIONS (BRANCH_0_SIZE + BRANCH_1_SIZE + BRANCH_2_SIZE)

/* ── Weight layout (flat f32 buffer) ─────────────────────────────────── */
/* W_embed [IQN_HIDDEN, IQN_EMBED_DIM], b_embed [IQN_HIDDEN]
 * W_b0 [BRANCH_0_SIZE, IQN_HIDDEN], b_b0 [BRANCH_0_SIZE]
 * W_b1 [BRANCH_1_SIZE, IQN_HIDDEN], b_b1 [BRANCH_1_SIZE]
 * W_b2 [BRANCH_2_SIZE, IQN_HIDDEN], b_b2 [BRANCH_2_SIZE] */

#define IQN_W_EMBED_SIZE  (IQN_HIDDEN * IQN_EMBED_DIM)
#define IQN_B_EMBED_SIZE  (IQN_HIDDEN)
#define IQN_W_B0_SIZE     (BRANCH_0_SIZE * IQN_HIDDEN)
#define IQN_B_B0_SIZE     (BRANCH_0_SIZE)
#define IQN_W_B1_SIZE     (BRANCH_1_SIZE * IQN_HIDDEN)
#define IQN_B_B1_SIZE     (BRANCH_1_SIZE)
#define IQN_W_B2_SIZE     (BRANCH_2_SIZE * IQN_HIDDEN)
#define IQN_B_B2_SIZE     (BRANCH_2_SIZE)

#define IQN_TOTAL_PARAMS (IQN_W_EMBED_SIZE + IQN_B_EMBED_SIZE \
    + IQN_W_B0_SIZE + IQN_B_B0_SIZE \
    + IQN_W_B1_SIZE + IQN_B_B1_SIZE \
    + IQN_W_B2_SIZE + IQN_B_B2_SIZE)

/* Offsets into the flat parameter buffer */
#define IQN_OFF_W_EMBED  0
#define IQN_OFF_B_EMBED  (IQN_OFF_W_EMBED + IQN_W_EMBED_SIZE)
#define IQN_OFF_W_B0     (IQN_OFF_B_EMBED + IQN_B_EMBED_SIZE)
#define IQN_OFF_B_B0     (IQN_OFF_W_B0 + IQN_W_B0_SIZE)
#define IQN_OFF_W_B1     (IQN_OFF_B_B0 + IQN_B_B0_SIZE)
#define IQN_OFF_B_B1     (IQN_OFF_W_B1 + IQN_W_B1_SIZE)
#define IQN_OFF_W_B2     (IQN_OFF_B_B1 + IQN_B_B1_SIZE)
#define IQN_OFF_B_B2     (IQN_OFF_W_B2 + IQN_W_B2_SIZE)

/* Distributed array size: ceil(dim / 32) elements per lane */
#define IQN_DIST(dim) (((dim) + 31) / 32)

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

/** Warp-reduce sum across all 32 lanes (butterfly). */
__device__ __forceinline__ float iqn_warp_sum(float val) {
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
    return val;
}

/** Quantile Huber loss element: ρ_τ(δ) = |τ - 𝟙{δ<0}| × L_κ(δ) */
__device__ __forceinline__ float quantile_huber_element(float tau, float delta, float kappa) {
    float abs_delta = fabsf(delta);
    /* L_κ(δ) = 0.5δ² if |δ|≤κ, κ(|δ| - 0.5κ) otherwise */
    float huber = (abs_delta <= kappa)
        ? 0.5f * delta * delta
        : kappa * (abs_delta - 0.5f * kappa);
    /* Asymmetric weight: |τ - 𝟙{δ<0}| */
    float indicator = (delta < 0.0f) ? 1.0f : 0.0f;
    float weight = fabsf(tau - indicator);
    return weight * huber;
}

/** Gradient of quantile Huber loss element w.r.t. predicted value.
 *  dρ_τ/d(pred) = -|τ - 𝟙{δ<0}| × dL_κ/dδ
 *  where dL_κ/dδ = δ if |δ|≤κ, κ·sign(δ) otherwise */
__device__ __forceinline__ float quantile_huber_grad(float tau, float delta, float kappa) {
    float abs_delta = fabsf(delta);
    float huber_grad = (abs_delta <= kappa)
        ? delta
        : kappa * ((delta > 0.0f) ? 1.0f : -1.0f);
    float indicator = (delta < 0.0f) ? 1.0f : 0.0f;
    float weight = fabsf(tau - indicator);
    /* Gradient w.r.t. predicted = -weight × huber_grad (negative because δ = target - pred) */
    return -weight * huber_grad;
}

/* ── Forward + Loss Kernel ──────────────────────────────────────────── */
/**
 * Per-sample IQN forward pass + quantile Huber loss.
 *
 * Reads h_s2 (shared trunk activation), computes quantile values for all
 * branches and τ samples, then computes quantile Huber loss against target
 * quantile values (from target network forward through same IQN head).
 *
 * τ values are pre-sampled and uploaded via the `taus` buffer.
 *
 * One warp (32 threads) per sample. Quantiles are processed sequentially
 * to conserve registers — each τ produces TOTAL_BRANCH_ACTIONS values.
 *
 * Saves per-quantile intermediate activations for backward pass.
 */
extern "C" __global__
void iqn_forward_loss_kernel(
    /* Inputs */
    const float* __restrict__ h_s2,           /* [B, IQN_HIDDEN] trunk activations */
    const float* __restrict__ target_h_s2,    /* [B, IQN_HIDDEN] target trunk activations */
    const float* __restrict__ taus,           /* [B, N] online τ samples ∈ (0,1) */
    const float* __restrict__ target_taus,    /* [B, N] target τ samples */
    const int*   __restrict__ actions,        /* [B, 3] branch actions (exposure, order, urgency) */
    const float* __restrict__ rewards,        /* [B] */
    const float* __restrict__ dones,          /* [B] */
    float gamma,
    /* Weights */
    const float* __restrict__ online_params,  /* [IQN_TOTAL_PARAMS] online IQN weights */
    const float* __restrict__ target_params,  /* [IQN_TOTAL_PARAMS] target IQN weights */
    /* Outputs */
    float* __restrict__ per_sample_loss,      /* [B] per-sample loss (for PER) */
    float* __restrict__ total_loss,           /* [1] batch-mean loss (atomicAdd) */
    /* Activation saves for backward */
    float* __restrict__ save_embed,           /* [B, N, IQN_HIDDEN] post-ReLU embeddings */
    float* __restrict__ save_combined,        /* [B, N, IQN_HIDDEN] h_s2 ⊙ embed */
    float* __restrict__ save_q_online,        /* [B, N, TOTAL_BRANCH_ACTIONS] quantile Q-values */
    float* __restrict__ save_q_target,        /* [B, N, TOTAL_BRANCH_ACTIONS] target Q-values */
    int batch_size
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;
    int lane = threadIdx.x;

    /* ── Pointers into this sample's data ── */
    const float* my_h_s2 = h_s2 + sample * IQN_HIDDEN;
    const float* my_target_h_s2 = target_h_s2 + sample * IQN_HIDDEN;
    const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
    const float* my_target_taus = target_taus + sample * IQN_NUM_QUANTILES;

    int a0 = actions[sample * 3 + 0];  /* exposure action */
    int a1 = actions[sample * 3 + 1];  /* order action */
    int a2 = actions[sample * 3 + 2];  /* urgency action */
    float reward = rewards[sample];
    float done = dones[sample];

    /* Weight pointers (same for all samples) */
    const float* w_embed = online_params + IQN_OFF_W_EMBED;
    const float* b_embed = online_params + IQN_OFF_B_EMBED;
    const float* w_b0 = online_params + IQN_OFF_W_B0;
    const float* b_b0 = online_params + IQN_OFF_B_B0;
    const float* w_b1 = online_params + IQN_OFF_W_B1;
    const float* b_b1 = online_params + IQN_OFF_B_B1;
    const float* w_b2 = online_params + IQN_OFF_W_B2;
    const float* b_b2 = online_params + IQN_OFF_B_B2;

    const float* tw_embed = target_params + IQN_OFF_W_EMBED;
    const float* tb_embed = target_params + IQN_OFF_B_EMBED;
    const float* tw_b0 = target_params + IQN_OFF_W_B0;
    const float* tb_b0 = target_params + IQN_OFF_B_B0;
    const float* tw_b1 = target_params + IQN_OFF_W_B1;
    const float* tb_b1 = target_params + IQN_OFF_B_B1;
    const float* tw_b2 = target_params + IQN_OFF_W_B2;
    const float* tb_b2 = target_params + IQN_OFF_B_B2;

    /* ── Load h_s2 into distributed registers ── */
    float h_dist[IQN_DIST(IQN_HIDDEN)];
    for (int d = lane; d < IQN_HIDDEN; d += 32)
        h_dist[d / 32] = my_h_s2[d];

    float h_target_dist[IQN_DIST(IQN_HIDDEN)];
    for (int d = lane; d < IQN_HIDDEN; d += 32)
        h_target_dist[d / 32] = my_target_h_s2[d];

    /* ── Process each ONLINE quantile ── */
    /* For each τ_i, compute quantile Q-values for the taken actions */
    float online_q_a[IQN_NUM_QUANTILES]; /* Q(s, a_taken, τ_i) per branch sum */

    for (int t = 0; t < IQN_NUM_QUANTILES; t++) {
        float tau = my_taus[t];

        /* 1. Cosine features: cos(π·(d+1)·τ) for d=0..IQN_EMBED_DIM-1 */
        /* 2. Linear embedding: embed[h] = Σ_d W_embed[h,d]·cos_feat[d] + b_embed[h] */
        /* Process via warp-cooperative matvec: each lane handles IQN_HIDDEN/32 outputs */
        float embed_dist[IQN_DIST(IQN_HIDDEN)];
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            float acc = b_embed[h];
            const float* w_row = w_embed + h * IQN_EMBED_DIM;
            for (int d = 0; d < IQN_EMBED_DIM; d++) {
                float cos_val = cosf(3.14159265f * (float)(d + 1) * tau);
                acc += w_row[d] * cos_val;
            }
            /* ReLU */
            embed_dist[h / 32] = fmaxf(acc, 0.0f);
        }

        /* 3. Element-wise product: combined = h_s2 ⊙ embed */
        float comb_dist[IQN_DIST(IQN_HIDDEN)];
        for (int h = lane; h < IQN_HIDDEN; h += 32)
            comb_dist[h / 32] = h_dist[h / 32] * embed_dist[h / 32];

        /* Save activations for backward pass */
        int save_offset = (sample * IQN_NUM_QUANTILES + t) * IQN_HIDDEN;
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            save_embed[save_offset + h] = embed_dist[h / 32];
            save_combined[save_offset + h] = comb_dist[h / 32];
        }

        /* 4. Per-branch output: q_d = W_bd × combined + b_bd */
        int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS;

        /* Branch 0 (exposure) */
        float q0_taken = 0.0f;
        for (int a = 0; a < BRANCH_0_SIZE; a++) {
            float acc = b_b0[a];
            const float* w_row = w_b0 + a * IQN_HIDDEN;
            float partial = 0.0f;
            for (int h = lane; h < IQN_HIDDEN; h += 32)
                partial += w_row[h] * comb_dist[h / 32];
            acc += iqn_warp_sum(partial);
            if (lane == 0) {
                save_q_online[q_save_offset + a] = acc;
                if (a == a0) q0_taken = acc;
            }
            if (a == a0) q0_taken = __shfl_sync(0xFFFFFFFF, q0_taken, 0);
        }

        /* Branch 1 (order) */
        float q1_taken = 0.0f;
        for (int a = 0; a < BRANCH_1_SIZE; a++) {
            float acc = b_b1[a];
            const float* w_row = w_b1 + a * IQN_HIDDEN;
            float partial = 0.0f;
            for (int h = lane; h < IQN_HIDDEN; h += 32)
                partial += w_row[h] * comb_dist[h / 32];
            acc += iqn_warp_sum(partial);
            if (lane == 0) {
                save_q_online[q_save_offset + BRANCH_0_SIZE + a] = acc;
                if (a == a1) q1_taken = acc;
            }
            if (a == a1) q1_taken = __shfl_sync(0xFFFFFFFF, q1_taken, 0);
        }

        /* Branch 2 (urgency) */
        float q2_taken = 0.0f;
        for (int a = 0; a < BRANCH_2_SIZE; a++) {
            float acc = b_b2[a];
            const float* w_row = w_b2 + a * IQN_HIDDEN;
            float partial = 0.0f;
            for (int h = lane; h < IQN_HIDDEN; h += 32)
                partial += w_row[h] * comb_dist[h / 32];
            acc += iqn_warp_sum(partial);
            if (lane == 0) {
                save_q_online[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + a] = acc;
                if (a == a2) q2_taken = acc;
            }
            if (a == a2) q2_taken = __shfl_sync(0xFFFFFFFF, q2_taken, 0);
        }

        /* Sum of taken-action Q-values across branches (for loss computation) */
        online_q_a[t] = q0_taken + q1_taken + q2_taken;
    }

    /* ── Process TARGET quantiles ── */
    float target_q_a[IQN_NUM_QUANTILES];

    for (int t = 0; t < IQN_NUM_QUANTILES; t++) {
        float tau = my_target_taus[t];

        /* Cosine embedding + linear (target weights) */
        float embed_dist[IQN_DIST(IQN_HIDDEN)];
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            float acc = tb_embed[h];
            const float* w_row = tw_embed + h * IQN_EMBED_DIM;
            for (int d = 0; d < IQN_EMBED_DIM; d++) {
                float cos_val = cosf(3.14159265f * (float)(d + 1) * tau);
                acc += w_row[d] * cos_val;
            }
            embed_dist[h / 32] = fmaxf(acc, 0.0f);
        }

        /* Element-wise product with target h_s2 */
        float comb_dist[IQN_DIST(IQN_HIDDEN)];
        for (int h = lane; h < IQN_HIDDEN; h += 32)
            comb_dist[h / 32] = h_target_dist[h / 32] * embed_dist[h / 32];

        /* Per-branch Q-values for taken actions */
        float q0 = 0.0f, q1 = 0.0f, q2 = 0.0f;

        /* Branch 0 */
        {
            float acc = tb_b0[a0];
            float partial = 0.0f;
            const float* w_row = tw_b0 + a0 * IQN_HIDDEN;
            for (int h = lane; h < IQN_HIDDEN; h += 32)
                partial += w_row[h] * comb_dist[h / 32];
            q0 = acc + iqn_warp_sum(partial);
        }
        /* Branch 1 */
        {
            float acc = tb_b1[a1];
            float partial = 0.0f;
            const float* w_row = tw_b1 + a1 * IQN_HIDDEN;
            for (int h = lane; h < IQN_HIDDEN; h += 32)
                partial += w_row[h] * comb_dist[h / 32];
            q1 = acc + iqn_warp_sum(partial);
        }
        /* Branch 2 */
        {
            float acc = tb_b2[a2];
            float partial = 0.0f;
            const float* w_row = tw_b2 + a2 * IQN_HIDDEN;
            for (int h = lane; h < IQN_HIDDEN; h += 32)
                partial += w_row[h] * comb_dist[h / 32];
            q2 = acc + iqn_warp_sum(partial);
        }

        /* Bellman target: r + γ(1-done) × Q_target */
        float raw_target = q0 + q1 + q2;
        target_q_a[t] = reward + gamma * (1.0f - done) * raw_target;

        /* Save target Q for backward */
        int q_save_offset = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS;
        if (lane == 0) {
            save_q_target[q_save_offset + a0] = target_q_a[t];
            save_q_target[q_save_offset + BRANCH_0_SIZE + a1] = target_q_a[t];
            save_q_target[q_save_offset + BRANCH_0_SIZE + BRANCH_1_SIZE + a2] = target_q_a[t];
        }
    }

    /* ── Quantile Huber Loss: (1/N) Σ_i (1/N') Σ_j ρ_τi(δ_ij) ── */
    float sample_loss = 0.0f;
    if (lane == 0) {
        for (int i = 0; i < IQN_NUM_QUANTILES; i++) {
            float tau_i = my_taus[i];
            float inner_sum = 0.0f;
            for (int j = 0; j < IQN_NUM_QUANTILES; j++) {
                float delta = target_q_a[j] - online_q_a[i];
                inner_sum += quantile_huber_element(tau_i, delta, IQN_KAPPA);
            }
            sample_loss += inner_sum / (float)IQN_NUM_QUANTILES;
        }
        sample_loss /= (float)IQN_NUM_QUANTILES;

        per_sample_loss[sample] = sample_loss;
        atomicAdd(total_loss, sample_loss / (float)batch_size);
    }
}

/* ── Backward Kernel ────────────────────────────────────────────────── */
/**
 * Backprop through IQN head. Computes gradients w.r.t. IQN weights only
 * (does NOT backprop into h_s2 — the shared trunk is trained by C51).
 *
 * For each sample, for each quantile τ_i:
 *   dL/dq_online_i = (1/N²) Σ_j dρ_τi(δ_ij)/dq_i
 *   dL/d(combined) = W_bd^T × dL/dq_d (for taken action a_d)
 *   dL/d(embed) = dL/d(combined) ⊙ h_s2 ⊙ 𝟙{pre_relu > 0}
 *   dL/d(W_embed) += outer(dL/d(embed), cos_features)
 *   dL/d(W_bd) += outer(dL/dq_d, combined)
 */
extern "C" __global__
void iqn_backward_kernel(
    /* Saved activations (from forward pass) */
    const float* __restrict__ h_s2,           /* [B, IQN_HIDDEN] */
    const float* __restrict__ save_embed,     /* [B, N, IQN_HIDDEN] */
    const float* __restrict__ save_combined,  /* [B, N, IQN_HIDDEN] */
    const float* __restrict__ save_q_online,  /* [B, N, TOTAL_BRANCH_ACTIONS] */
    const float* __restrict__ save_q_target,  /* [B, N, TOTAL_BRANCH_ACTIONS] */
    const float* __restrict__ taus,           /* [B, N] */
    const int*   __restrict__ actions,        /* [B, 3] */
    /* Weights (for backward through linear layers) */
    const float* __restrict__ online_params,
    /* Gradient output */
    float* __restrict__ grad_buf,             /* [IQN_TOTAL_PARAMS] accumulated gradients */
    float* __restrict__ d_h_s2_out,           /* [B, IQN_HIDDEN] dL/d(h_s2) for trunk gradient */
    int batch_size
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;
    int lane = threadIdx.x;

    const float* my_h_s2 = h_s2 + sample * IQN_HIDDEN;
    const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
    int a0 = actions[sample * 3 + 0];
    int a1 = actions[sample * 3 + 1];
    int a2 = actions[sample * 3 + 2];

    const float* w_embed = online_params + IQN_OFF_W_EMBED;
    const float* w_b0 = online_params + IQN_OFF_W_B0;
    const float* w_b1 = online_params + IQN_OFF_W_B1;
    const float* w_b2 = online_params + IQN_OFF_W_B2;

    /* Load h_s2 into distributed registers */
    float h_dist[IQN_DIST(IQN_HIDDEN)];
    for (int h = lane; h < IQN_HIDDEN; h += 32)
        h_dist[h / 32] = my_h_s2[h];

    /* Pre-compute online and target Q-values for taken actions across all quantiles */
    /* (reading from saved buffers) */
    float online_q[IQN_NUM_QUANTILES];
    float target_q[IQN_NUM_QUANTILES];
    for (int t = 0; t < IQN_NUM_QUANTILES; t++) {
        int qoff = (sample * IQN_NUM_QUANTILES + t) * TOTAL_BRANCH_ACTIONS;
        if (lane == 0) {
            online_q[t] = save_q_online[qoff + a0]
                        + save_q_online[qoff + BRANCH_0_SIZE + a1]
                        + save_q_online[qoff + BRANCH_0_SIZE + BRANCH_1_SIZE + a2];
            target_q[t] = save_q_target[qoff + a0]; /* target_q already contains Bellman target */
        }
        online_q[t] = __shfl_sync(0xFFFFFFFF, online_q[t], 0);
        target_q[t] = __shfl_sync(0xFFFFFFFF, target_q[t], 0);
    }

    /* Process each quantile τ_i */
    float inv_n_sq = 1.0f / ((float)IQN_NUM_QUANTILES * (float)IQN_NUM_QUANTILES);

    for (int ti = 0; ti < IQN_NUM_QUANTILES; ti++) {
        float tau_i = my_taus[ti];
        int save_offset = (sample * IQN_NUM_QUANTILES + ti) * IQN_HIDDEN;

        /* Compute dL/dq_online for this quantile:
         * dL/dq_i = (1/N²) Σ_j dρ_τi(δ_ij)/dq_i */
        float dL_dq = 0.0f;
        if (lane == 0) {
            for (int tj = 0; tj < IQN_NUM_QUANTILES; tj++) {
                float delta = target_q[tj] - online_q[ti];
                dL_dq += quantile_huber_grad(tau_i, delta, IQN_KAPPA);
            }
            dL_dq *= inv_n_sq;
        }
        dL_dq = __shfl_sync(0xFFFFFFFF, dL_dq, 0);

        /* Load saved activations */
        float embed_dist[IQN_DIST(IQN_HIDDEN)];
        float comb_dist[IQN_DIST(IQN_HIDDEN)];
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            embed_dist[h / 32] = save_embed[save_offset + h];
            comb_dist[h / 32] = save_combined[save_offset + h];
        }

        /* ── Gradient through branch output layers (taken action only) ── */
        /* dL/d(combined) = Σ_d W_bd[a_d, :] × dL/dq (only for taken actions) */
        float dL_dcomb[IQN_DIST(IQN_HIDDEN)];
        for (int h = lane; h < IQN_HIDDEN; h += 32)
            dL_dcomb[h / 32] = 0.0f;

        /* Branch 0: accumulate gradient */
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            dL_dcomb[h / 32] += w_b0[a0 * IQN_HIDDEN + h] * dL_dq;
            /* dL/dW_b0[a0, h] += dL/dq × combined[h] */
            atomicAdd(&grad_buf[IQN_OFF_W_B0 + a0 * IQN_HIDDEN + h],
                      dL_dq * comb_dist[h / 32]);
        }
        if (lane == 0)
            atomicAdd(&grad_buf[IQN_OFF_B_B0 + a0], dL_dq);

        /* Branch 1 */
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            dL_dcomb[h / 32] += w_b1[a1 * IQN_HIDDEN + h] * dL_dq;
            atomicAdd(&grad_buf[IQN_OFF_W_B1 + a1 * IQN_HIDDEN + h],
                      dL_dq * comb_dist[h / 32]);
        }
        if (lane == 0)
            atomicAdd(&grad_buf[IQN_OFF_B_B1 + a1], dL_dq);

        /* Branch 2 */
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            dL_dcomb[h / 32] += w_b2[a2 * IQN_HIDDEN + h] * dL_dq;
            atomicAdd(&grad_buf[IQN_OFF_W_B2 + a2 * IQN_HIDDEN + h],
                      dL_dq * comb_dist[h / 32]);
        }
        if (lane == 0)
            atomicAdd(&grad_buf[IQN_OFF_B_B2 + a2], dL_dq);

        /* ── Gradient through element-wise product ── */
        /* combined = h_s2 ⊙ embed
         * dL/d(embed) = dL/d(combined) ⊙ h_s2
         * dL/d(h_s2) = dL/d(combined) ⊙ embed — NOW COMPUTED for trunk gradient */
        float dL_dembed[IQN_DIST(IQN_HIDDEN)];
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            dL_dembed[h / 32] = dL_dcomb[h / 32] * h_dist[h / 32];
            /* Accumulate dL/d(h_s2) across all quantiles.
             * This flows IQN's bounded Huber gradient to the shared trunk. */
            atomicAdd(&d_h_s2_out[sample * IQN_HIDDEN + h],
                      dL_dcomb[h / 32] * embed_dist[h / 32]);
        }

        /* ── Gradient through ReLU ── */
        /* embed = ReLU(pre_relu) → dL/d(pre_relu) = dL/d(embed) × 𝟙{embed > 0} */
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            if (embed_dist[h / 32] <= 0.0f)
                dL_dembed[h / 32] = 0.0f;
        }

        /* ── Gradient through embedding linear layer ── */
        /* pre_relu[h] = Σ_d W_embed[h,d]·cos(π(d+1)τ) + b_embed[h]
         * dL/dW_embed[h,d] += dL/d(pre_relu)[h] × cos(π(d+1)τ)
         * dL/db_embed[h] += dL/d(pre_relu)[h] */
        float tau = my_taus[ti];
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            float dL_dpre = dL_dembed[h / 32];
            /* Bias gradient */
            atomicAdd(&grad_buf[IQN_OFF_B_EMBED + h], dL_dpre);
            /* Weight gradient: outer product with cosine features */
            for (int d = 0; d < IQN_EMBED_DIM; d++) {
                float cos_val = cosf(3.14159265f * (float)(d + 1) * tau);
                atomicAdd(&grad_buf[IQN_OFF_W_EMBED + h * IQN_EMBED_DIM + d],
                          dL_dpre * cos_val);
            }
        }
    }
}

/* ── Gradient Norm Kernel ───────────────────────────────────────────── */
extern "C" __global__
void iqn_grad_norm_kernel(
    const float* __restrict__ grads,
    float* __restrict__ norm_out,  /* [1] */
    int total_params
)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;

    /* Each thread accumulates its own partial sum */
    float partial = 0.0f;
    for (int i = tid; i < total_params; i += gridDim.x * blockDim.x) {
        float g = grads[i];
        partial += g * g;
    }

    /* Block-level reduction via shared memory */
    __shared__ float shared[256];
    shared[threadIdx.x] = partial;
    __syncthreads();

    /* Tree reduction in shared memory */
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (threadIdx.x < s)
            shared[threadIdx.x] += shared[threadIdx.x + s];
        __syncthreads();
    }

    if (threadIdx.x == 0)
        atomicAdd(norm_out, sqrtf(shared[0]));
}

/* ── Adam Update Kernel ─────────────────────────────────────────────── */
extern "C" __global__
void iqn_adam_kernel(
    float* __restrict__ params,
    float* __restrict__ grads,
    float* __restrict__ m,          /* first moment */
    float* __restrict__ v,          /* second moment */
    const float* __restrict__ norm, /* [1] gradient L2 norm */
    float lr, float beta1, float beta2, float eps,
    float weight_decay, float max_grad_norm,
    int adam_t,                     /* step counter (1-indexed) */
    int total_params
)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;
    if (tid >= total_params) return;

    /* Gradient clipping */
    float grad_norm = norm[0];
    float clip_scale = (grad_norm > max_grad_norm && grad_norm > 0.0f)
        ? max_grad_norm / grad_norm : 1.0f;
    float g = grads[tid] * clip_scale;

    /* AdamW update */
    float mi = beta1 * m[tid] + (1.0f - beta1) * g;
    float vi = beta2 * v[tid] + (1.0f - beta2) * g * g;
    m[tid] = mi;
    v[tid] = vi;

    /* Bias correction */
    float m_hat = mi / (1.0f - powf(beta1, (float)adam_t));
    float v_hat = vi / (1.0f - powf(beta2, (float)adam_t));

    /* Parameter update with decoupled weight decay */
    float p = params[tid];
    p -= lr * (m_hat / (sqrtf(v_hat) + eps) + weight_decay * p);
    params[tid] = p;

    /* Zero gradient for next step */
    grads[tid] = 0.0f;
}

/* ── Forward-Only Kernel (Inference) ────────────────────────────────── */
/**
 * IQN forward pass for inference — no loss, no activation saves.
 *
 * Computes expected Q-values per action per branch by averaging over
 * quantile samples. Used for action selection blending with C51.
 *
 * Outputs: expected_q [B, TOTAL_BRANCH_ACTIONS] — mean Q per action.
 */
extern "C" __global__
void iqn_forward_kernel(
    const float* __restrict__ h_s2,     /* [B, IQN_HIDDEN] */
    const float* __restrict__ taus,     /* [B, N] */
    const float* __restrict__ params,   /* [IQN_TOTAL_PARAMS] */
    float* __restrict__ expected_q,     /* [B, TOTAL_BRANCH_ACTIONS] */
    int batch_size
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;
    int lane = threadIdx.x;

    const float* my_h_s2 = h_s2 + sample * IQN_HIDDEN;
    const float* my_taus = taus + sample * IQN_NUM_QUANTILES;
    const float* w_embed = params + IQN_OFF_W_EMBED;
    const float* b_embed = params + IQN_OFF_B_EMBED;
    const float* w_b0 = params + IQN_OFF_W_B0;
    const float* b_b0 = params + IQN_OFF_B_B0;
    const float* w_b1 = params + IQN_OFF_W_B1;
    const float* b_b1 = params + IQN_OFF_B_B1;
    const float* w_b2 = params + IQN_OFF_W_B2;
    const float* b_b2 = params + IQN_OFF_B_B2;

    /* Load h_s2 */
    float h_dist[IQN_DIST(IQN_HIDDEN)];
    for (int h = lane; h < IQN_HIDDEN; h += 32)
        h_dist[h / 32] = my_h_s2[h];

    /* Accumulate Q-values across quantiles (for mean) */
    float q_acc[TOTAL_BRANCH_ACTIONS];
    for (int a = 0; a < TOTAL_BRANCH_ACTIONS; a++)
        q_acc[a] = 0.0f;

    for (int t = 0; t < IQN_NUM_QUANTILES; t++) {
        float tau = my_taus[t];

        /* Cosine embedding + linear + ReLU */
        float embed_dist[IQN_DIST(IQN_HIDDEN)];
        for (int h = lane; h < IQN_HIDDEN; h += 32) {
            float acc = b_embed[h];
            const float* w_row = w_embed + h * IQN_EMBED_DIM;
            for (int d = 0; d < IQN_EMBED_DIM; d++) {
                float cos_val = cosf(3.14159265f * (float)(d + 1) * tau);
                acc += w_row[d] * cos_val;
            }
            embed_dist[h / 32] = fmaxf(acc, 0.0f);
        }

        /* Element-wise product */
        float comb_dist[IQN_DIST(IQN_HIDDEN)];
        for (int h = lane; h < IQN_HIDDEN; h += 32)
            comb_dist[h / 32] = h_dist[h / 32] * embed_dist[h / 32];

        /* Branch outputs */
        /* Branch 0 */
        for (int a = 0; a < BRANCH_0_SIZE; a++) {
            float partial = 0.0f;
            const float* w_row = w_b0 + a * IQN_HIDDEN;
            for (int h = lane; h < IQN_HIDDEN; h += 32)
                partial += w_row[h] * comb_dist[h / 32];
            float q_val = iqn_warp_sum(partial) + b_b0[a];
            q_acc[a] += q_val;
        }
        /* Branch 1 */
        for (int a = 0; a < BRANCH_1_SIZE; a++) {
            float partial = 0.0f;
            const float* w_row = w_b1 + a * IQN_HIDDEN;
            for (int h = lane; h < IQN_HIDDEN; h += 32)
                partial += w_row[h] * comb_dist[h / 32];
            float q_val = iqn_warp_sum(partial) + b_b1[a];
            q_acc[BRANCH_0_SIZE + a] += q_val;
        }
        /* Branch 2 */
        for (int a = 0; a < BRANCH_2_SIZE; a++) {
            float partial = 0.0f;
            const float* w_row = w_b2 + a * IQN_HIDDEN;
            for (int h = lane; h < IQN_HIDDEN; h += 32)
                partial += w_row[h] * comb_dist[h / 32];
            float q_val = iqn_warp_sum(partial) + b_b2[a];
            q_acc[BRANCH_0_SIZE + BRANCH_1_SIZE + a] += q_val;
        }
    }

    /* Write mean Q-values (averaged over quantiles) */
    float inv_n = 1.0f / (float)IQN_NUM_QUANTILES;
    if (lane == 0) {
        for (int a = 0; a < TOTAL_BRANCH_ACTIONS; a++)
            expected_q[sample * TOTAL_BRANCH_ACTIONS + a] = q_acc[a] * inv_n;
    }
}

/* ── Trunk Forward Kernel ──────────────────────────────────────────── */
/**
 * Compute h_s2 from raw states using DQN shared trunk weights.
 *
 * state → LeakyReLU(W_s1 × state + b_s1) → LeakyReLU(W_s2 × h1 + b_s2) → h_s2
 *
 * Used to compute target_h_s2 from next_states + target DQN weights.
 * One warp per sample. Uses shared memory for h1 intermediate.
 *
 * Compile-time defines: IQN_STATE_DIM, IQN_SHARED_H1, IQN_HIDDEN (=SHARED_H2)
 */
#ifndef IQN_STATE_DIM
#define IQN_STATE_DIM 48
#endif
#ifndef IQN_SHARED_H1
#define IQN_SHARED_H1 256
#endif

extern "C" __global__
void iqn_trunk_forward_kernel(
    const float* __restrict__ states,   /* [B, IQN_STATE_DIM] */
    const float* __restrict__ w_s1,     /* [IQN_SHARED_H1, IQN_STATE_DIM] */
    const float* __restrict__ b_s1,     /* [IQN_SHARED_H1] */
    const float* __restrict__ w_s2,     /* [IQN_HIDDEN, IQN_SHARED_H1] */
    const float* __restrict__ b_s2,     /* [IQN_HIDDEN] */
    float* __restrict__ h_s2_out,       /* [B, IQN_HIDDEN] */
    int batch_size
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;
    int lane = threadIdx.x;

    extern __shared__ float shmem[];
    float* h1 = shmem; /* [IQN_SHARED_H1] */

    const float* my_state = states + sample * IQN_STATE_DIM;

    /* Layer 1: h1 = leaky_relu(W_s1 @ state + b_s1) */
    for (int h = lane; h < IQN_SHARED_H1; h += 32) {
        float acc = b_s1[h];
        const float* w_row = w_s1 + h * IQN_STATE_DIM;
        for (int d = 0; d < IQN_STATE_DIM; d++)
            acc += w_row[d] * my_state[d];
        h1[h] = (acc > 0.0f) ? acc : 0.01f * acc;
    }
    __syncwarp();

    /* Layer 2: h2 = leaky_relu(W_s2 @ h1 + b_s2) */
    for (int h = lane; h < IQN_HIDDEN; h += 32) {
        float acc = b_s2[h];
        const float* w_row = w_s2 + h * IQN_SHARED_H1;
        for (int d = 0; d < IQN_SHARED_H1; d++)
            acc += w_row[d] * h1[d];
        float val = (acc > 0.0f) ? acc : 0.01f * acc;
        h_s2_out[sample * IQN_HIDDEN + h] = val;
    }
}

/* ── EMA Target Update Kernel ──────────────────────────────────────── */
/**
 * Polyak EMA: target[i] = (1 - tau) * target[i] + tau * online[i]
 */
extern "C" __global__
void iqn_ema_kernel(
    float* __restrict__ target,
    const float* __restrict__ online,
    float tau,
    int n
)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n)
        target[i] = (1.0f - tau) * target[i] + tau * online[i];
}

/* ── Flat action → branch action decode kernel ─────────────────────────
 *
 * Decodes flat action index ∈ [0, B0×B1×B2) into 3 branch indices.
 * Grid: (ceil(batch_size/256), 1, 1), Block: (256, 1, 1)
 *
 * Input:  flat_actions[B]   -- i32 flat action indices
 * Output: branch_actions[B*3] -- i32 [exposure, order, urgency] per sample
 */
extern "C" __global__
void iqn_decode_actions_kernel(
    const int* __restrict__ flat_actions,
    int* __restrict__ branch_actions,
    int batch_size
)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= batch_size) return;

    int flat = flat_actions[i];
    /* Bounds check: clamp to valid action range [0, total_actions) */
    int max_action = BRANCH_0_SIZE * BRANCH_1_SIZE * BRANCH_2_SIZE;
    if (flat < 0 || flat >= max_action) flat = 0;
    int stride = BRANCH_1_SIZE * BRANCH_2_SIZE;
    int exposure = flat / stride;
    int remainder = flat % stride;
    int order     = remainder / BRANCH_2_SIZE;
    int urgency   = remainder % BRANCH_2_SIZE;

    branch_actions[i * 3 + 0] = exposure;
    branch_actions[i * 3 + 1] = order;
    branch_actions[i * 3 + 2] = urgency;
}

/* ── GPU τ sampling kernel (Philox-based PRNG) ─────────────────────────
 *
 * Generates uniform random τ ∈ (0.01, 0.99) directly on GPU.
 * Uses Philox 4×32 counter-based PRNG — deterministic, high quality,
 * no global state. Each thread produces one τ value from (thread_id, seed).
 *
 * Grid: (ceil(total/256), 1, 1), Block: (256, 1, 1)
 * total = batch_size × num_quantiles
 */

/* Philox 4×32-10 round function — same PRNG used by PyTorch/JAX */
__device__ __forceinline__
unsigned int iqn_philox_single(unsigned int counter, unsigned int key) {
    unsigned int hi, lo;
    /* Philox S-box constants */
    lo = counter * 0xD2511F53u;
    hi = __umulhi(counter, 0xD2511F53u);
    /* 10 rounds of mixing */
    for (int r = 0; r < 10; r++) {
        unsigned int t = hi ^ key;
        hi = lo * 0xCD9E8D57u;
        lo = __umulhi(lo, 0xCD9E8D57u);
        lo ^= t;
        key += 0x9E3779B9u; /* golden ratio */
    }
    return lo ^ hi;
}

extern "C" __global__
void iqn_sample_taus_kernel(
    float* __restrict__ taus,
    unsigned int seed,
    int total
)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= total) return;

    /* Counter = thread index, key = seed — unique per (step, online/target) */
    unsigned int bits = iqn_philox_single((unsigned int)i, seed);
    /* Map to (0, 1) then clamp to (0.01, 0.99) */
    float u = (float)(bits >> 8) * (1.0f / 16777216.0f); /* 24-bit mantissa */
    taus[i] = 0.01f + u * 0.98f;
}
