/**
 * 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)
 *
 * Architecture choices:
 *   H = 4 heads (standard for small dimensions)
 *   D_head = D / H = 72 / 4 = 18 per head
 *   Total attention params: 3 * D * D + D * D + 2 * D = 4 * D^2 + 2D
 *     For D=72: 4 * 5184 + 144 = 20,880 params (~82KB)
 *
 * Grid: (B, 1, 1), Block: (32, 1, 1). One warp per sample.
 * Maximum D = 128 (register pressure). For D=72, well within limits.
 *
 * This teaches the model WHICH features matter for the current state:
 *   - Trending market: ADX + momentum heads activate
 *   - Volatile market: CUSUM + volatility heads activate
 *   - In a trade: portfolio context heads activate (drawdown, hold_time)
 *   - Multi-timeframe: higher timeframe heads provide context
 */

/* ── 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)

/**
 * 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
 *
 * Total params: 3*D*D + 3*D + D*D + D + D + D = 4*D^2 + 5*D
 *   For D=72: 4*5184 + 360 = 21,096 params
 *
 * @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,
    int B
) {
    const int D = ATTN_STATE_DIM;
    const int H = ATTN_NUM_HEADS;
    const int Dh = ATTN_HEAD_DIM;

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

    const float* x = states + sample * D;
    float* out = output + sample * 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] */

    /* ── Step 1: Compute Q, K, V projections ── */
    /* Each thread handles D/32 ≈ 2-3 features */
    float q_local[4]; /* max Dh per head that this thread handles */
    float k_local[4];
    float v_local[4];
    float head_out[4]; /* attention output per head */

    /* Accumulate multi-head output into a shared buffer */
    __shared__ float shared_concat[128]; /* max D */

    /* 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 += 32) {
            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 += xj * W_Q[j * D + global_f];
                k += xj * W_K[j * D + global_f];
                v += xj * W_V[j * D + global_f];
            }

            /* 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 += 32) {
            float s = shared_concat[head_start + f];
            if (s > max_s) max_s = s;
        }
        for (int mask = 16; mask >= 1; mask >>= 1)
            max_s = fmaxf(max_s, __shfl_xor_sync(0xFFFFFFFF, max_s, mask));
        max_s = __shfl_sync(0xFFFFFFFF, max_s, 0);

        float local_sum = 0.0f;
        for (int f = tid; f < Dh; f += 32) {
            float e = expf(shared_concat[head_start + f] - max_s);
            shared_concat[head_start + f] = e;
            local_sum += e;
        }
        for (int mask = 16; mask >= 1; mask >>= 1)
            local_sum += __shfl_xor_sync(0xFFFFFFFF, local_sum, mask);
        float inv_sum = __shfl_sync(0xFFFFFFFF, 1.0f / (local_sum + 1e-8f), 0);

        /* Apply attention to V, store in shared_concat */
        for (int f = tid; f < Dh; f += 32) {
            float attn = shared_concat[head_start + f] * inv_sum;

            /* Recompute V[f] */
            int global_f = head_start + f;
            float v = b_V[global_f];
            for (int j = 0; j < D; j++) {
                v += x[j] * W_V[j * D + global_f];
            }
            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 += 32) {
        float proj = b_O[f];
        for (int j = 0; j < D; j++) {
            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 += 32) local_mean += out[f];
    for (int mask = 16; mask >= 1; mask >>= 1)
        local_mean += __shfl_xor_sync(0xFFFFFFFF, local_mean, mask);
    float mean = __shfl_sync(0xFFFFFFFF, local_mean / (float)D, 0);

    /* variance = avg((out[f] - mean)^2) */
    float local_var = 0.0f;
    for (int f = tid; f < D; f += 32) {
        float diff = out[f] - mean;
        local_var += diff * diff;
    }
    for (int mask = 16; mask >= 1; mask >>= 1)
        local_var += __shfl_xor_sync(0xFFFFFFFF, local_var, mask);
    float inv_std = __shfl_sync(0xFFFFFFFF, 1.0f / sqrtf(local_var / (float)D + 1e-5f), 0);

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

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