/**
 * IQL (Implicit Q-Learning) Value Network Training Kernel
 *
 * Kostrikov et al. 2021 -- trains V(s) with asymmetric expectile regression:
 *
 *   L_tau(u) = |tau - 1(u < 0)| * u^2,  where u = Q(s,a) - V(s)
 *
 * Architecture: 2-hidden-layer MLP with SiLU activation.
 *   state -> Linear(STATE_DIM, H) -> SiLU -> Linear(H, H) -> SiLU -> Linear(H, 1) -> V(s)
 *
 * Launch configs:
 *   Forward + loss:         grid=(B, 1, 1),                   block=(256, 1, 1)
 *   Backward (per-sample):  grid=(B, 1, 1),                   block=(256, 1, 1)
 *   Weight grad reduce:     grid=(ceil(P/256), 1, 1),         block=(256, 1, 1)
 *   Grad norm phase 1:      grid=(ceil(P/256), 1, 1),         block=(256, 1, 1)
 *   Grad norm phase 2:      grid=(1, 1, 1),                   block=(1, 1, 1)
 *   Adam:                   grid=(ceil(P/256), 1, 1),         block=(256, 1, 1)
 *   Forward-only:           grid=(B, 1, 1),                   block=(256, 1, 1)
 *   Advantage weights:      grid=(ceil(B/256), 1, 1),         block=(256, 1, 1)
 *
 * Compile-time defines (injected via build.rs dim_overrides):
 *   STATE_DIM          -- tensor-core-aligned state width (e.g. 48 or 56)
 *   VALUE_HIDDEN_DIM   -- hidden layer width (default 128)
 *
 * ALL storage and computation in native f32. Zero atomicAdd — fully deterministic.
 */

#ifndef VALUE_HIDDEN_DIM
#define VALUE_HIDDEN_DIM 128
#endif

/* SP4 Layer B: per-branch atom-position bound from ISV[ATOM_POS_BOUND_BASE + branch].
 * Mirrors `crates/ml/src/cuda_pipeline/sp4_isv_slots.rs::ATOM_POS_BOUND_BASE = 132`.
 */
#ifndef ATOM_POS_BOUND_BASE
#define ATOM_POS_BOUND_BASE 132
#endif

/* Parameter sizes that don't depend on state_dim */
#define V_B1_SIZE  (VALUE_HIDDEN_DIM)
#define V_W2_SIZE  (VALUE_HIDDEN_DIM * VALUE_HIDDEN_DIM)
#define V_B2_SIZE  (VALUE_HIDDEN_DIM)
#define V_W3_SIZE  (VALUE_HIDDEN_DIM * 2)   /* output layer: 2 x H (V_short + V_long) */
#define V_B3_SIZE  (2)

/* Runtime offset computation helper -- state_dim is a kernel parameter.
 * W1 size = VALUE_HIDDEN_DIM * state_dim (varies with state_dim). */
__device__ __forceinline__ void iql_compute_offsets(
    int state_dim,
    int* off_w1, int* off_b1, int* off_w2, int* off_b2, int* off_w3, int* off_b3
) {
    int w1_size = VALUE_HIDDEN_DIM * state_dim;
    *off_w1 = 0;
    *off_b1 = w1_size;
    *off_w2 = *off_b1 + V_B1_SIZE;
    *off_b2 = *off_w2 + V_W2_SIZE;
    *off_w3 = *off_b2 + V_B2_SIZE;
    *off_b3 = *off_w3 + V_W3_SIZE;
}

/* ------------------------------------------------------------------ */
/*  SiLU activation: x * sigmoid(x)                                    */
/* ------------------------------------------------------------------ */
__device__ __forceinline__ float silu(float x) {
    return x / (1.0f + expf(-x));
}

__device__ __forceinline__ float silu_grad(float x) {
    float s = 1.0f / (1.0f + expf(-x));
    return s * (1.0f + x * (1.0f - s));
}

/* Block-level f32 sum reduction: 256 threads -> single sum.
 * Returns the result broadcast to ALL threads in the block. */
__device__ __forceinline__ float iql_block_sum(
    float val, float* warp_sums /* [8] in shared memory */
) {
    int tid = threadIdx.x;
    int warp_id = tid / 32;
    int warp_lane = tid % 32;
    /* Intra-warp sum via shfl_xor */
    for (int offset = 16; offset > 0; offset >>= 1)
        val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
    if (warp_lane == 0) warp_sums[warp_id] = val;
    __syncthreads();
    if (warp_id == 0) {
        float v = (warp_lane < (blockDim.x / 32)) ? warp_sums[warp_lane] : 0.0f;
        for (int off = 16; off > 0; off >>= 1)
            v += __shfl_xor_sync(0xFFFFFFFF, v, off);
        warp_sums[0] = v;
    }
    __syncthreads();
    return warp_sums[0];
}

/* ------------------------------------------------------------------ */
/*  Loss Reduce (deterministic sequential sum)                         */
/* ------------------------------------------------------------------ */
/**
 * Reduce per-sample losses to a single batch-mean scalar.
 *
 * total_loss[0] = (1/B) * sum(loss_out[b])
 *
 * Sequential for determinism. Launch: grid=1, block=1.
 */
extern "C" __global__
void iql_loss_reduce(
    const float* __restrict__ loss_out,    /* [B * num_heads] per-element loss */
    float* __restrict__ total_loss,        /* [1] output */
    int batch_size,
    int num_heads
)
{
    float sum = 0.0f;
    int total = batch_size * num_heads;
    for (int i = 0; i < total; i++) {
        sum += loss_out[i];
    }
    /* Normalise by B so the loss scale is consistent with the single-head case.
     * Summing over num_heads and dividing by B gives the mean over samples;
     * each head contributes equally. */
    total_loss[0] = sum / (float)batch_size;
}

