/**
 * Backward pass for Multi-Head Feature-Level Self-Attention.
 *
 * Backpropagates gradients through:
 *   1. Recompute forward pass (concat, pre_ln) from saved_input + params
 *   2. LayerNorm backward (d_gamma, d_beta, chain through normalization)
 *   3. Residual split (d_prenorm -> d_projection + d_residual)
 *   4. Output projection backward (W_O, b_O gradients)
 *   5. Multi-head attention backward (Q, K, V, softmax, attention value)
 *   6. Input projection backward (W_Q, W_K, W_V, bias gradients)
 *
 * Weight layout (same as forward):
 *   W_Q:      [D, D]     offset 0
 *   W_K:      [D, D]     offset D*D
 *   W_V:      [D, D]     offset 2*D*D
 *   b_Q:      [D]        offset 3*D*D
 *   b_K:      [D]        offset 3*D*D + D
 *   b_V:      [D]        offset 3*D*D + 2*D
 *   W_O:      [D, D]     offset 3*D*D + 3*D
 *   b_O:      [D]        offset 4*D*D + 3*D
 *   ln_gamma: [D]        offset 4*D*D + 4*D
 *   ln_beta:  [D]        offset 4*D*D + 5*D
 *   Total: 4*D^2 + 5*D
 *
 * Grid: (B, 1, 1), Block: (32, 1, 1) -- one warp per sample.
 * Weight gradients accumulated via atomicAdd across batch dimension.
 *
 * The kernel recomputes the full forward pass from saved_input + params
 * to recover all intermediate activations (concat, pre_ln). This avoids
 * storing large intermediate buffers and is acceptable since the attention
 * layer runs once per training step.
 */

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

/**
 * @param d_output     [B, D] gradient from downstream (DQN trunk backward d_h_s2)
 * @param states_input [B, D] saved input states from forward pass
 * @param params       [total_params] attention weights (read-only)
 * @param d_input      [B, D] gradient to upstream (zeroed before launch, atomicAdd within)
 * @param d_params     [total_params] weight gradients (zeroed before launch, atomicAdd across batch)
 * @param B            batch size
 */
