#ifndef __ATTENTION_KV_CACHE_TEMPLATE_CUH__
#define __ATTENTION_KV_CACHE_TEMPLATE_CUH__

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

extern bool g_enable_kv_cache; // Toggles KV-caching logic
extern float* g_kv_cache;      // Pointer to the GPU-allocated KV cache buffer
extern int g_layer_idx;        // Tracks the current transformer layer, useful for knowing when to switch to decode
extern bool g_is_prefill;      // True during the initial prompt, False during token generation
extern int g_current_pos;      // Tracks the absolute sequence position during Decode, helps index into the KV cache
constexpr int MAX_SEQ_LEN = 1024; // gpt 2 default
// --- Suggested Kernel Placeholders (TODO: Add your regular attention kernel here, and any other kernels you may need) ---
//all the tests are running out of mem to store KV so implementing dynamic alloc
static int s_cache_b = 0;
static int s_cache_nh = 0;
static int s_cache_hs = 0;


__global__ void permute_kernel(float* __restrict__ q, float* __restrict__ k, float*  __restrict__ v, const float* __restrict__ 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(const float* __restrict__ inp, float* __restrict__ 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_];
    }

}

// base line kernels

__global__ void preattention_scores_kernel(const float* __restrict__ q, const float* __restrict__ k, int B, int T, int NH, int HS, float* __restrict__ 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;
        const float* __restrict__ q_ptr = q + b * NH * T * HS + nh * T * HS + t1 * HS;
        const float* __restrict__ k_ptr = k + b * NH * T * HS + nh * T * HS + t2 * HS;
        for (int i = 0; i < HS; i++) {
            sum += q_ptr[i] * k_ptr[i];
        }
        preatt[b * NH * T * T + nh * T * T + t1 * T + t2] = sum;
    }   
}

__global__ void attention_v_kernel(float* __restrict__ vaccum, const float* __restrict__ att, const float* __restrict__ 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;
        const float* __restrict__ att_ptr = att + b * NH * T * T + nh * T * T + t * T;
        const float* __restrict__ v_ptr   = v + b * NH * T * HS + nh * T * HS;
        for (int i = 0; i < T; i++) {
            sum += att_ptr[i] * v_ptr[i * HS + hs];
        }
        vaccum[b * NH * T * HS + nh * T * HS + t * HS + hs] = sum;
    }
}



// helper kernel to calculate index into 1D cache buffer
__device__ inline int get_cache_index(int layer, int is_v, int b, int nh, int t, int B, int NH, int HS) {
    // our layout layer -> k/v -> batch -> head -> time -> head_size
    return ((((layer * 2 + is_v) * B + b) * NH + nh) * MAX_SEQ_LEN + t) * HS;
}


/**
 * Kernel to store the newly computed Key and Value into the global KV cache. 
 * CAN BE FUSED WITH QKV MATMUL
 * In Prefill: Stores T tokens for the current layer.
 * In Decode: Stores 1 token at g_current_pos for the current layer.
 */
__global__ void store_kv_kernel(float* cache, const float* k, const float* v,int layer, int B, int NH, int HS, int T, int current_pos, bool is_prefill) {
    // Implement logic to copy K and V from the input QKV buffer into g_kv_cache
    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;

        //check where to write in the cache
        int cache_idx;
        if(is_prefill){
            cache_idx = t;
        } else {
            cache_idx = current_pos;
        }

        // Write K and V to the cache
        int cache_source_idx = b * NH * T * HS + nh * T * HS + t * HS + hs;
        cache[get_cache_index(layer, 0, b, nh, cache_idx, B, NH, HS) + hs] = k[cache_source_idx]; // is_v = 0 for K
        cache[get_cache_index(layer, 1, b, nh, cache_idx, B, NH, HS) + hs] = v[cache_source_idx]; // is_v = 1 for V
    }
}

/**
 * Kernel to compute the attention for the Decode phase (T=1).
 * This involves:
 * 1. Extracting the current Query (Q)
 * 2. Gathering all previous Keys (K) and Values (V) from the cache
 * 3. Computing attention using the new Q and cached KV values.
 */