/* ------------------------------------------------------------------ */
/*  Gradient Norm — Phase 1 (per-block partial sums, no atomicAdd)     */
/* ------------------------------------------------------------------ */
/**
 * Each block computes a partial sum-of-squares and writes to block_sums.
 * No atomicAdd — one output per block.
 */
extern "C" __global__
void iql_grad_norm_phase1(
    const float* __restrict__ grads,
    float* __restrict__ block_sums,    /* [num_blocks] */
    int total_params
)
{
    int tid = blockIdx.x * blockDim.x + threadIdx.x;

    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();

    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)
        block_sums[blockIdx.x] = shared[0];
}

/* ------------------------------------------------------------------ */
/*  Gradient Norm — Phase 2 (sequential sum of partials)               */
/* ------------------------------------------------------------------ */
/**
 * Sequential sum of block partial sums → sum-of-squares output.
 * Adam kernel reads this as raw sum-of-squares (takes sqrt internally).
 *
 * Launch: grid=1, block=1.
 */
extern "C" __global__
void iql_grad_norm_phase2(
    const float* __restrict__ block_sums,
    float* __restrict__ norm_out,      /* [1] sum-of-squares */
    int num_blocks
)
{
    float sum = 0.0f;
    for (int i = 0; i < num_blocks; i++)
        sum += block_sums[i];
    norm_out[0] = sum;
}

/* ------------------------------------------------------------------ */
/*  Adam Update Kernel                                                 */
/* ------------------------------------------------------------------ */
/**
 * AdamW optimizer step over the flat parameter buffer.
 *
 * grid = ceil(total_params / 256), block = 256.
 * Each thread updates one parameter.
 *
 * grad_norm_sq_buf contains raw sum-of-squares from phase1+phase2.
 * Gradient clipping: scale = max_grad_norm / sqrt(grad_norm_sq).
 */
extern "C" __global__
void iql_adam_kernel(
    float* __restrict__ params,
    float* __restrict__ grads,
    float* __restrict__ m,          /* first moment */
    float* __restrict__ v,          /* second moment */
    const float* __restrict__ grad_norm_sq_buf,  /* [1] sum-of-squares */
    float lr,
    float beta1,
    float beta2,
    float eps,
    float weight_decay,
    float max_grad_norm,
    const int* __restrict__ t_buf,  /* [1] Adam step on device */
    int total_params,
    /* SP4 Layer B: ISV-driven |p| bound (`ISV[WEIGHT_BOUND[IqlHigh|IqlLow]].max(EPS_CLAMP_FLOOR)`)
     * — same pattern wired into all five Adam kernels. 0=disabled (debug only). */
    float weight_clamp_max_abs,
    /* SP4 Task A14 (2026-04-30): Pearl C engagement counter. Same warp+
     * block tree-reduce pattern as `dqn_adam_update_kernel`. Caller
     * passes `engage_buf_offset = group_idx × MAX_BLOCKS_PER_ADAM` (group
     * = ParamGroup::IqlHigh or IqlLow depending on which IQL trainer
     * launched). `engage_buf_offset < 0` skips the writeback. */
    int* __restrict__ nan_flags_buf,
    int diag_slot,
    int* __restrict__ clamp_engage_per_block_buf,
    int engage_buf_offset
)
{
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    bool active = (i < total_params);
    int local_engage = 0;

    if (active) {
        /* Gradient clipping from sum-of-squares */
        float gn = sqrtf(grad_norm_sq_buf[0]);
        float clip_scale = (gn > max_grad_norm && gn > 0.0f) ? (max_grad_norm / gn) : 1.0f;
        float g = grads[i] * clip_scale;

        /* AdamW: decoupled weight decay */
        float p = params[i];
        p -= lr * weight_decay * p;

        /* Adam moment updates */
        float m_i = beta1 * m[i] + (1.0f - beta1) * g;
        float v_i = beta2 * v[i] + (1.0f - beta2) * g * g;

        /* Bias correction */
        int t = t_buf[0];
        float m_hat = m_i / (1.0f - powf(beta1, (float)t));
        float v_hat = v_i / (1.0f - powf(beta2, (float)t));

        /* Parameter update */
        p -= lr * m_hat / (sqrtf(v_hat) + eps);

        /* SP3 Mech 9 + SP4 Task A14: post-Adam |p| clamp + engagement
         * counter. Last step before writeback — applies after all
         * decay/decoupled updates so the bound holds for the value the
         * next forward sees. */
        if (weight_clamp_max_abs > 0.0f) {
            if (fabsf(p) > weight_clamp_max_abs) {
                if (nan_flags_buf != nullptr && diag_slot >= 0) {
                    nan_flags_buf[diag_slot] = 1;
                }
                local_engage = 1;
            }
            p = fminf(fmaxf(p, -weight_clamp_max_abs), weight_clamp_max_abs);
        }

        params[i] = p;
        m[i] = m_i;
        v[i] = v_i;
        grads[i] = 0.0f;  /* zero gradient for next step */
    }

    /* SP4 Task A14: warp + block tree-reduce of `local_engage`. */
    int e = local_engage;
    for (int offset = 16; offset > 0; offset >>= 1)
        e += __shfl_xor_sync(0xFFFFFFFF, e, offset);

    __shared__ int warp_engage[8]; /* blockDim.x = 256 ⇒ 8 warps */
    int warp_id   = threadIdx.x / 32;
    int warp_lane = threadIdx.x % 32;
    if (warp_lane == 0) warp_engage[warp_id] = e;
    __syncthreads();

    if (warp_id == 0) {
        int val = (warp_lane < blockDim.x / 32) ? warp_engage[warp_lane] : 0;
        for (int offset = 16; offset > 0; offset >>= 1)
            val += __shfl_xor_sync(0xFFFFFFFF, val, offset);
        if (warp_lane == 0 && engage_buf_offset >= 0 && clamp_engage_per_block_buf != nullptr) {
            clamp_engage_per_block_buf[engage_buf_offset + blockIdx.x] = val;
        }
    }
}

