#include <metal_stdlib>
using namespace metal;

// ── RoPE (Rotary Position Embedding) ──
// Applies rotary position embeddings to Q and K tensors
// Input: [seqLen, hiddenSize], positions: [seqLen]
// Output: in-place modified Q/K

kernel void apply_rope(
    device float *q           [[buffer(0)]],
    device float *k           [[buffer(1)]],
    constant uint &seqLen     [[buffer(2)]],
    constant uint &headDim    [[buffer(3)]],
    constant uint &numHeads   [[buffer(4)]],
    constant float &ropeTheta [[buffer(5)]],
    uint tid [[thread_position_in_grid]],
    uint gid [[threadgroup_position_in_grid]]
) {
    uint totalThreads = numHeads * headDim / 2;
    if (tid >= totalThreads) return;

    uint headIdx = tid / (headDim / 2);
    uint dimIdx = tid % (headDim / 2);

    // Each thread handles one (head, dim/2) pair across all positions
    for (uint pos = 0; pos < seqLen; pos++) {
        float theta = pow(ropeTheta, -2.0 * float(dimIdx) / float(headDim));
        float freq = float(pos) * theta;
        float cosFreq = cos(freq);
        float sinFreq = sin(freq);

        uint qBase = pos * numHeads * headDim + headIdx * headDim;
        uint kBase = pos * numHeads * headDim + headIdx * headDim;

        // Q rotation
        float q0 = q[qBase + dimIdx];
        float q1 = q[qBase + dimIdx + headDim / 2];
        q[qBase + dimIdx] = q0 * cosFreq - q1 * sinFreq;
        q[qBase + dimIdx + headDim / 2] = q0 * sinFreq + q1 * cosFreq;

        // K rotation
        float k0 = k[kBase + dimIdx];
        float k1 = k[kBase + dimIdx + headDim / 2];
        k[kBase + dimIdx] = k0 * cosFreq - k1 * sinFreq;
        k[kBase + dimIdx + headDim / 2] = k0 * sinFreq + k1 * cosFreq;
    }
}

// ── Q/K RMSNorm ──
// Applies RMSNorm to each head's Q/K vector

kernel void rms_norm_head(
    device const float *input   [[buffer(0)]],
    device const float *weight  [[buffer(1)]],
    device float *output        [[buffer(2)]],
    constant uint &headDim      [[buffer(3)]],
    constant uint &numHeads     [[buffer(4)]],
    constant float &eps         [[buffer(5)]],
    uint tid [[thread_position_in_grid]]
) {
    if (tid >= numHeads) return;

    uint base = tid * headDim;
    float sumSq = 0.0;
    for (uint i = 0; i < headDim; i++) {
        sumSq += input[base + i] * input[base + i];
    }
    float rms = sqrt(sumSq / float(headDim) + eps);

    for (uint i = 0; i < headDim; i++) {
        output[base + i] = input[base + i] * weight[i] / rms;
    }
}

// ── Bidirectional Sliding Window Attention ──
// Computes softmax(Q*K^T/sqrt(d)) * V with sliding window mask

kernel void bidirectional_sliding_attn(
    device const float *q           [[buffer(0)]],
    device const float *k           [[buffer(1)]],
    device const float *v           [[buffer(2)]],
    device float *output            [[buffer(3)]],
    constant uint &seqLen           [[buffer(4)]],
    constant uint &headDim          [[buffer(5)]],
    constant uint &numHeads         [[buffer(6)]],
    constant uint &numKVHeads       [[buffer(7)]],
    constant uint &slidingWindow    [[buffer(8)]],
    constant float &scale           [[buffer(9)]],
    threadgroup float *shared_mem   [[threadgroup(0)]],
    uint tid [[thread_position_in_grid]],
    uint tgSize [[threads_per_threadgroup]]
) {
    // Each thread handles one (query_position, head) pair
    uint totalQueries = seqLen * numHeads;
    if (tid >= totalQueries) return;

    uint qPos = tid / numHeads;
    uint headIdx = tid % numHeads;
    uint kvHeadIdx = headIdx * numKVHeads / numHeads;  // GQA

    float sqrtD = sqrt(float(headDim));
    uint kvBase = kvHeadIdx * headDim;

    // Compute attention scores with sliding window
    float maxScore = -1e30f;
    float scores[2048];  // max seqLen
    uint validCount = 0;

    uint windowStart = qPos > slidingWindow ? qPos - slidingWindow : 0;
    uint windowEnd = min(qPos + slidingWindow + 1, seqLen);

    for (uint kPos = windowStart; kPos < windowEnd; kPos++) {
        float dot = 0.0;
        for (uint d = 0; d < headDim; d++) {
            dot += q[qPos * numHeads * headDim + headIdx * headDim + d] *
                   k[kPos * numKVHeads * headDim + kvBase + d];
        }
        scores[validCount] = dot * scale / sqrtD;
        if (scores[validCount] > maxScore) maxScore = scores[validCount];
        validCount++;
    }

    // Softmax
    float sumExp = 0.0;
    for (uint i = 0; i < validCount; i++) {
        scores[i] = exp(scores[i] - maxScore);
        sumExp += scores[i];
    }
    if (sumExp > 0) {
        for (uint i = 0; i < validCount; i++) {
            scores[i] /= sumExp;
        }
    }

    // Weighted sum of V
    uint outBase = qPos * numHeads * headDim + headIdx * headDim;
    for (uint d = 0; d < headDim; d++) {
        float val = 0.0;
        uint kPosIdx = windowStart;
        for (uint i = 0; i < validCount; i++) {
            val += scores[i] * v[kPosIdx * numKVHeads * headDim + kvBase + d];
            kPosIdx++;
        }
        output[outBase + d] = val;
    }
}

// ── GELU ──

kernel void gelu_kernel(
    device const float *input  [[buffer(0)]],
    device float *output       [[buffer(1)]],
    constant uint &count       [[buffer(2)]],
    uint tid [[thread_position_in_grid]]
) {
    if (tid >= count) return;
    float x = input[tid];
    float absv = abs(x);
    float gelu;
    if (absv > 10.0f) {
        gelu = x > 0 ? x : 0.0f;
    } else {
        float x3 = x * x * x;
        gelu = 0.5f * x * (1.0f + tanh(0.7978845608028654f * (x + 0.044715f * x3)));
    }
    output[tid] = gelu;
}

// ── GELU + Multiply ──

kernel void gelu_mul_kernel(
    device const float *gate   [[buffer(0)]],
    device const float *up     [[buffer(1)]],
    device float *output       [[buffer(2)]],
    constant uint &count       [[buffer(3)]],
    uint tid [[thread_position_in_grid]]
) {
    if (tid >= count) return;
    float g = gate[tid];
    float absv = abs(g);
    float gelu;
    if (absv > 10.0f) {
        gelu = g > 0 ? g : 0.0f;
    } else {
        float g3 = g * g * g;
        gelu = 0.5f * g * (1.0f + tanh(0.7978845608028654f * (g + 0.044715f * g3)));
    }
    output[tid] = gelu * up[tid];
}