/**
 * Multi-Head Feature-Level Self-Attention for DQN State Preprocessing.
 *
 * Computes H-head self-attention over the feature dimension:
 *
 *   For each head h (0..H-1):
 *     Q_h = state @ W_Q_h    [B, D] @ [D, D_head] -> [B, D_head]
 *     K_h = state @ W_K_h    [B, D] @ [D, D_head] -> [B, D_head]
 *     V_h = state @ W_V_h    [B, D] @ [D, D_head] -> [B, D_head]
 *     attn_h = softmax(Q_h @ K_h^T / sqrt(D_head))  [B, D_head]
 *     head_h = V_h * attn_h                           [B, D_head]
 *
 *   Concat all heads -> [B, D] (H * D_head = D)
 *   Output projection: concat @ W_O + b_O -> [B, D]
 *   Residual: output = state + projected
 *   LayerNorm: output = layernorm(output)
 *
 * Grid: (B, 1, 1), Block: (256, 1, 1). 256 threads (8 warps) per sample.
 */

/* -- Configuration constants (overridden by NVRTC injection) ---------- */
#ifndef ATTN_NUM_HEADS
#define ATTN_NUM_HEADS 4
#endif
#ifndef ATTN_STATE_DIM
#define ATTN_STATE_DIM 72
#endif
#define ATTN_HEAD_DIM (ATTN_STATE_DIM / ATTN_NUM_HEADS)

/* Block-level max reduction: 256 threads -> single max value.
 * Returns the result broadcast to ALL threads in the block. */
__device__ __forceinline__ float block_max_f(
    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 max */
    for (int offset = 16; offset > 0; offset >>= 1)
        val = fmaxf(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] : -1e30f;
        for (int off = 16; off > 0; off >>= 1)
            v = fmaxf(v, __shfl_xor_sync(0xFFFFFFFF, v, off));
        warp_sums[0] = v;
    }
    __syncthreads();
    return warp_sums[0];
}

/* Block-level sum reduction: 256 threads -> single sum.
 * Returns the result broadcast to ALL threads in the block. */
__device__ __forceinline__ float block_sum_f(
    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 */
    for (int offset = 16; offset > 0; offset >>= 1)
        val = 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 = v + __shfl_xor_sync(0xFFFFFFFF, v, off);
        warp_sums[0] = v;
    }
    __syncthreads();
    return warp_sums[0];
}

/**
 * Forward pass: multi-head self-attention + residual + layer norm.
 *
 * Weight layout (flat buffer):
 *   W_QKV: [3, D, D] -- packed Q, K, V projections
 *   b_QKV: [3, D]    -- packed biases
 *   W_O:   [D, D]    -- output projection
 *   b_O:   [D]       -- output bias
 *   ln_gamma: [D]    -- layer norm scale
 *   ln_beta:  [D]    -- layer norm shift
 *
 * @param states       [B, D] input
 * @param params       flat weight buffer (layout above)
 * @param output       [B, D] output
 * @param B            batch size
 */