/* ------------------------------------------------------------------ */
/*  Gather Q(s, a_taken) from q_out_buf                                */
/* ------------------------------------------------------------------ */
/**
 * Extract Q-value for the taken factored action from branch Q-values.
 *
 * q_out is [B, b0+b1+b2+b3] — per-branch expected Q-values.
 * actions[b] is a factored action index: dir*b1*b2*b3 + mag*b2*b3 + ord*b3 + urg.
 * Q(s,a) = q_out[b, dir_offset + dir] + q_out[b, mag_offset + mag]
 *         + q_out[b, ord_offset + ord] + q_out[b, urg_offset + urg]
 *
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_gather_q_taken(
    const float* __restrict__ q_out,       /* [B, b0+b1+b2+b3] */
    const int* __restrict__ actions,        /* [B] factored action indices */
    float* __restrict__ q_taken,            /* [B] output */
    int batch_size,
    int b0, int b1, int b2, int b3
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    int a = actions[b];
    int total_per_row = b0 + b1 + b2 + b3;
    const float* row = q_out + b * total_per_row;

    /* Decode factored action: a = dir*(b1*b2*b3) + mag*(b2*b3) + ord*b3 + urg */
    int urg = a % b3;           a /= b3;
    int ord = a % b2;           a /= b2;
    int mag = a % b1;           a /= b1;
    int dir = a;

    q_taken[b] = row[dir]
               + row[b0 + mag]
               + row[b0 + b1 + ord]
               + row[b0 + b1 + b2 + urg];
}

/* ------------------------------------------------------------------ */
/*  Forward-Only Kernel (for inference / advantage computation)        */
/* ------------------------------------------------------------------ */
/**
 * V(s) forward pass only -- no loss, no backward.
 *
 * Used for computing advantages A(s,a) = Q(s,a) - V(s) at inference time.
 * Same architecture as forward_loss but skips loss and activation saves.
 * 256 threads (8 warps) per sample.
 */
extern "C" __global__
void iql_forward_kernel(
    const float* __restrict__ states,
    const float* __restrict__ params,
    float* __restrict__ v_out,
    int batch_size,
    int state_dim
)
{
    int sample = blockIdx.x;
    if (sample >= batch_size) return;

    int tid = threadIdx.x; /* 0..255 */

    int off_w1, off_b1, off_w2, off_b2, off_w3, off_b3;
    iql_compute_offsets(state_dim, &off_w1, &off_b1, &off_w2, &off_b2, &off_w3, &off_b3);

    const float* w1 = params + off_w1;
    const float* b1 = params + off_b1;
    const float* w2 = params + off_w2;
    const float* b2 = params + off_b2;
    const float* w3 = params + off_w3;
    const float* b3 = params + off_b3;

    const float* x = states + sample * state_dim;

    /* Shared memory for intermediate activations + block reduction */
    __shared__ float sh1[VALUE_HIDDEN_DIM];
    __shared__ float sh2[VALUE_HIDDEN_DIM];
    __shared__ float warp_sums[8];

    /* Layer 1: SiLU */
    for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) {
        float acc = b1[j];
        for (int k = 0; k < state_dim; k++) {
            acc += w1[j * state_dim + k] * x[k];
        }
        sh1[j] = silu(acc);
    }
    __syncthreads();

    /* Layer 2: SiLU */
    for (int j = tid; j < VALUE_HIDDEN_DIM; j += 256) {
        float acc = b2[j];
        for (int k = 0; k < VALUE_HIDDEN_DIM; k++) {
            acc += w2[j * VALUE_HIDDEN_DIM + k] * sh1[k];
        }
        sh2[j] = silu(acc);
    }
    __syncthreads();

    /* Output: two heads (V_short, V_long) — dot products with rows 0 and 1
     * of W3[2, H] stored row-major in params. v_out layout [2, B] col-major:
     * v_out[0*B + sample] = V_short, v_out[1*B + sample] = V_long.
     * Consumers sum the two: V(s) = v_out[sample*2+0] + v_out[sample*2+1]
     * using the [B*2] stride-2 view produced by the cuBLAS forward path. */
    float v0_acc = 0.0f;
    float v1_acc = 0.0f;
    for (int k = tid; k < VALUE_HIDDEN_DIM; k += 256) {
        v0_acc += w3[k]                    * sh2[k];
        v1_acc += w3[VALUE_HIDDEN_DIM + k] * sh2[k];
    }
    v0_acc = iql_block_sum(v0_acc, warp_sums);
    __syncthreads();
    v1_acc = iql_block_sum(v1_acc, warp_sums);

    if (tid == 0) {
        /* col-major [2, B]: row r, col b stored at r + b*2 */
        v_out[sample * 2 + 0] = v0_acc + b3[0];
        v_out[sample * 2 + 1] = v1_acc + b3[1];
    }
}