extern "C" __global__ void attention_backward_kernel(
    const float* __restrict__ d_output,
    const float* __restrict__ states_input,
    const float* __restrict__ params,
    float*       d_input,
    float*       d_params,
    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_input + sample * D;
    const float* dy = d_output + sample * D;
    float* dx = d_input + sample * D;

    /* Weight offsets in flat params buffer (same layout as forward) */
    const float* W_Q = params;
    const float* W_K = W_Q + D * D;
    const float* W_V = W_K + D * D;
    const float* b_Q = W_V + D * D;
    const float* b_K = b_Q + D;
    const float* b_V = b_K + D;
    const float* W_O = b_V + D;
    const float* b_O = W_O + D * D;
    const float* ln_gamma = b_O + D;
    /* ln_beta = ln_gamma + D; (not needed in backward computation) */

    /* Gradient offsets in d_params (same layout) */
    float* dW_Q = d_params;
    float* dW_K = dW_Q + D * D;
    float* dW_V = dW_K + D * D;
    float* db_Q = dW_V + D * D;
    float* db_K = db_Q + D;
    float* db_V = db_K + D;
    float* dW_O = db_V + D;
    float* db_O_g = dW_O + D * D;
    float* d_ln_gamma = db_O_g + D;
    float* d_ln_beta = d_ln_gamma + D;

    /* Shared memory layout: concat[D] + d_proj[D] */
    __shared__ float concat[128];  /* recomputed multi-head attention output */
    __shared__ float d_proj[128];  /* d_projection = d_prenorm from LN backward */

    /* ===================================================================
     * Step 1: Recompute forward pass to recover concat and pre_ln
     *
     * Same logic as forward kernel: for each head, compute Q, K, V,
     * attention scores, softmax, weighted V. Accumulate into concat[D].
     * Then compute pre_ln = x + W_O @ concat + b_O.
     * =================================================================== */

    for (int h = 0; h < H; h++) {
        int head_start = h * Dh;

        /* Compute Q, K for attention scores */
        for (int f = tid; f < Dh; f += 32) {
            int gf = head_start + f;
            float q = b_Q[gf];
            float k = b_K[gf];
            for (int j = 0; j < D; j++) {
                float xj = x[j];
                q += xj * W_Q[j * D + gf];
                k += xj * W_K[j * D + gf];
            }
            concat[gf] = q * k / sqrtf((float)Dh);
        }
        __syncthreads();

        /* Softmax over head */
        float max_s = -1e30f;
        for (int f = tid; f < Dh; f += 32) {
            float s = 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(concat[head_start + f] - max_s);
            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 as concat */
        for (int f = tid; f < Dh; f += 32) {
            float attn = concat[head_start + f] * inv_sum;
            int gf = head_start + f;
            float v = b_V[gf];
            for (int j = 0; j < D; j++)
                v += x[j] * W_V[j * D + gf];
            concat[gf] = attn * v;
        }
        __syncthreads();
    }

    /* Now concat[0..D) is the recomputed multi-head attention output.
     * Compute pre_ln[f] = x[f] + sum_j(concat[j] * W_O[j*D+f]) + b_O[f]
     * Store pre_ln in d_proj temporarily (we'll overwrite with d_projection later). */
    for (int f = tid; f < D; f += 32) {
        float proj = b_O[f];
        for (int j = 0; j < D; j++)
            proj += concat[j] * W_O[j * D + f];
        d_proj[f] = x[f] + proj;  /* pre_ln = states + projection */
    }
    __syncthreads();

    /* ===================================================================
     * Step 2: LayerNorm backward
     *
     * Forward was: out = gamma * (pre_ln - mean) / std + beta
     * pre_ln is in d_proj[0..D), dy is the upstream gradient.
     * =================================================================== */

    /* Compute mean of pre_ln */
    float local_mean = 0.0f;
    for (int f = tid; f < D; f += 32)
        local_mean += d_proj[f];  /* d_proj holds pre_ln here */
    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);

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

    /* LN backward:
     * d_prenorm[f] = (1/D) * inv_std * (D * dy[f]*gamma[f]
     *               - sum(dy*gamma) - x_hat[f] * sum(dy*gamma*x_hat))
     * where x_hat = (pre_ln - mean) * inv_std */

    /* First pass: accumulate sum(dy*gamma) and sum(dy*gamma*x_hat) */
    float sum_dy_gamma = 0.0f;
    float sum_dy_gamma_xhat = 0.0f;
    for (int f = tid; f < D; f += 32) {
        float x_hat = (d_proj[f] - mean) * inv_std;  /* d_proj holds pre_ln */
        float dy_g = dy[f] * ln_gamma[f];
        sum_dy_gamma += dy_g;
        sum_dy_gamma_xhat += dy_g * x_hat;
    }
    for (int mask = 16; mask >= 1; mask >>= 1) {
        sum_dy_gamma += __shfl_xor_sync(0xFFFFFFFF, sum_dy_gamma, mask);
        sum_dy_gamma_xhat += __shfl_xor_sync(0xFFFFFFFF, sum_dy_gamma_xhat, mask);
    }
    sum_dy_gamma = __shfl_sync(0xFFFFFFFF, sum_dy_gamma, 0);
    sum_dy_gamma_xhat = __shfl_sync(0xFFFFFFFF, sum_dy_gamma_xhat, 0);

    /* Second pass: compute d_prenorm, accumulate LN param grads.
     * Overwrite d_proj with d_prenorm (we're done with pre_ln). */
    for (int f = tid; f < D; f += 32) {
        float x_hat = (d_proj[f] - mean) * inv_std;
        float dy_g = dy[f] * ln_gamma[f];

        float d_prenorm = inv_std / (float)D *
            ((float)D * dy_g - sum_dy_gamma - x_hat * sum_dy_gamma_xhat);

        d_proj[f] = d_prenorm;  /* overwrite: now d_proj holds d_prenorm */

        /* LN parameter gradients */
        atomicAdd(&d_ln_gamma[f], dy[f] * x_hat);
        atomicAdd(&d_ln_beta[f], dy[f]);
    }
    __syncthreads();

    /* Now d_proj[0..D) = d_prenorm = d_projection (through residual split).
     *
     * ===================================================================
     * Step 3: Residual split
     *
     * Forward: pre_ln = x + projection
     * d_projection = d_prenorm   (already in d_proj)
     * d_residual   = d_prenorm   (same gradient goes to both paths)
     * =================================================================== */

    /* Initialize d_input with residual gradient */
    for (int f = tid; f < D; f += 32)
        atomicAdd(&dx[f], d_proj[f]);  /* d_residual via atomicAdd (heads also add) */
    __syncthreads();

    /* ===================================================================
     * Step 4: Output projection backward
     *
     * Forward: projection[f] = sum_j(concat[j] * W_O[j*D + f]) + b_O[f]
     * d_concat[j] = sum_f(d_proj[f] * W_O[j*D + f])
     * dW_O[j*D+f] += concat[j] * d_proj[f]
     * db_O[f] += d_proj[f]
     *
     * concat[0..D) still holds the recomputed attention output from Step 1.
     * d_proj[0..D) holds d_projection.
     * =================================================================== */

    /* b_O gradient */
    for (int f = tid; f < D; f += 32)
        atomicAdd(&db_O_g[f], d_proj[f]);

    /* W_O gradient and d_concat computation.
     * We need d_concat in shared memory for the next step.
     * Repurpose: compute d_concat into a register per feature, then
     * swap shared buffers. Since we need both concat and d_concat
     * simultaneously, and we have d_proj available, we compute d_concat
     * into d_proj (overwriting d_projection which we no longer need). */

    /* Phase 1: accumulate W_O gradients (concat[j] * d_proj[f]) */
    for (int f = tid; f < D; f += 32) {
        float dp = d_proj[f];
        for (int j = 0; j < D; j++)
            atomicAdd(&dW_O[j * D + f], concat[j] * dp);
    }
    __syncthreads();

    /* Phase 2: compute d_concat and store back in d_proj */
    for (int j = tid; j < D; j += 32) {
        float d_c = 0.0f;
        for (int f = 0; f < D; f++)
            d_c += d_proj[f] * W_O[j * D + f];
        /* Temporarily store in register; we need to sync before overwriting d_proj */
        concat[j] = d_c;  /* repurpose concat as d_concat temporarily */
    }
    __syncthreads();

    /* Copy d_concat from concat back to d_proj, restore concat for head backward */
    for (int f = tid; f < D; f += 32)
        d_proj[f] = concat[f];  /* d_proj now holds d_concat */
    __syncthreads();

    /* ===================================================================
     * Step 5: Multi-head attention backward (per head)
     *
     * For each head, recompute Q, K, V, attn (same as Step 1).
     * Then compute:
     *   d_attn[f] = d_concat[gf] * V[f]
     *   d_V[f]    = d_concat[gf] * attn[f]
     *   d_score   = softmax_backward(d_attn, attn)
     *   d_score_scaled = d_score / sqrt(Dh)
     *   d_Q[f]    = d_score_scaled * K[f]
     *   d_K[f]    = d_score_scaled * Q[f]
     * =================================================================== */

    for (int h = 0; h < H; h++) {
        int head_start = h * Dh;

        /* Recompute Q, K, V and attention for this head.
         * Use concat[head_start..] as scratch for attn values. */
        for (int f = tid; f < Dh; f += 32) {
            int gf = head_start + f;
            float q = b_Q[gf];
            float k = b_K[gf];
            for (int j = 0; j < D; j++) {
                float xj = x[j];
                q += xj * W_Q[j * D + gf];
                k += xj * W_K[j * D + gf];
            }
            concat[gf] = q * k / sqrtf((float)Dh);
        }
        __syncthreads();

        /* Softmax recompute */
        float max_s = -1e30f;
        for (int f = tid; f < Dh; f += 32) {
            float s = 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 sm_sum = 0.0f;
        for (int f = tid; f < Dh; f += 32) {
            float e = expf(concat[head_start + f] - max_s);
            concat[head_start + f] = e;
            sm_sum += e;
        }
        for (int mask = 16; mask >= 1; mask >>= 1)
            sm_sum += __shfl_xor_sync(0xFFFFFFFF, sm_sum, mask);
        float inv_sum = __shfl_sync(0xFFFFFFFF, 1.0f / (sm_sum + 1e-8f), 0);

        /* concat[gf] now has unnorm exp. attn[f] = concat[gf] * inv_sum */

        /* Compute sum(d_attn * attn) for softmax backward */
        float sum_da_a = 0.0f;
        for (int f = tid; f < Dh; f += 32) {
            int gf = head_start + f;
            float attn = concat[gf] * inv_sum;

            /* Recompute V */
            float v = b_V[gf];
            for (int j = 0; j < D; j++)
                v += x[j] * W_V[j * D + gf];

            float d_attn = d_proj[gf] * v;  /* d_proj holds d_concat */
            sum_da_a += d_attn * attn;
        }
        for (int mask = 16; mask >= 1; mask >>= 1)
            sum_da_a += __shfl_xor_sync(0xFFFFFFFF, sum_da_a, mask);
        sum_da_a = __shfl_sync(0xFFFFFFFF, sum_da_a, 0);

        /* Compute d_Q, d_K, d_V and accumulate weight/bias/input gradients */
        for (int f = tid; f < Dh; f += 32) {
            int gf = head_start + f;
            float attn = concat[gf] * inv_sum;

            /* Recompute Q, K, V */
            float q = b_Q[gf];
            float k = b_K[gf];
            float v = b_V[gf];
            for (int j = 0; j < D; j++) {
                float xj = x[j];
                q += xj * W_Q[j * D + gf];
                k += xj * W_K[j * D + gf];
                v += xj * W_V[j * D + gf];
            }

            float d_attn_v = d_proj[gf];  /* d_concat[gf] */
            float d_attn = d_attn_v * v;
            float d_v = d_attn_v * attn;

            /* Softmax backward: d_score = attn * (d_attn - sum(d_attn*attn)) */
            float d_score = attn * (d_attn - sum_da_a);
            float d_score_scaled = d_score / sqrtf((float)Dh);

            /* d_Q = d_score_scaled * K */
            float d_q = d_score_scaled * k;
            /* d_K = d_score_scaled * Q */
            float d_k = d_score_scaled * q;

            /* -- Bias gradients -- */
            atomicAdd(&db_Q[gf], d_q);
            atomicAdd(&db_K[gf], d_k);
            atomicAdd(&db_V[gf], d_v);

            /* -- Weight gradients: dW[j*D+gf] += x[j] * d_proj -- */
            for (int j = 0; j < D; j++) {
                float xj = x[j];
                atomicAdd(&dW_Q[j * D + gf], xj * d_q);
                atomicAdd(&dW_K[j * D + gf], xj * d_k);
                atomicAdd(&dW_V[j * D + gf], xj * d_v);
            }

            /* -- Input gradients: dx[j] += d_q*W_Q + d_k*W_K + d_v*W_V -- */
            for (int j = 0; j < D; j++) {
                float grad_j = d_q * W_Q[j * D + gf]
                             + d_k * W_K[j * D + gf]
                             + d_v * W_V[j * D + gf];
                atomicAdd(&dx[j], grad_j);
            }
        }
        __syncthreads();
    }
}

/* =====================================================================
 * ATTENTION GRADIENT NORM KERNEL
 *
 * Computes L2 norm (as sqrt of sum of squares) of attention parameter
 * gradients. Same pattern as iqn_grad_norm_kernel.
 *
 * Launch: grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1)
 * ===================================================================== */
extern "C" __global__ void attn_grad_norm_kernel(
    const float* __restrict__ grads,
    float* __restrict__ norm_out,  /* [1] */
    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;
    }

    __shared__ float shared[256];
    shared[threadIdx.x] = partial;
    __syncthreads();

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

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

/* =====================================================================
 * ATTENTION ADAM UPDATE KERNEL
 *
 * AdamW with gradient clipping for attention parameters.
 * Same algorithm as iqn_adam_kernel.
 *
 * Launch: grid=(ceil(total_params/256), 1, 1), block=(256, 1, 1)
 * ===================================================================== */
extern "C" __global__ void attn_adam_kernel(
    float* __restrict__ params,
    float* __restrict__ grads,
    float* __restrict__ m,
    float* __restrict__ v,
    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,
    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: weight decay applied to params directly */
    params[tid] *= (1.0f - lr * weight_decay);

    /* Moment updates */
    float m_val = beta1 * m[tid] + (1.0f - beta1) * g;
    float v_val = beta2 * v[tid] + (1.0f - beta2) * g * g;
    m[tid] = m_val;
    v[tid] = v_val;

    /* Bias correction */
    float bc1 = 1.0f - powf(beta1, (float)adam_t);
    float bc2 = 1.0f - powf(beta2, (float)adam_t);
    float m_hat = m_val / (bc1 + 1e-12f);
    float v_hat = v_val / (bc2 + 1e-12f);

    /* Parameter update */
    params[tid] -= lr * m_hat / (sqrtf(v_hat) + eps);

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