#include <cub/cub.cuh>
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
using std::cout;
using std::endl;

void dev_cumsum (const float *dev_inData, float *dev_outData, int n) {
    // Set up the working space on the device
    void* dev_temp_storage = NULL;
    size_t temp_storage_bytes = 0;
    cub::DeviceScan::ExclusiveSum(dev_temp_storage, temp_storage_bytes, const_cast<float*>(dev_inData), dev_outData, n);
    cudaMalloc(&dev_temp_storage, temp_storage_bytes);

    // Run the calculation
    cudaError_t error = cub::DeviceScan::ExclusiveSum(dev_temp_storage, temp_storage_bytes, const_cast<float*>(dev_inData), dev_outData, n);
    cudaFree(dev_temp_storage);

    // Abort on error
    if (error != 0) {
        cout << "Aborting with error: " << error << endl;
        exit(-1);
    }
}

int test (float* host, float* device, unsigned int length) {
    for (unsigned int i = 0; i < length; i++) {
        if (host[i] != device[i]) {
            return i;
        }
    }
    return -1;
}

int main(){

    for (unsigned int i = 72; i < 1024; i++) {
        float raw_data[i];
        float simple_calculation[i];
        float device_calculation[i];

        raw_data[0] = rand()%10;
        simple_calculation[0] = 0;
        for (unsigned int j = 1; j < i; j++) {
            raw_data[j] = rand()%10;
            simple_calculation[j] = raw_data[j-1] + simple_calculation[j-1];
        }

        size_t size = i * sizeof(float);

        float* device_input;
        cudaMalloc(&device_input, size);
        cudaMemcpy(device_input, raw_data, size, cudaMemcpyHostToDevice);

        float* device_output;
        cudaMalloc(&device_output, size);
        dev_cumsum(device_input, device_output, i);
        cudaMemcpy(device_calculation, device_output, size, cudaMemcpyDeviceToHost);

        cudaFree(device_input);
        cudaFree(device_output);

        int failIndex = test(simple_calculation, device_calculation, i);

        if (failIndex < 0) {
            cout << i << " succeeded" << endl;
        }
        else {
            cout << i << " failed in position " << failIndex << endl;
        }
    }
}