/* ------------------------------------------------------------------ */
/*  Advantage Weight Kernel                                            */
/* ------------------------------------------------------------------ */
/**
 * Compute IQL advantage weights for policy modulation:
 *   w[b] = clamp(exp(beta * (Q_taken[b] - V(s)[b])), 0.01, 100)
 *
 * Q_taken is gathered from q_out_buf using the taken action index.
 * Weights are clamped for numerical stability.
 *
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_compute_advantage_weights(
    const float* __restrict__ q_out,       /* [B, b0+b1+b2+b3] branch Q-values */
    const int* __restrict__ actions,        /* [B] factored action indices */
    const float* __restrict__ v_out,        /* [B*2] V(s) two heads col-major [2,B] */
    const float* __restrict__ readiness_buf,/* [1] CV-based readiness */
    float* __restrict__ adv_weights,        /* [B] output weights */
    float beta,
    int batch_size,
    int b0, int b1, int b2, int b3
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    float r = readiness_buf[0];
    int total_per_row = b0 + b1 + b2 + b3;
    const float* row = q_out + b * total_per_row;
    int a = actions[b];
    int urg = a % b3; a /= b3;
    int ord = a % b2; a /= b2;
    int mag = a % b1; a /= b1;
    int dir = a;
    float q_taken = row[dir] + row[b0 + mag] + row[b0 + b1 + ord] + row[b0 + b1 + b2 + urg];
    /* V(s) = V_short + V_long; col-major [2, B]: element (head, b) at head + b*2 */
    float adv = q_taken - (v_out[b * 2 + 0] + v_out[b * 2 + 1]);
    float raw_w = expf(beta * adv);
    float clamped_w = fminf(fmaxf(raw_w, 0.01f), 100.0f);
    /* Blend: readiness=0 → neutral weight 1.0, readiness=1 → IQL weight */
    adv_weights[b] = r * clamped_w + (1.0f - r) * 1.0f;
}

/* ------------------------------------------------------------------ */
/*  TD-Error Modulation Kernel                                          */
/* ------------------------------------------------------------------ */
/**
 * Modulate TD errors by advantage weights and staleness decay:
 *   K     = max(exp(beta * 3 * max(sigma_adv, 1e-6)), 1.1)
 *   w     = clamp(adv_weights[b], 1/K, K)
 *   age   = (write_pos - indices[b] + capacity) % capacity
 *   decay = exp(-staleness_lambda * age / max(staleness_tau, 1.0))
 *   td_errors[b] *= w * decay
 *
 * In-place update of td_errors for PER priority recomputation.
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_modulate_td_errors(
    float* __restrict__ td_errors,           /* [B] in-place modulated */
    const float* __restrict__ adv_weights,   /* [B] advantage weights */
    const int* __restrict__ indices,          /* [B] buffer indices for staleness */
    const float* __restrict__ sigma_adv_buf, /* [1] EMA of advantage std */
    const float* __restrict__ readiness_buf, /* [6] drift readiness state */
    float beta,
    int write_pos,
    int capacity,
    int batch_size
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    float r = readiness_buf[0];

    float sigma_adv = sigma_adv_buf[0];
    float K = expf(beta * 3.0f * fmaxf(sigma_adv, 1e-6f));
    K = fmaxf(K, 1.1f);
    float inv_K = 1.0f / K;

    float w = fminf(fmaxf(adv_weights[b], inv_K), K);

    /* Buffer-relative staleness: age normalized to [0,1] by capacity.
     * decay_k adapts to buffer utilization: ln(20) ≈ 3.0 gives exp(-3)=0.05
     * for oldest when buffer is full. This is the natural constant:
     * "oldest transition should contribute ~5% weight" → k = -ln(0.05) = 3.0.
     * No tuning — 5% is the threshold where a transition is effectively dead. */
    int age = (write_pos - indices[b] + capacity) % capacity;
    float age_norm = (float)age / fmaxf((float)capacity, 1.0f);
    float decay = expf(-logf(20.0f) * age_norm);  /* oldest gets 1/20 = 5% weight */

    /* Blend: readiness=0 → td_errors unchanged, readiness=1 → full modulation */
    float modulation = r * (w * decay) + (1.0f - r) * 1.0f;
    td_errors[b] *= modulation;
}

/* ------------------------------------------------------------------ */
/*  Advantage Sigma EMA + Adaptive CV Readiness (GPU-side, zero DtoH)  */
/* ------------------------------------------------------------------ */
/**
 * Adaptive readiness from advantage CV improvement.
 *
 * readiness = (cv_initial - cv_current) / cv_initial
 * Measures fractional improvement from initial noise level.
 * Self-calibrating to any batch size, reward scale, Q-value magnitude.
 *
 * readiness_buf layout [4]:
 *   [0] readiness scalar
 *   [1] cv_initial (first observed CV)
 *   [2] cv_max (running max — regime-change anchor)
 *   [3] reserved
 *
 * Launch: grid=1, block=1.
 */
extern "C" __global__
void iql_adv_sigma_ema_update(
    const float* __restrict__ adv_stats,     /* [2]: mean, variance */
    float* __restrict__ sigma_ema,           /* [1] running EMA of advantage std */
    float* __restrict__ readiness_buf,       /* [4] CV readiness state */
    float ema_beta                           /* 0.99 for sigma EMA */
)
{
    /* ── 1. Advantage sigma EMA ── */
    float mean = adv_stats[0];
    float var = adv_stats[1];
    float sigma = sqrtf(fmaxf(var, 0.0f));
    float prev_sigma = sigma_ema[0];
    if (prev_sigma < 1e-8f) {
        sigma_ema[0] = sigma;
    } else {
        sigma_ema[0] = ema_beta * prev_sigma + (1.0f - ema_beta) * sigma;
    }

    /* ── 2. CV readiness: improvement from decaying max ── */
    float cv = sigma_ema[0] / fmaxf(fabsf(mean), 1e-6f);
    float cv_initial = readiness_buf[1];
    float cv_max = readiness_buf[2];

    if (cv_initial < 1e-8f) {
        /* First step: capture initial CV */
        readiness_buf[1] = cv;
        readiness_buf[2] = cv;
        readiness_buf[0] = 0.0f;
    } else {
        /* cv_max: decays slowly (0.999/step ≈ 1000-step half-life) so stale
         * spikes don't permanently anchor readiness. Updates on new spikes. */
        cv_max = cv_max * 0.999f;
        if (cv > cv_max) cv_max = cv;
        readiness_buf[2] = cv_max;

        float anchor = fmaxf(cv_max, cv_initial);
        float improvement = fminf(fmaxf((anchor - cv) / fmaxf(anchor, 1e-6f), 0.0f), 1.0f);
        /* Adaptive-rate EMA: smooth readiness transitions to prevent
           per-sample support from snapping discontinuously. Tracks fast
           when readiness changes rapidly, smooths when stable. */
        float prev_r = readiness_buf[0];
        float r_err = fabsf(improvement - prev_r);
        float r_alpha = r_err / (r_err + 0.05f);
        r_alpha = fminf(fmaxf(r_alpha, 0.01f), 0.3f);
        readiness_buf[0] = (1.0f - r_alpha) * prev_r + r_alpha * improvement;
    }
}