extern "C" __global__ void multihead_feature_attention(
    const float* __restrict__ states,
    const float* __restrict__ params,
    float*       output,
    float*       saved_qkv,  /* [B, 3, D] save buffer for Q, K, V projections */
    int B,
    int state_dim,       /* runtime state dimension (replaces ATTN_STATE_DIM) */
    int num_heads        /* runtime num_heads (replaces ATTN_NUM_HEADS) */
) {
    const int D = state_dim;
    const int H = num_heads;
    const int Dh = D / H;

    int sample = blockIdx.x;
    if (sample >= B) return;
    int tid = threadIdx.x; /* 0-255 */

    const float* x = states + sample * D;
    float* out = output + sample * D;
    /* Save buffer pointers: Q, K, V each [D] per sample */
    float* save_Q = saved_qkv + sample * 3 * D;
    float* save_K = save_Q + D;
    float* save_V = save_K + D;

    /* Weight offsets in flat params buffer */
    const float* W_Q = params;                        /* [D, D] */
    const float* W_K = W_Q + D * D;                   /* [D, D] */
    const float* W_V = W_K + D * D;                   /* [D, D] */
    const float* b_Q = W_V + D * D;                   /* [D] */
    const float* b_K = b_Q + D;                        /* [D] */
    const float* b_V = b_K + D;                        /* [D] */
    const float* W_O = b_V + D;                        /* [D, D] */
    const float* b_O = W_O + D * D;                    /* [D] */
    const float* ln_gamma = b_O + D;                   /* [D] */
    const float* ln_beta = ln_gamma + D;               /* [D] */

    /* Shared memory layout: concat[D] + reduction scratch [8 floats] */
    extern __shared__ float shmem_fwd[];
    float* shared_concat = shmem_fwd;          /* [D] */
    float* warp_reduce   = shmem_fwd + D;      /* [8] for block reductions */

    /* Process each head sequentially (H=4, each head has Dh=18 dims) */
    for (int h = 0; h < H; h++) {
        int head_start = h * Dh;

        /* Compute Q_h, K_h, V_h for features this thread handles */
        for (int f = tid; f < Dh; f += 256) {
            int global_f = head_start + f;
            float q = b_Q[global_f];
            float k = b_K[global_f];
            float v = b_V[global_f];
            for (int j = 0; j < D; j++) {
                float xj = x[j];
                q = q + xj * W_Q[j * D + global_f];
                k = k + xj * W_K[j * D + global_f];
                v = v + xj * W_V[j * D + global_f];
            }

            /* Save Q, K, V projections for backward pass */
            save_Q[global_f] = q;
            save_K[global_f] = k;
            save_V[global_f] = v;

            /* Attention score: Q[f] * K[f] / sqrt(Dh) */
            float score = (q * k) / sqrtf((float)Dh);

            /* Store for softmax */
            shared_concat[head_start + f] = score;
        }
        __syncthreads();

        /* Softmax over this head's Dh dimensions */
        float max_s = -1e30f;
        for (int f = tid; f < Dh; f += 256) {
            float s = shared_concat[head_start + f];
            max_s = fmaxf(max_s, s);
        }
        max_s = block_max_f(max_s, warp_reduce);

        float local_sum = 0.0f;
        for (int f = tid; f < Dh; f += 256) {
            float e = expf(shared_concat[head_start + f] - max_s);
            shared_concat[head_start + f] = e;
            local_sum = local_sum + e;
        }
        float total_sum = block_sum_f(local_sum, warp_reduce);
        float inv_sum = 1.0f / (total_sum + 1e-8f);

        /* Apply attention to V (read from saved projection), store in shared_concat */
        for (int f = tid; f < Dh; f += 256) {
            float attn = shared_concat[head_start + f] * inv_sum;
            int global_f = head_start + f;
            float v = save_V[global_f];  /* read saved V instead of recomputing */
            shared_concat[head_start + f] = attn * v;
        }
        __syncthreads();
    }

    /* -- Step 2: Output projection + residual -- */
    /* out[f] = x[f] + sum_j(concat[j] * W_O[j * D + f]) + b_O[f] */
    for (int f = tid; f < D; f += 256) {
        float proj = b_O[f];
        for (int j = 0; j < D; j++) {
            proj = proj + shared_concat[j] * W_O[j * D + f];
        }
        out[f] = x[f] + proj; /* residual connection */
    }
    __syncthreads();

    /* -- Step 3: Layer normalization -- */
    /* mean = avg(out[f]) */
    float local_mean = 0.0f;
    for (int f = tid; f < D; f += 256) local_mean = local_mean + out[f];
    float total_mean = block_sum_f(local_mean, warp_reduce);
    float mean = total_mean / (float)D;

    /* variance = avg((out[f] - mean)^2) */
    float local_var = 0.0f;
    for (int f = tid; f < D; f += 256) {
        float diff = out[f] - mean;
        local_var = local_var + diff * diff;
    }
    float total_var = block_sum_f(local_var, warp_reduce);
    float inv_std = 1.0f / sqrtf(total_var / (float)D + 1e-5f);

    /* Normalize + affine */
    for (int f = tid; f < D; f += 256) {
        out[f] = ln_gamma[f] * (out[f] - mean) * inv_std + ln_beta[f];
    }
}

/* Backward pass is in attention_backward_kernel.cu (separate compilation unit). */
