#ifndef __ATTENTION_CUH__
#define __ATTENTION_CUH__

#include <cuda_runtime.h>
#include "../utils/cuda_utils.cuh"
#include "softmax.cuh"

__global__ void permute_kernel(float* q, float* k, float* v, const float* inp, int B, int N, int NH, int d) {
    // Implement this

    int x = blockDim.x * blockIdx.x + threadIdx.x;

    if( x < B*NH*N*d){
        int b = (x / d / N / NH);
        int nh = (x / d / N) % NH;
        int t = (x / d) % N;
        int d_ = x % d;

        int inp_idx = (b * N * 3 * NH * d) + (t * 3 * NH * d) + (0 * NH * d) + (nh * d) + d_;
        q[b * NH * N * d + nh * N * d + t * d + d_] = inp[inp_idx];
        k[b * NH * N * d + nh * N * d + t * d + d_] = inp[inp_idx + NH * d];
        v[b * NH * N * d + nh * N * d + t * d + d_] = inp[inp_idx + 2 * (NH * d)];
    }
}


__global__ void unpermute_kernel(float* inp, float *out, int B, int N, int NH, int d) {
    // Implement this
    int x = blockDim.x * blockIdx.x + threadIdx.x;

    if( x < B*NH*N*d){
        int b = (x / d / N / NH);
        int nh = (x / d / N) % NH;
        int t = (x / d) % N;
        int d_ = x % d;

        out[b*N*NH*d + t*NH*d + nh*d + d_] = inp[b*NH*N*d + nh*N*d + t*d + d_];
    }

}

__global__ void preattention_scores_kernel(float* q, float* k, int B, int T, int NH, int HS, float* preatt){
    int x = blockDim.x * blockIdx.x + threadIdx.x;

    if(x < B* T* NH* T){
        int b = (x / T / T / NH);
        int nh = (x / T / T) % NH;
        int t1 = (x / T) % T;
        int t2 = x % T;

        float sum = 0.0f;
        for (int i = 0; i < HS; i++) {
            sum += q[b*NH*T*HS + nh*T*HS + t1*HS + i] * k[b*NH*T*HS + nh*T*HS + t2*HS + i];
        }
        preatt[b * NH * T * T + nh * T * T + t1 * T + t2] = sum;
    }   
}

__global__ void attention_v_kernel(float* vaccum, float* att, float* v, int B, int T, int NH, int HS){
    int x = blockDim.x * blockIdx.x + threadIdx.x;

    if(x < B*NH*T*HS){
        int b   = (x / HS / T / NH);
        int nh  = (x / HS / T) % NH;
        int t   = (x / HS) % T;
        int hs  = x % HS;
        float sum = 0.0f;
        for (int t2 = 0; t2 < T; t2++) {
            sum += att[b*NH*T*T + nh*T*T + t*T + t2] * v[b*NH*T*HS + nh*T*HS + t2*HS + hs];
        }

        vaccum[b * NH * T * HS + nh * T * HS + t * HS + hs] = sum;
    }
}

// Launch all kernels related to attention here 
void attention_forward(float* out, float* qkvr, float* att, float* inp, int B, int T, int C, int NH) {
    // Implement this
    int HS = C / NH;

    float *q, *k, *v;
    q = qkvr + 0 * B * T * C;
    k = qkvr + 1 * B * T * C;
    v = qkvr + 2 * B * T * C;

    dim3 block(32,1,1);
    dim3 grid((B * NH * T * HS + block.x-1) / block.x,1,1);

    permute_kernel<<<grid, block>>>(q, k, v, inp, B, T, NH, HS);

    float* preatt = inp;
    dim3 grid_preatt((B * NH * T * T + block.x-1) / block.x,1,1);
    preattention_scores_kernel<<<grid_preatt, block>>>(q, k, B, T, NH, HS, preatt);

    float scale = 1.0 / sqrtf(HS);
    int softmax_rows = B * NH * T;
    softmax_forward_kernel<<<softmax_rows, 256>>>(att, scale, preatt, B * NH, T);
    
    float* vaccum = inp;
    attention_v_kernel<<<grid_preatt,block>>>(vaccum, att, v, B, T, NH, HS);

    unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
}

#endif // __ATTENTION_CUH__