/* ------------------------------------------------------------------ */
/*  Advantage Variance Reduction Kernel                                 */
/* ------------------------------------------------------------------ */
/**
 * Compute mean and variance of advantage weights over a batch.
 * Single-thread kernel — launched with grid=1, block=1.
 *   stats[0] = mean(adv_weights)
 *   stats[1] = var(adv_weights)
 *
 * Used to adaptively modulate the TD-error clamp range K.
 * Launch: grid=1, block=1.
 */
extern "C" __global__
void iql_adv_variance_reduce(
    const float* __restrict__ adv_weights,   /* [B] */
    float* __restrict__ stats,               /* [2]: mean, variance */
    int batch_size
)
{
    float sum = 0.0f;
    for (int b = 0; b < batch_size; b++)
        sum += adv_weights[b];
    float mean = sum / (float)batch_size;

    float var_sum = 0.0f;
    for (int b = 0; b < batch_size; b++) {
        float d = adv_weights[b] - mean;
        var_sum += d * d;
    }
    stats[0] = mean;
    stats[1] = var_sum / (float)batch_size;
}

/* ------------------------------------------------------------------ */
/*  Per-Sample C51 Atom Support Kernel — per-branch [B, 4, 3] tile      */
/* ------------------------------------------------------------------ */
/**
 * Compute per-sample, per-branch [v_min, v_max, delta_z] for C51 atom
 * projection. Spec: docs/superpowers/specs/2026-04-23-isv-v-range-unification.md
 * (Phase 2d).
 *
 *   spread   = max_{a} |Q(s,a) - V(s)|
 *   half_w   = spread * (1 + gamma)                  (shared across branches)
 *   centre_b = isv_signals[V_CENTER_DIR + 2*b]       (per-branch from ISV)
 *   v_min_b  = centre_b - half_w
 *   v_max_b  = centre_b + half_w
 *   delta_z  = (v_max - v_min) / (num_atoms - 1)
 *
 * Rationale: the V(s)-derived half_w captures per-sample Q magnitude
 * (how wide the atoms need to span around the centre). The centre
 * differs per branch because direction/magnitude/order/urgency have
 * genuinely different Q scales — the ISV bus broadcasts each branch's
 * centre EMA.
 *
 * Bootstrap safety: when ISV centres are 0 and readiness=0 blend falls
 * back to [-1, 1], all 4 branch triples become [-1, 1, 2/(NA-1)] —
 * byte-identical to the pre-spec single-range behaviour at epoch 1.
 *
 * Output layout stride-12: per_sample_support[b*12 + d*3 + {0,1,2}] per
 * branch d in 0..4. The 4-branch tile is read by the C51 loss / grad
 * kernels and downstream experience kernels that migrated in lockstep.
 *
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_compute_per_sample_support(
    const float* __restrict__ v_out,              /* [B*2] two heads col-major [2,B] */
    const float* __restrict__ q_out,              /* [B, total_actions] */
    float* __restrict__ per_sample_support,       /* [B, 4, 3] */
    const float* __restrict__ readiness_buf,      /* [1] CV-based readiness */
    const float* __restrict__ isv_signals,        /* [ISV_TOTAL_DIM]; NULL → 0 centres */
    float gamma,
    int batch_size,
    int total_actions,
    int num_atoms
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    float r = readiness_buf[0];
    /* V(s) = V_short + V_long; col-major [2, B]: element (head, b) at head + b*2 */
    float v = v_out[b * 2 + 0] + v_out[b * 2 + 1];
    const float* q = q_out + b * total_actions;

    float spread = 0.0f;
    for (int a = 0; a < total_actions; a++) {
        float dist = fabsf(q[a] - v);
        spread = fmaxf(spread, dist);
    }

    float default_dz = 2.0f / (float)(num_atoms - 1);

    /* Per-branch centre from ISV slots 23 + 2*d. When ISV is null (smoke
     * paths / pre-init) the centre collapses to 0; in that regime the
     * readiness blend below falls back to the uniform [-1, 1] default
     * while r ramps from 0 → 1. */
    for (int d = 0; d < 4; d++) {
        float isv_centre = (isv_signals != nullptr)
            ? isv_signals[23 + 2 * d]
            : 0.0f;

        /* Plan 2 D.2: per-branch γ from ISV[43..47). Each branch uses its own
         * discount horizon to scale the Q-spread half-width. Falls back to the
         * scalar gamma parameter when ISV is NULL (smoke paths / pre-init).
         * ISV[43]=DIR, ISV[44]=MAG, ISV[45]=ORD, ISV[46]=URG. */
        float gamma_b = (isv_signals != nullptr)
            ? isv_signals[43 + d]
            : gamma;
        float half_w = spread * (1.0f + gamma_b);

        /* Per-branch IQL range: centre from ISV, half from per-sample
         * V(s)-derived Q spread. Direction/magnitude/order/urgency
         * use independent centres and per-branch-scaled half-widths. */
        float iql_vmin = isv_centre - half_w;
        float iql_vmax = isv_centre + half_w;
        float iql_dz   = (iql_vmax - iql_vmin) / (float)(num_atoms - 1);

        /* Blend: readiness=0 → default [-1,1], readiness=1 → adaptive. */
        float v_min   = r * iql_vmin + (1.0f - r) * (-1.0f);
        float v_max   = r * iql_vmax + (1.0f - r) * ( 1.0f);
        float delta_z = r * iql_dz   + (1.0f - r) * default_dz;

        /* SP4 Layer B (Mech 2): bound per-sample support [v_min, v_max] to
         * ±ISV[ATOM_POS_BOUND[d]]. Per-branch p99(|atom_positions|) populated
         * by SP4 producer `launch_sp4_atom_pos_p99` (Layer A Task A6).
         * ε floor on the bound itself (`max(.., 1.0)`) per the SP1 cold-start
         * pearl. delta_z is recomputed from the clamped span so the
         * (v_min, v_max, delta_z) triple stays self-consistent. */
        if (isv_signals != nullptr) {
            const float max_atom_abs = fmaxf(isv_signals[ATOM_POS_BOUND_BASE + d], 1.0f);
            v_min = fminf(fmaxf(v_min, -max_atom_abs), max_atom_abs);
            v_max = fminf(fmaxf(v_max, -max_atom_abs), max_atom_abs);
            delta_z = (v_max - v_min) / (float)(num_atoms - 1);
        }

        int base = b * 12 + d * 3;
        per_sample_support[base + 0] = v_min;
        per_sample_support[base + 1] = v_max;
        per_sample_support[base + 2] = delta_z;
    }
}

