// MP Scan
// Given a list (lst) of length n
// Output its prefix sum = {lst[0], lst[0] + lst[1], lst[0] + lst[1] + ... + lst[n-1]}

#include    <wb.h>

#define BLOCK_SIZE 512 //@@ You can change this

#define wbCheck(stmt) do {                                                    \
        cudaError_t err = stmt;                                               \
        if (err != cudaSuccess) {                                             \
            wbLog(ERROR, "Failed to run stmt ", #stmt);                       \
            wbLog(ERROR, "Got CUDA error ...  ", cudaGetErrorString(err));    \
            return -1;                                                        \
        }                                                                     \
    } while(0)

__global__ void scan(float * input, float * output, int len) {
    //@@ Modify the body of this function to complete the functionality of
    //@@ the scan on the device
    //@@ You may need multiple kernel calls; write your kernels before this
    //@@ function and call them from here

    __shared__ float sList[2*BLOCK_SIZE];

    int t = threadIdx.x;
    int b = blockIdx.x*(BLOCK_SIZE*2);

    sList[t]            = ( b   + t < len ? input[b  +t] : 0);
    sList[t+BLOCK_SIZE] = ( b+BLOCK_SIZE+t < len ? input[b+BLOCK_SIZE+t] : 0);

    __syncthreads();

    for( int s = 1; s <= BLOCK_SIZE; s *= 2)
    {
        int i = (t+1)*s*2 -1;
        if( i < 2*BLOCK_SIZE)
            sList[i] += sList[i-s];
        __syncthreads();
    }

    for( int s = BLOCK_SIZE/2; s > 0; s /= 2)
    {
        __syncthreads();
        int i = (t+1)*s*2 -1;
        if( i+s < 2*BLOCK_SIZE)
            sList[i+s] += sList[i];
    }

    __syncthreads();

    if( b+t < len) output[b+t] = sList[t];
    if( b+BLOCK_SIZE+t < len) output[b+BLOCK_SIZE+t] = sList[t+BLOCK_SIZE];
}

int main(int argc, char ** argv) {
    wbArg_t args;
    float * hostInput; // The input 1D list
    float * hostOutput; // The output list
    float * deviceInput;
    float * deviceOutput;
    int numElements; // number of elements in the list

    args = wbArg_read(argc, argv);

    wbTime_start(Generic, "Importing data and creating memory on host");
    hostInput = (float *) wbImport(wbArg_getInputFile(args, 0), &numElements);
    hostOutput = (float*) malloc(numElements * sizeof(float));
    wbTime_stop(Generic, "Importing data and creating memory on host");

    wbLog(TRACE, "The number of input elements in the input is ", numElements);

    wbTime_start(GPU, "Allocating GPU memory.");
    wbCheck(cudaMalloc((void**)&deviceInput, numElements*sizeof(float)));
    wbCheck(cudaMalloc((void**)&deviceOutput, numElements*sizeof(float)));
    wbTime_stop(GPU, "Allocating GPU memory.");

    wbTime_start(GPU, "Clearing output memory.");
    wbCheck(cudaMemset(deviceOutput, 0, numElements*sizeof(float)));
    wbTime_stop(GPU, "Clearing output memory.");

    wbTime_start(GPU, "Copying input memory to the GPU.");
    wbCheck(cudaMemcpy(deviceInput, hostInput, numElements*sizeof(float), cudaMemcpyHostToDevice));
    wbTime_stop(GPU, "Copying input memory to the GPU.");

    //@@ Initialize the grid and block dimensions here

    dim3 Block( BLOCK_SIZE, 1, 1);
    dim3 Grid((numElements-1) / (BLOCK_SIZE*2) +1, 1, 1);

    wbTime_start(Compute, "Performing CUDA computation");
    //@@ Modify this to complete the functionality of the scan
    //@@ on the deivce

    scan<<< Grid, Block>>>( deviceInput, deviceOutput, numElements);

    cudaDeviceSynchronize();
    wbTime_stop(Compute, "Performing CUDA computation");

    wbTime_start(Copy, "Copying output memory to the CPU");
    wbCheck(cudaMemcpy(hostOutput, deviceOutput, numElements*sizeof(float), cudaMemcpyDeviceToHost));
    wbTime_stop(Copy, "Copying output memory to the CPU");

    wbTime_start(GPU, "Freeing GPU Memory");
    cudaFree(deviceInput);
    cudaFree(deviceOutput);
    wbTime_stop(GPU, "Freeing GPU Memory");

    for( int i = BLOCK_SIZE*2; i < numElements; i++)
    {
        hostOutput[i]+= hostOutput[(i/(BLOCK_SIZE*2))*BLOCK_SIZE*2 -1];
    }

    wbSolution(args, hostOutput, numElements);

    free(hostInput);
    free(hostOutput);

    return 0;
}