#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;
//     // }   
//      DONT NEED 
// }

// __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;
//     // }

//     //DONT NEED 
// }

// 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;
    float alpha = 1.0f;
    float beta = 0.0f;
    
    for(int b = 0; b < B;b ++){
        for(int nh = 0; nh < NH; nh++){
            int head = b * NH + nh;
            const float* q_head = q+ head * T * HS;
            const float* k_head = k+ head * T * HS;
            float* preatt_head = preatt + head * T * T;

            cublasSgemm(cublas_handle, CUBLAS_OP_T, CUBLAS_OP_N, T,T,HS, &alpha, k_head, HS, q_head, HS, &beta, preatt_head, T);
        }
    }

    float scale = 1.0 / sqrtf(HS);
    dim3 grid_softmax (( B * NH * T + block.x -1 )/ block.x,1,1);
    softmax_forward_kernel<<<grid_softmax, block>>>(att, scale, preatt, B * NH, T);

    float* vaccum = inp;
    for(int b = 0; b < B;b ++){
        for(int nh = 0; nh < NH; nh++){
            int head = b * NH + nh;
            const float* att_head = att + head * T * T;
            const float* v_head = v + head * T * HS;
            float* out_head = vaccum + head * T * HS;

            cublasSgemm(cublas_handle, CUBLAS_OP_N, CUBLAS_OP_N, HS,T,T, &alpha, v_head, HS, att_head, T, &beta, out_head, HS);
        }
    }


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

#endif // __ATTENTION_CUH__