/* ------------------------------------------------------------------ */
/*  Support Floor — Frugal p5 quantile estimator                       */
/* ------------------------------------------------------------------ */
/**
 * Ensures no sample has zero delta_z by flooring to 10% of the batch's
 * p5 quantile of half-widths. Uses Frugal-1U streaming quantile estimator
 * (one GPU scalar, O(B) per step, scale-free).
 *
 * Phase 2d layout: per_sample_support is now [B, 4, 3] stride-12; the
 * p5 estimator aggregates half-widths across ALL (sample, branch) pairs
 * (one scalar p5 for the entire tile) and the floor is applied per
 * (sample, branch) independently. Sharing a single p5 across branches
 * keeps the estimator statistically stable and preserves the anti-
 * collapse guarantee for every branch's smallest half-width.
 *
 * Launch: grid=1, block=1 (after iql_compute_per_sample_support).
 */
extern "C" __global__
void iql_support_floor(
    float* __restrict__ per_sample_support,  /* [B, 4, 3] in-place */
    float* __restrict__ p5_state,            /* [2]: [0]=p5_estimate, [1]=step_count */
    int batch_size,
    int num_atoms
)
{
    float est = p5_state[0];
    float step_count = p5_state[1];
    int total_branches = batch_size * 4;

    /* First call: init from aggregated mean spread across all (sample, branch). */
    if (step_count < 0.5f) {
        float hw_sum = 0.0f;
        for (int b = 0; b < batch_size; b++) {
            for (int d = 0; d < 4; d++) {
                int base = b * 12 + d * 3;
                float hw = (per_sample_support[base + 1] - per_sample_support[base + 0]) * 0.5f;
                hw_sum += hw;
            }
        }
        est = hw_sum / fmaxf((float)total_branches, 1.0f);
        est = fmaxf(est, 1e-6f);
    }

    /* Frugal-1U with adaptive step: 1/sqrt(step_count) → convergence guarantee.
     * Asymmetric ratio 0.05/0.95 is mathematically correct for p5 quantile. */
    float adaptive_rate = 1.0f / fmaxf(sqrtf(step_count + 1.0f), 1.0f);
    for (int b = 0; b < batch_size; b++) {
        for (int d = 0; d < 4; d++) {
            int base = b * 12 + d * 3;
            float hw = (per_sample_support[base + 1] - per_sample_support[base + 0]) * 0.5f;
            float step = fmaxf(est * adaptive_rate, 1e-8f);
            if (hw < est) {
                est -= step;
            } else {
                est += step * (0.05f / 0.95f);
            }
            est = fmaxf(est, 1e-8f);
        }
    }
    p5_state[0] = est;
    p5_state[1] = step_count + 1.0f;

    /* Floor: 1/num_atoms of p5 — guarantees at least 1 atom of resolution.
     * Adaptive to atom count: 51 atoms → 2% of p5, 11 atoms → 9% of p5.
     * Applied per (sample, branch) so each branch's floor is independent. */
    float floor_fraction = 1.0f / fmaxf((float)num_atoms, 1.0f);
    float floor_hw = est * floor_fraction;
    float floor_dz = (2.0f * floor_hw) / fmaxf((float)(num_atoms - 1), 1.0f);

    for (int b = 0; b < batch_size; b++) {
        for (int d = 0; d < 4; d++) {
            int base = b * 12 + d * 3;
            if (per_sample_support[base + 2] < floor_dz) {
                float center = (per_sample_support[base + 0] + per_sample_support[base + 1]) * 0.5f;
                per_sample_support[base + 0] = center - floor_hw;
                per_sample_support[base + 1] = center + floor_hw;
                per_sample_support[base + 2] = floor_dz;
            }
        }
    }
}

