/**
 * IQN CVaR (Conditional Value-at-Risk) position sizing kernel.
 *
 * Computes CVaR on GPU -- zero CPU readback.
 * One thread per sample: insertion-sort quantile values in registers,
 * compute CVaR as mean of lowest alpha_count values.
 *
 * Plan 4 Task 3 (E.3): N_TAU now 5 (was 32). With fixed
 * FIXED_TAUS = {0.05, 0.25, 0.50, 0.75, 0.95}, the discrete CVaR cases
 * are:
 *   ALPHA = 0.05  → alpha_count = floor(0.05 × 5) = 0 → clamped to 1
 *                  → returns the single τ=0.05 quantile (worst-5% estimator)
 *   ALPHA = 0.25  → alpha_count = floor(0.25 × 5) = 1
 *                  → again returns the τ=0.05 quantile (single worst-25%
 *                    sample given the 5-point grid; mathematically
 *                    consistent with discrete quantile aggregation)
 *   ALPHA = 0.5+  → alpha_count grows toward N_TAU; mean of bottom slots
 *
 * The `sorted[8]` register array stays oversized for the new grid (max
 * alpha_count = 5) — no resizing needed; same kernel handles the smaller
 * τ axis without modification.
 *
 * Scale: [0.25, 1.0]. CVaR >= 0 -> full size. CVaR < 0 -> reduce.
 *
 * Launch config: grid=(ceil(B/256), 1, 1), block=(256, 1, 1).
 */

extern "C" __global__ void iqn_cvar_kernel(
    const float* __restrict__ q_values,  // [B, N_TAU, TBA]
    const int*   __restrict__ actions,   // [B] factored action indices
    float* scales_out,           // [B] output CVaR scales
    int B, int N_TAU, int TBA, int B0_SIZE, int B1B2, float ALPHA
) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i >= B) return;

    // Decode exposure index from factored action
    int exposure_idx = actions[i] / B1B2;
    if (exposure_idx >= B0_SIZE) exposure_idx = B0_SIZE - 1;

    // Extract quantile values for this sample's exposure action
    // q_values layout: [B, N_TAU, TBA]
    int alpha_count = (int)(ALPHA * (float)N_TAU);
    if (alpha_count < 1) alpha_count = 1;

    // Find the alpha_count smallest quantile values using partial sort.
    float cvar_sum = 0.0f;
    if (alpha_count == 1) {
        // Fast path: just find the minimum
        float min_val = 1e30f;
        for (int t = 0; t < N_TAU; t++) {
            float v = q_values[i * N_TAU * TBA + t * TBA + exposure_idx];
            min_val = fminf(min_val, v);
        }
        cvar_sum = min_val;
    } else {
        // General path: find alpha_count smallest values.
        // sorted[8] sizing: with Plan 4 Task 3 (E.3) N_TAU=5 the maximum
        // possible alpha_count is N_TAU itself = 5 (when ALPHA=1.0). The
        // legacy 32-quantile kernel sized this for ALPHA≤0.25 → 8 slots;
        // 8 still covers the new grid completely. Kept oversized rather
        // than tightened to 5 so future N_TAU bumps below 8 don't require
        // a kernel-source edit.
        float sorted[8];
        int sorted_len = 0;
        for (int t = 0; t < N_TAU; t++) {
            float v = q_values[i * N_TAU * TBA + t * TBA + exposure_idx];
            if (sorted_len < alpha_count) {
                // Insert into sorted array
                int pos = sorted_len;
                while (pos > 0 && sorted[pos-1] > v) {
                    sorted[pos] = sorted[pos-1];
                    pos--;
                }
                sorted[pos] = v;
                sorted_len++;
            } else if (v < sorted[sorted_len-1]) {
                // Replace the largest in sorted
                int pos = sorted_len - 1;
                while (pos > 0 && sorted[pos-1] > v) {
                    sorted[pos] = sorted[pos-1];
                    pos--;
                }
                sorted[pos] = v;
            }
        }
        for (int j = 0; j < alpha_count; j++) cvar_sum = cvar_sum + sorted[j];
    }

    float cvar = cvar_sum / (float)alpha_count;

    // Scale: [0.25, 1.0]. CVaR >= 0 -> full size. CVaR < 0 -> reduce.
    float scale = (cvar >= 0.0f)
        ? 1.0f
        : fmaxf(0.25f, 1.0f + cvar * 5.0f);
    scales_out[i] = scale;
}
