#ifndef __SOFTMAX_KERNEL_CUH__
#define __SOFTMAX_KERNEL_CUH__

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

__global__ void local_softmax_forward_kernel(float* out, float inv_temperature, const float* inp, int N, int T) {
    // Implement this
    int x = blockDim.x * blockIdx.x + threadIdx.x;

    if(x < N*T){
        //row offset based on WIndow size
        int rowoff = x * WINDOW;

        float maxval = -FLT_MAX;
        for (int i = 0; i < WINDOW; i++) {
            maxval = fmaxf(maxval, inp[i + rowoff]);
        }

        float sumval = 0.0f;
        //Traverse Window Length Again
        for (int i = 0; i < WINDOW; i++) {
            float ev = expf(inv_temperature * (inp[i + rowoff] - maxval));
            sumval += ev;
            out[i + rowoff] = ev;
        }

        float norm = 1.0f / sumval;
        for (int i = 0; i < WINDOW; i++) {
            out[i + rowoff] *= norm;
        }
    }
}

#endif // __SOFTMAX_KERNEL_CUH__