/* ------------------------------------------------------------------ */
/*  Per-Branch Advantage Scale Kernel                                   */
/* ------------------------------------------------------------------ */
/**
 * Compute per-branch normalized advantage scales for 4-branch DQN:
 *   For each branch d, marginalise Q over all joints where sub_d == a_d.
 *   a_branch[d] = |Q_marginal - V(s)| normalised by max across branches.
 *
 * Output: branch_scales[b*4+d] in [0,1] for each sample and branch.
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_per_branch_advantage(
    const float* __restrict__ q_out,         /* [B, total_actions] */
    const float* __restrict__ v_out,         /* [B*2] two heads col-major [2,B] */
    const int* __restrict__ actions,          /* [B] factored action indices */
    float* __restrict__ branch_scales,       /* [B*4] */
    const float* __restrict__ readiness_buf, /* [1] CV-based readiness */
    int batch_size,
    int total_actions,
    int b0_size, int b1_size, int b2_size, int b3_size,
    /* Per-sample branch_scales floor fix (2026-04-23): IQL advantage
     * weighting collapses the losing branch's per-sample scale to ~0
     * when readiness→1 and `a_branch[d]≈0` (e.g. direction Hold/Flat
     * samples have near-zero direction advantage). That zero multiplies
     * the C51 gradient to zero for that (b, d) pair, starving the
     * branch of learning signal. This kernel now reads the floor value
     * from ISV slot `IQL_BRANCH_SCALE_FLOOR_INDEX=36` and takes
     * `max(scale, floor)` so every branch always contributes at least
     * floor × gradient. NULL → floor is 0 (legacy behaviour).
     * See `feedback_isv_for_adaptive_bounds.md`. */
    const float* __restrict__ isv_signals,
    int iql_branch_scale_floor_idx
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    /* V(s) = V_short + V_long; col-major [2, B]: element (head, b) at head + b*2 */
    float v = v_out[b * 2 + 0] + v_out[b * 2 + 1];
    const float* q = q_out + b * total_actions;

    int factored = actions[b];
    int max_act = b0_size * b1_size * b2_size * b3_size;
    if (factored < 0 || factored >= max_act) factored = 0;
    int a0 = factored / (b1_size * b2_size * b3_size);
    int a1 = (factored / (b2_size * b3_size)) % b1_size;
    int a2 = (factored / b3_size) % b2_size;
    int a3 = factored % b3_size;
    int branch_actions[4] = { a0, a1, a2, a3 };
    int branch_sizes[4] = { b0_size, b1_size, b2_size, b3_size };

    float a_branch[4];
    float max_a = 0.0f;

    for (int d = 0; d < 4; d++) {
        int a_d = branch_actions[d];
        float q_sum = 0.0f;
        int count = 0;

        for (int joint = 0; joint < total_actions; joint++) {
            int sub_d;
            if (d == 0) sub_d = joint / (b1_size * b2_size * b3_size);
            else if (d == 1) sub_d = (joint / (b2_size * b3_size)) % b1_size;
            else if (d == 2) sub_d = (joint / b3_size) % b2_size;
            else sub_d = joint % b3_size;

            if (sub_d == a_d) {
                q_sum += q[joint];
                count++;
            }
        }

        float q_marginal = (count > 0) ? (q_sum / (float)count) : v;
        a_branch[d] = fabsf(q_marginal - v);
        max_a = fmaxf(max_a, a_branch[d]);
    }

    float r = readiness_buf[0];
    float inv_max = (max_a > 1e-8f) ? (1.0f / max_a) : 1.0f;
    /* Per-sample branch_scales floor: read from ISV, guarded on NULL so
     * the kernel falls back to pre-fix behaviour (floor = 0) when the
     * ISV bus is not wired (smoke tests). */
    float floor_val = 0.0f;
    if (isv_signals != NULL) {
        floor_val = isv_signals[iql_branch_scale_floor_idx];
        /* Defensive clamp: the floor is a safety bound, not a scale.
         * [0, 1] is the only defensible range — negative is nonsense,
         * > 1 would over-ride the IQL weighting entirely. */
        if (floor_val < 0.0f) floor_val = 0.0f;
        if (floor_val > 1.0f) floor_val = 1.0f;
    }
    for (int d = 0; d < 4; d++) {
        /* Blend: readiness=0 → uniform 0.25, readiness=1 → IQL-derived scale */
        float iql_scale = a_branch[d] * inv_max;
        float scale = r * iql_scale + (1.0f - r) * 0.25f;
        if (scale < floor_val) scale = floor_val;
        branch_scales[b * 4 + d] = scale;
    }
}

/* ------------------------------------------------------------------ */
/*  Expectile Gap Kernels                                               */
/* ------------------------------------------------------------------ */
/**
 * Compute element-wise gap between high and low IQL expectile values:
 *   gap[b] = v_high[b] - v_low[b]
 *
 * Used for per-sample adaptive epsilon computation.
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_expectile_gap(
    const float* __restrict__ v_high,   /* [B] high-expectile value estimates */
    const float* __restrict__ v_low,    /* [B] low-expectile value estimates */
    float* __restrict__ gap,             /* [B] output gaps */
    int batch_size
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;
    gap[b] = v_high[b] - v_low[b];
}

/**
 * Reduce gap array to scalar mean.
 * Single-thread kernel — launched with grid=1, block=1.
 *   gap_mean[0] = mean(gap[0..batch_size-1])
 *
 * Launch: grid=1, block=1.
 */
extern "C" __global__
void iql_gap_mean_reduce(
    const float* __restrict__ gap,       /* [B] */
    float* __restrict__ gap_mean,        /* [1] scalar output */
    int batch_size
)
{
    float sum = 0.0f;
    for (int b = 0; b < batch_size; b++)
        sum += gap[b];
    gap_mean[0] = sum / (float)batch_size;
}

/* ------------------------------------------------------------------ */
/*  Per-Sample Adaptive Epsilon Kernel                                  */
/* ------------------------------------------------------------------ */
/**
 * Compute per-sample exploration epsilon from normalised expectile gap:
 *   x   = gap[b] / max(gap_mean[0], 1e-8) - 1.0
 *   sig = sigmoid(x)
 *   per_sample_eps[b] = base_epsilon * sig
 *
 * Samples with wider value uncertainty get higher exploration probability.
 * Launch: grid=ceil(B/256), block=256.
 */