__global__ void decode_attention_kernel(float* vaccum, const float* q, const float* cache, int layer, int B, int NH, int HS, int current_pos) {
    // Implement T=1 attention logic using the KV cache
    int x = blockDim.x * blockIdx.x + threadIdx.x;

    if(x >= B*NH*HS){
        return;
    }

    int b = (x / HS / NH);
    int nh = (x / HS) % NH;
    int hs = x % HS;

    float max_val = -FLT_MAX;
    float scale = rsqrtf((float)HS);
    const float* __restrict__ q_ptr = q + b * NH * HS + nh * HS;

    // Q * K^T for all previous tokens in the cache
    for(int t = 0; t <= current_pos; t++){
        const float* __restrict__ k_ptr = cache + get_cache_index(layer, 0, b, nh, t, B, NH, HS); // is_v = 0 for K
        float score = 0.0f;
        for(int i = 0; i < HS; i++){
            score += q_ptr[i] * k_ptr[i];
        }
        score *= scale;
        if(score > max_val){
            max_val = score;
        }
    }

    // Compute softmax
    float sum_exp = 0.0f;
    float output = 0.0f;
    
    for(int t = 0; t <= current_pos; t++){
        const float* __restrict__ k_ptr = cache + get_cache_index(layer, 0, b, nh, t, B, NH, HS); 
        float score = 0.0f;
        for(int i = 0; i < HS; i++){
            score += q_ptr[i] * k_ptr[i];
        }
        score *= scale;
        
        float exp_score = expf(score - max_val);
        sum_exp += exp_score;
        
        const float* __restrict__ v_ptr = cache + get_cache_index(layer, 1, b, nh, t, B, NH, HS); 
        output += exp_score * v_ptr[hs];
    }
    vaccum[b * NH * HS + nh * HS + hs] = output / sum_exp;
}

//alloc and realloc kv cache
static void allocate_kv_cache(int B, int NH, int HS){
    if(g_kv_cache != nullptr && s_cache_b == B && s_cache_nh == NH && s_cache_hs == HS){
        // If the cache exists and the dimensions match, do nothing
        return;
    }

    // If the cache exists but the dimensions don't match, free and reallocate
    if(g_kv_cache != nullptr){
        cudaFree(g_kv_cache);
        g_kv_cache = nullptr;
    }

    size_t total_cache_size = 12 * 2 * B * NH * MAX_SEQ_LEN * HS * sizeof(float); // 12 layers, K and V, max seq len
    cudaMalloc(&g_kv_cache, total_cache_size);
    cudaMemset(g_kv_cache, 0, total_cache_size); // 0 the cache

    s_cache_b = B;
    s_cache_nh = NH;
    s_cache_hs = HS;
}


void attention_forward(float* out, float* qkvr, float* att, float* inp, int B, int T, int C, int NH) {
    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(256, 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);

    // 1. Check if KV-caching is enabled, used by the verification script
    if (!g_enable_kv_cache) {
        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.0f / sqrtf((float)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((B * NH * T * HS + block.x - 1) / block.x);   // <<< ADD
        attention_v_kernel<<<grid_attv, block>>>(vaccum, att, v, B, T, NH, HS);  // <<< grid_attv not grid_preatt
        unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);
        return;
    }

    allocate_kv_cache(B, NH, HS); // Ensure KV cache is allocated with the right dimensions
    // store K and V in the cache
    store_kv_kernel<<<grid, block>>>(g_kv_cache, k, v, g_layer_idx, B, NH, HS, T, g_current_pos, g_is_prefill);


    if (g_is_prefill) {
        // --- PREFILL PHASE ---
        // During prefill, we process the entire input sequence (T > 1).
        
        // A. Store the K and V results from this layer into the cache.
        //    The projections are available in 'inp' (the QKV matmul output).
        
        // B. Run standard baseline attention for the initial prompt tokens.

        // C. Update the global layer index to track progress through the model.

        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.0f / sqrtf((float)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((B * NH * T * HS + block.x - 1) / block.x);
        attention_v_kernel<<<grid_attv, block>>>(vaccum, att, v, B, T, NH, HS);
        unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);

        g_layer_idx++;
        if (g_layer_idx == 12) {
            g_layer_idx = 0;
            g_is_prefill = false; // Transition to Decode phase for the next token
            g_current_pos = T; // Set current position to the end of the prefill sequence
        }
    } else {
        // --- DECODE PHASE ---
        // During decode, we only process the single newest token (T = 1).
        
        // A. Store the single new Key and Value for this step into the cache at g_current_pos.
        
        // B. Run specialized Decode Attention.
        //    This requires reading the history of Keys and Values from the cache
        //    to compute the output for the current token.
        //    TODO: Don't forget to scale your attention scores
        //    TODO: Ensure the output shapes match the next layer
        
        // C. Update the global state.

        float *vaccum = inp;
        dim3 grid_dec((B * NH * HS + block.x - 1) / block.x);
        decode_attention_kernel<<<grid_dec, block>>>(vaccum, q, g_kv_cache, g_layer_idx, B, NH, HS, g_current_pos);
        unpermute_kernel<<<grid, block>>>(vaccum, out, B, T, NH, HS);   

        g_layer_idx++;
        if (g_layer_idx == 12) {
            g_layer_idx = 0;
            g_current_pos++; // Advance the sequence position
        }
    }
}

#endif // __ATTENTION_KV_CACHE_TEMPLATE_CUH__