#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_];
    }

}

#define PREATT_T 128
#define PREATT_U 8
#define PREATT_S 16

__global__ void preattention_scores_kernel(float* q, float* k, int B, int T, int NH, int HS, float* preatt){
    
    // each block handles a (b,nh) head and a tile of (t1, t2)
    int head = blockIdx.z; //whcih pair
    int b = head / NH;
    int nh = head % NH;

    int t1 = blockIdx.x * PREATT_T + threadIdx.x; // row of output query positin
    int t2 = blockIdx.y * PREATT_U; //start column key position

    //head pointers 
    const float* q_head = q + (b * NH + nh) * T * HS; //[T,HS]
    const float* k_head = k + (b * NH + nh) * T * HS; //[T,HS]
    float* preatt_head = preatt + (b * NH +nh) * T * T; //[T,T]

    //reg tile for accumulation 
    float accumulation[PREATT_U] = {0.0f};

    __shared__ float k_tile[PREATT_S][PREATT_U];

    for(int tile_start = 0; tile_start < HS; tile_start+=PREATT_S){
        // shared loading
        // K is [T, HS], K[t2..t2+U, tile_start..tile_start+S]
        
        int s_l = threadIdx.x / PREATT_U;
        int u_l = threadIdx.x % PREATT_U;

        int hs_idx = tile_start + s_l;
        int t2_idx = t2 + u_l;

        if (hs_idx < HS && t2_idx < T) {
            k_tile[s_l][u_l] = k_head[t2_idx * HS + hs_idx];
        } else {
            k_tile[s_l][u_l] = 0.0f;
        }

        __syncthreads();

        //each threads now loads its row of queries from global memory into registers and accumulates
        if (t1 < T){
            #pragma unroll
            for(int s = 0; s < PREATT_S; s++){
                int c = tile_start + s;
                float q_val = (c < HS) ? q_head[t1 * HS +c] : 0.0f;
                #pragma unroll
                for(int u = 0; u < PREATT_U; u++){
                    accumulation[u] += q_val * k_tile[s][u];
                }
            }
        }
        __syncthreads();
    }

    if(t1 < T){
        #pragma unroll
        for(int u = 0; u < PREATT_U; u ++){
            int t2_u = t2 + u;
            if(t2_u < T){
                preatt_head[t1 * T + t2_u] = accumulation[u];
            }
        }
    }
}

#define ATTV_T 128
#define ATTV_U 8
#define ATTV_S 16

__global__ void attention_v_kernel(float* vaccum, float* att, float* v, int B, int T, int NH, int HS){
    int head = blockIdx.z;
    int b = head / NH;
    int nh = head % NH;

    int t = blockIdx.x * ATTV_T + threadIdx.x; // row of output query positin
    int hs_base = blockIdx.y * ATTV_U; //start column key position

    const float* att_head = att + (b * NH + nh) * T * T;
    const float* v_head = v + (b * NH + nh) * T * HS;
    float* out_head = vaccum + (b * NH + nh) * T * HS;

    float accumulation[ATTV_U] = {0.0f};

    __shared__ float v_tile[ATTV_S][ATTV_U];

    for(int tile_start = 0; tile_start < T; tile_start += ATTV_S){
        int s_l = threadIdx.x / ATTV_U;
        int u_l = threadIdx.x % ATTV_U;

        int t2_idx = tile_start + s_l;
        int hs_idx = hs_base + u_l;

        if (t2_idx < T && hs_idx < HS) {
            v_tile[s_l][u_l] = v_head[t2_idx * HS + hs_idx];
        } else {
            v_tile[s_l][u_l] = 0.0f;
        }
        __syncthreads();

        // Each thread loads att values from registers and accumulates
        if (t < T) {
            #pragma unroll
            for (int s = 0; s < ATTV_S; s++) {
                int t2 = tile_start + s;
                float att_val = (t2 < T) ? att_head[t * T + t2] : 0.0f;
                #pragma unroll
                for (int u = 0; u < ATTV_U; u++) {
                    accumulation[u] += att_val * v_tile[s][u];
                }
            }
        }
        __syncthreads();
    }

    if (t < T) {
        #pragma unroll
        for (int u = 0; u < ATTV_U; u++) {
            int hs = hs_base + u;
            if (hs < HS) {
                out_head[t * HS + hs] = accumulation[u];
            }
        }
    }
}

// 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 block_preatt(PREATT_T);
    dim3 grid_preatt((T + PREATT_T - 1) / PREATT_T,(T+PREATT_U-1)/PREATT_U,B*NH);
    preattention_scores_kernel<<<grid_preatt, block_preatt>>>(q, k, B, T, NH, HS, preatt);

    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;
    dim3 block_attv(PREATT_T);
    dim3 grid_attv((T + ATTV_T - 1) / ATTV_T,(HS+ATTV_U-1)/ATTV_U,B*NH);
    attention_v_kernel<<<grid_attv,block_attv>>>(vaccum, att, v, B, T, NH, HS);

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

#endif // __ATTENTION_CUH__