extern "C" __global__
void iql_compute_per_sample_epsilon(
    const float* __restrict__ gap,            /* [B] per-sample expectile gaps */
    const float* __restrict__ gap_mean,       /* [1] scalar mean gap */
    const float* __restrict__ readiness_buf,  /* [1] CV-based readiness */
    float* __restrict__ per_sample_eps,       /* [B] output epsilon values */
    float base_epsilon,
    int batch_size
)
{
    int b = blockIdx.x * blockDim.x + threadIdx.x;
    if (b >= batch_size) return;

    float r = readiness_buf[0];
    float gm = fmaxf(gap_mean[0], 1e-8f);
    /* Scale gap by readiness — at r=0, gap=0 → sigmoid(-1)=0.27 → uniform ~0.27*base_eps */
    float scaled_gap = r * gap[b];
    float x = scaled_gap / gm - 1.0f;
    float sig = 1.0f / (1.0f + expf(-x));
    per_sample_eps[b] = base_epsilon * sig;
}

/* ============================================================
 * cuBLAS rewrite helpers — element-wise kernels for batched GEMM pipeline
 * All kernels: zero atomicAdd, fully deterministic, extern "C" for cubin export.
 * Grid/block dims are set by the Rust launcher.
 * ============================================================ */

/* iql_silu_fwd — Element-wise SiLU activation forward
 *   out[i] = x[i] / (1 + exp(-x[i]))
 *   Launch: grid=(ceil(N/256),1,1), block=(256,1,1)
 */
extern "C" __global__
void iql_silu_fwd(
    const float* __restrict__ x,
    float* __restrict__ out,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float val = x[idx];
        out[idx] = val / (1.0f + expf(-val));
    }
}

/* iql_silu_bwd — Element-wise SiLU backward
 *   d_x[i] = d_out[i] * s*(1 + x*(1-s)),  s = sigmoid(x)
 *   Launch: grid=(ceil(N/256),1,1), block=(256,1,1)
 */
extern "C" __global__
void iql_silu_bwd(
    const float* __restrict__ x,
    const float* __restrict__ d_out,
    float* __restrict__ d_x,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float val = x[idx];
        float s = 1.0f / (1.0f + expf(-val));
        float grad = s * (1.0f + val * (1.0f - s));
        d_x[idx] = d_out[idx] * grad;
    }
}

/* iql_expectile_loss — Expectile loss + gradient (fused), horizon-decomposed.
 *
 * Operates over B*num_heads elements. For idx in [0, B*num_heads):
 *   b         = idx / num_heads   (sample index)
 *   v_out[idx] = V_head for that (sample, head) pair
 *   q_taken[b] = same Q-target for both heads of the same sample
 *
 *   loss[idx] = w*(Q[b]-V[idx])^2,  w = tau if u>=0 else (1-tau)
 *   dv[idx]   = -2*w*(Q[b]-V[idx])
 *
 * For the 2-head case (num_heads=2): v_out is [B*2] col-major [2,B];
 * both heads regress the same expectile target. The sum V(s)=V_short+V_long
 * is consumed by downstream kernels.
 *
 * Launch: grid=(ceil(B*num_heads/256),1,1), block=(256,1,1)
 */
extern "C" __global__
void iql_expectile_loss(
    const float* __restrict__ v_out,
    const float* __restrict__ q_taken,
    float tau,
    float* __restrict__ loss_out,
    float* __restrict__ dv_out,
    int B,
    int num_heads
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < B * num_heads) {
        int b = idx / num_heads;
        float v = v_out[idx];
        float q = q_taken[b];
        float u = q - v;
        float weight = (u >= 0.0f) ? tau : (1.0f - tau);
        loss_out[idx] = weight * u * u;
        dv_out[idx] = -2.0f * weight * u;
    }
}

/* iql_bias_add — Broadcast bias along batch dimension
 *   x[i] += bias[i % out_dim]
 *   Launch: grid=(ceil(N/256),1,1), block=(256,1,1),  N = B*out_dim
 */
extern "C" __global__
void iql_bias_add(
    float* __restrict__ x,
    const float* __restrict__ bias,
    int out_dim,
    int N
) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        x[idx] += bias[idx % out_dim];
    }
}

/* iql_bias_grad_reduce_p1 — Phase 1: block-level shared-memory reduce
 *   Grid: (num_blocks, out_dim), Block: 256, shared: 256*sizeof(float)
 */
extern "C" __global__
void iql_bias_grad_reduce_p1(
    const float* __restrict__ d_pre,
    float* __restrict__ partials,
    int out_dim,
    int B
) {
    extern __shared__ float sdata[];
    int j   = blockIdx.y;
    int tid = threadIdx.x;
    int gid = blockIdx.x * blockDim.x + tid;
    float val = (gid < B) ? d_pre[j + (long)gid * out_dim] : 0.0f;
    sdata[tid] = val;
    __syncthreads();
    for (int s = blockDim.x / 2; s > 0; s >>= 1) {
        if (tid < s) sdata[tid] += sdata[tid + s];
        __syncthreads();
    }
    if (tid == 0) partials[blockIdx.x * out_dim + j] = sdata[0];
}

/* iql_bias_grad_reduce_p2 — Phase 2: warp-shuffle final reduce + scale
 *   Grid: ceil(out_dim/256), Block: 256
 */
extern "C" __global__
void iql_bias_grad_reduce_p2(
    const float* __restrict__ partials,
    float* __restrict__ d_bias,
    int out_dim,
    int num_blocks,
    float inv_batch
) {
    int j = blockIdx.x * blockDim.x + threadIdx.x;
    if (j >= out_dim) return;
    float sum = 0.0f;
    for (int i = 0; i < num_blocks; i++) sum += partials[i * out_dim + j];
    d_bias[j] = sum * inv_batch;
}
