#include <metal_stdlib>
using namespace metal;

// ═══════════════════════════════════════════════════════════════
// Batch Metal Kernels - Process multiple tokens simultaneously
// ═══════════════════════════════════════════════════════════════

// Batch quantized matmul - process N tokens with shared weights
// Expected improvement: 8-15x for batch inference
kernel void quantized_matmul_batch(
    device float* inputs [[buffer(0)]],        // [batchSize, inDim]
    device uint8_t* weights [[buffer(1)]],     // [outDim, inDim] packed
    device float* scales [[buffer(2)]],        // [outDim, groups]
    device float* biases [[buffer(3)]],        // [outDim]
    device float* outputs [[buffer(4)]],       // [batchSize, outDim]
    constant uint32_t& inDim [[buffer(5)]],
    constant uint32_t& outDim [[buffer(6)]],
    constant uint32_t& groupSize [[buffer(7)]],
    constant uint32_t& batchSize [[buffer(8)]],
    uint3 gid [[thread_position_in_grid]])
{
    // Each thread processes one output dimension for one batch element
    uint batchIdx = gid.x;  // [0, batchSize)
    uint outIdx = gid.y;    // [0, outDim)
    
    if (batchIdx >= batchSize || outIdx >= outDim) return;
    
    // Get input for this batch element
    device float* input = inputs + batchIdx * inDim;
    
    // Compute dot product for this output dimension
    float sum = biases[outIdx];
    uint groupIdx = outIdx * (inDim / groupSize);
    
    for (uint i = 0; i < inDim; i += 4) {
        // Load 4 input values
        float4 inVals = float4(input[i], input[i+1], input[i+2], input[i+3]);
        
        // Load 4 packed weights (uint8 packed as uint32)
        uint packedWeight = weights[outIdx * inDim + i];
        uint8_t w0 = (packedWeight >> 0) & 0xFF;
        uint8_t w1 = (packedWeight >> 8) & 0xFF;
        uint8_t w2 = (packedWeight >> 16) & 0xFF;
        uint8_t w3 = (packedWeight >> 24) & 0xFF;
        
        // Get scale for this group
        uint g0 = (i + 0) / groupSize;
        uint g1 = (i + 1) / groupSize;
        uint g2 = (i + 2) / groupSize;
        uint g3 = (i + 3) / groupSize;
        
        float scale0 = scales[groupIdx + g0];
        float scale1 = scales[groupIdx + g1];
        float scale2 = scales[groupIdx + g2];
        float scale3 = scales[groupIdx + g3];
        
        // Dequantize and multiply
        sum += inVals.x * (w0 - 128) * scale0;
        sum += inVals.y * (w1 - 128) * scale1;
        sum += inVals.z * (w2 - 128) * scale2;
        sum += inVals.w * (w3 - 128) * scale3;
    }
    
    outputs[batchIdx * outDim + outIdx] = sum;
}

// Batch RMS norm - process N tokens simultaneously
kernel void rms_norm_batch(
    device float* inputs [[buffer(0)]],         // [batchSize, N]
    device float* weights [[buffer(1)]],        // [N]
    device float* outputs [[buffer(2)]],        // [batchSize, N]
    constant uint32_t& N [[buffer(3)]],
    constant float& eps [[buffer(4)]],
    constant uint32_t& batchSize [[buffer(5)]],
    uint3 gid [[thread_position_in_grid]])
{
    uint batchIdx = gid.x;
    uint elemIdx = gid.y;
    
    if (batchIdx >= batchSize || elemIdx >= N) return;
    
    // Compute sum of squares for this batch element
    threadgroup float sharedSqSum[256];
    uint threadIdx = elemIdx % 256;
    
    device float* input = inputs + batchIdx * N;
    
    float sqSum = 0.0;
    for (uint i = 0; i < N; i++) {
        sqSum += input[i] * input[i];
    }
    
    // RMS
    float rms = sqrt(sqSum / float(N) + eps);
    
    // Normalize
    outputs[batchIdx * N + elemIdx] = input[elemIdx] / rms * weights[elemIdx];
}

// Batch attention - process N tokens with shared KV cache
// This is the most complex batch operation
kernel void sliding_attention_batch(
    device float* queries [[buffer(0)],        // [batchSize, nHeads, headDim]
    device float* kvCache [[buffer(1)],        // [maxSeqLen, 2, nKvHeads, headDim]
    device float* outputs [[buffer(2)],        // [batchSize, nHeads, headDim]
    constant uint32_t& positions [[buffer(3)], // [batchSize]
    constant uint32_t& nHeads [[buffer(4)],
    constant uint32_t& nKvHeads [[buffer(5]],
    constant uint32_t& headDim [[buffer(6]],
    constant uint32_t& batchSize [[buffer(7]],
    constant uint32_t& windowSize [[buffer(8]],
    uint3 gid [[thread_position_in_grid]])
{
    uint batchIdx = gid.x;
    uint headIdx = gid.y;
    uint dimIdx = gid.z;
    
    if (batchIdx >= batchSize || headIdx >= nHeads || dimIdx >= headDim) return;
    
    uint pos = positions[batchIdx];
    uint kvHeadIdx = headIdx / (nHeads / nKvHeads);
    
    device float* query = queries + batchIdx * nHeads * headDim + headIdx * headDim;
    
    // Sliding window attention
    uint start = max(0u, pos - windowSize);
    uint end = min(pos, maxSeqLen);
    
    float sum = 0.0;
    float maxScore = -1e10;
    
    // Compute attention scores
    for (uint t = start; t < end; t++) {
        device float* key = kvCache + t * 2 * nKvHeads * headDim + kvHeadIdx * headDim;
        
        float score = 0.0;
        for (uint d = 0; d < headDim; d++) {
            score += query[d] * key[d];
        }
        
        score /= sqrt(float(headDim));
        maxScore = max(maxScore, score);
    }
    
    // Softmax
    float expSum = 0.0;
    for (uint t = start; t < end; t++) {
        device float* key = kvCache + t * 2 * nKvHeads * headDim + kvHeadIdx * headDim;
        
        float score = 0.0;
        for (uint d = 0; d < headDim; d++) {
            score += query[d] * key[d];
        }
        
        score /= sqrt(float(headDim));
        expSum += exp(score - maxScore);
    }
    
    // Compute weighted sum of values
    float output = 0.0;
    for (uint t = start; t < end; t++) {
        device float* key = kvCache + t * 2 * nKvHeads * headDim + kvHeadIdx * headDim;
        device float* value = kvCache + t * 2 * nKvHeads * headDim + nKvHeads * headDim + kvHeadIdx * headDim;
        
        float score = 0.0;
        for (uint d = 0; d < headDim; d++) {
            score += query[d] * key[d];
        }
        
        score /= sqrt(float(headDim));
        float weight = exp(score - maxScore) / expSum;
        
        output += weight * value[dimIdx];
    }
    
    outputs[batchIdx * nHeads * headDim + headIdx * headDim + dimIdx] = output;
}