#ifndef __ATTENTION_CUH__
#define __ATTENTION_CUH__

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

using namespace nvcuda;

#define WMMA_M 16
#define WMMA_N 16
#define WMMA_K 8

__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_WARPS_M 2
#define PREATT_WARPS_N 4
#define PREATT_BM (PREATT_WARPS_M * WMMA_M)
#define PREATT_BN (PREATT_WARPS_N * WMMA_N)

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

    //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]

    int start_row = blockIdx.x * PREATT_BM;
    int start_col = blockIdx.y * PREATT_BN;

    int warp_id = threadIdx.x / 32;
    int warp_row = warp_id / PREATT_WARPS_N;
    int warp_col = warp_id % PREATT_WARPS_N;

    int warp_out_row = start_row + warp_row * WMMA_M;
    int warp_out_col = start_col + warp_col * WMMA_N;

    // Q * K^T: A=Q (row_major), B=K (col_major gives K^T)
    wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::col_major> b_frag;
    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> c_frag;

    wmma::fill_fragment(c_frag, 0.0f);

    for(int kk = 0; kk < HS; kk += WMMA_K){
        if(warp_out_row < T && kk < HS){
            wmma::load_matrix_sync(a_frag, q_head + warp_out_row * HS + kk, HS);
            #pragma unroll
            for(int i = 0; i < a_frag.num_elements; i++){
                a_frag.x[i] = wmma::__float_to_tf32(a_frag.x[i]);
            }
        }

        if(warp_out_col < T && kk < HS) {
            wmma::load_matrix_sync(b_frag, k_head + warp_out_col * HS + kk, HS);
            #pragma unroll
            for (int i = 0; i < b_frag.num_elements; i++)
                b_frag.x[i] = wmma::__float_to_tf32(b_frag.x[i]);
        }

        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
    }

    //write to global mem
    if(warp_out_row < T && warp_out_col < T){
        wmma::store_matrix_sync(preatt_head + warp_out_row * T + warp_out_col, c_frag, T, wmma::mem_row_major);
    }
}

#define ATTV_WARPS_M 2
#define ATTV_WARPS_N 4
#define ATTV_BM 32
#define ATTV_BN 64

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

    //head pointers 
    const float* att_head = att + (b * NH + nh) * T * T; //[T,HS]
    const float* v_head = v + (b * NH + nh) * T * HS; //[T,HS]
    float* out_head = vaccum + (b * NH +nh) * T * HS; //[T,T]

    int row_start = blockIdx.x * ATTV_BM;
    int col_start = blockIdx.y * ATTV_BN;

    int warpId = threadIdx.x / 32;
    int warp_row = warpId / ATTV_WARPS_N;
    int warp_col = warpId % ATTV_WARPS_N;

    int warp_out_row = row_start + warp_row * WMMA_M;
    int warp_out_col = col_start + warp_col * WMMA_N;

    // att * V: both row-major; standard matmul
    wmma::fragment<wmma::matrix_a, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::row_major> a_frag;
    wmma::fragment<wmma::matrix_b, WMMA_M, WMMA_N, WMMA_K, wmma::precision::tf32, wmma::row_major> b_frag;
    wmma::fragment<wmma::accumulator, WMMA_M, WMMA_N, WMMA_K, float> c_frag;

    wmma::fill_fragment(c_frag, 0.0f);

    for(int kk = 0; kk < T; kk+=WMMA_K){
        if (warp_out_row < T && kk < T) {
            wmma::load_matrix_sync(a_frag, att_head + warp_out_row * T + kk, T);
            #pragma unroll
            for (int i = 0; i < a_frag.num_elements; i++)
                a_frag.x[i] = wmma::__float_to_tf32(a_frag.x[i]);
        }

        if (kk < T && warp_out_col < HS) {
            wmma::load_matrix_sync(b_frag, v_head + kk * HS + warp_out_col, HS);
            #pragma unroll
            for (int i = 0; i < b_frag.num_elements; i++)
                b_frag.x[i] = wmma::__float_to_tf32(b_frag.x[i]);
        }

        wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
    }

    if (warp_out_row < T && warp_out_col < HS) {
        wmma::store_matrix_sync(out_head + warp_out_row * HS + warp_out_col, c_frag, HS, wmma::mem_row_major);
    }
}

// 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((T + PREATT_BM - 1) / PREATT_BM, (T + PREATT_BN - 1) / PREATT_BN, B * NH);
    preattention_scores_kernel<<<grid_preatt, 256>>>(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 grid_attv((T + ATTV_BM - 1) / ATTV_BM, (HS + ATTV_BN - 1) / ATTV_BN, B * NH);
    attention_v_kernel<<<grid_attv,256>>>(vaccum, att, v, B, T, NH, HS);

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

#endif // __ATTENTION_CUH__