#define DATA_SIZE (64 * 1024 * 1024)

#define MAX_BATCH_SIZE (16 * 1024)
#define GROUPSIZE 128
#define GPUMEM_GB 4.0


#include <cmath>
#include <cstring>
#include <cstdlib>
#include <vector>
#include <iostream>
#include <exception>
#include <fstream>
#include <chrono>

#include <CL/cl.h>

cl_device_id get_device()
{
	int errcode;

	cl_platform_id platform;
	errcode = clGetPlatformIDs(1, &platform, NULL);
	if (errcode < 0)
	{
		std::cerr << "Couldn't determine platform" << std::endl;
		exit(1);
	}

	cl_device_id device;
	errcode = clGetDeviceIDs(platform, CL_DEVICE_TYPE_GPU, 1, &device, NULL);
	if (errcode == CL_DEVICE_NOT_FOUND)
	{
		std::cerr << "No GPU devices found" << std::endl;
		exit(1);
	}

	return device;
}

cl_program build_program(cl_context ctx, cl_device_id dev)
{
	int errcode;

	const char* source = "                                          \
		__kernel void elementwise_mult(__global float* data,       \
					__global float* data2, __global float* result,            \
						int arrsize) {                                  \
																\
					uint addr = get_global_id(0);                       \
					if (addr < arrsize)                                 \
				result[addr] = data[addr] * data2[addr];        \
		}";

	std::cout << "Creating program..." << std::endl;

        const char** sources = &source;
        size_t srcsize = strlen(source);
	cl_program program = clCreateProgramWithSource(ctx, 1, sources, (const size_t*)&srcsize, &errcode);
	if (errcode < 0)
	{
		std::cerr << "Program creation failed" << std::endl;
		exit(1);
	}

	std::cout << "Building program..." << std::endl;
	errcode = clBuildProgram(program, 0, NULL, NULL, NULL, NULL);
	if (errcode < 0)
	{
		size_t logsize;

		std::cerr << "Program build error. Getting log..." << std::endl;

		clGetProgramBuildInfo(program, dev, CL_PROGRAM_BUILD_LOG, 0, NULL, &logsize);

		char* logcontent;
		logcontent = (char*)malloc(logsize + 1);
		clGetProgramBuildInfo(program, dev, CL_PROGRAM_BUILD_LOG, logsize + 1, logcontent, NULL);

		logcontent[logsize] = '\0';

		std::cerr << logcontent << std::endl;

		free(logcontent);
		exit(1);
	}
	std::cout << "Program successfully built" << std::endl;

	return program;
}

int main(int argc, char* argv[])
{
	int errcode;
	cl_device_id device = get_device();
	cl_context context = clCreateContext(NULL, 1, &device, NULL, NULL, &errcode);
	if (errcode < 0)
	{
		std::cerr << "OpenCL context creation failed" << std::endl;
		return 1;
	}

	cl_program program = build_program(context, device);

	std::vector<cl_queue_properties> commprops;
	commprops.push_back(0);		// the list must be terminated with zero

	cl_command_queue queue = clCreateCommandQueueWithProperties(context, device, commprops.data(), &errcode);
	if (errcode < 0)
	{
		std::cerr << "Command queue creation failed" << std::endl;
		return 1;
	}

	cl_kernel kernel = clCreateKernel(program, "elementwise_mult", &errcode);
	if (errcode < 0)
	{
		std::cerr << "Kernel creation failed" << std::endl;
		return 1;
	}


	std::vector<float> data1(DATA_SIZE);
	std::vector<float> data2(DATA_SIZE);
	std::vector<float> result(DATA_SIZE);

	std::cout << "Populating data..." << std::endl;

	for (uint64_t i = 0; i < data1.size(); i++)
	{
		data1[i] = i + 1;
		data2[i] = i + 2;
	}

	std::cout << "Calculating..." << std::endl;

	int64_t batch_size = MAX_BATCH_SIZE;
	if (data1.size() < MAX_BATCH_SIZE)
		batch_size = data1.size();


	cl_mem input_buffer, input_buffer2, mult_buffer;
	input_buffer = clCreateBuffer(context, CL_MEM_READ_ONLY, batch_size * sizeof(float), NULL, &errcode);
	input_buffer2 = clCreateBuffer(context, CL_MEM_READ_ONLY, batch_size * sizeof(float), NULL, &errcode);
	mult_buffer = clCreateBuffer(context, CL_MEM_WRITE_ONLY, batch_size * sizeof(float), NULL, &errcode);
	if (errcode < 0)
	{
		std::cerr << "Buffer creation failed" << std::endl;
		return 1;
	}


	long long maxmem_elements = (uint64_t)GPUMEM_GB * 1024 * 1024 * 1024 / sizeof(float);
	long long maxmem_elements_per_buffer = maxmem_elements / 3;  // divide by 3 because we have 3 buffers involved in the operation

	if (batch_size > maxmem_elements_per_buffer)
	{
		batch_size = maxmem_elements_per_buffer;
	}


	size_t group_work_items = GROUPSIZE;
	size_t total_work_items = batch_size + group_work_items - (batch_size % group_work_items);

	auto total_start = std::chrono::high_resolution_clock::now();
	int64_t passes = 0;

	for (int64_t pos = 0; pos < data1.size(); pos += batch_size)
	{
		int64_t cur_batch_size = batch_size < data1.size() - pos ? batch_size : data1.size() - pos;

		cl_event writeevent1, writeevent2;
		clEnqueueWriteBuffer(queue, input_buffer, CL_FALSE, 0, cur_batch_size * sizeof(float), data1.data() + pos, 0, NULL, &writeevent1);
		clEnqueueWriteBuffer(queue, input_buffer2, CL_FALSE, 0, cur_batch_size * sizeof(float), data2.data() + pos, 0, NULL, &writeevent2);

		cl_int arrsize = cur_batch_size;

		errcode = clSetKernelArg(kernel, 0, sizeof(cl_mem), &input_buffer) |
			clSetKernelArg(kernel, 1, sizeof(cl_mem), &input_buffer2) |
			clSetKernelArg(kernel, 2, sizeof(cl_mem), &mult_buffer) |
			clSetKernelArg(kernel, 3, sizeof(cl_int), &arrsize);
		if (errcode < 0)
		{
			std::cerr << "Kernel argument creation failed" << std::endl;
			return 1;
		}

		cl_event kernelevent;
		errcode = clEnqueueNDRangeKernel(queue, kernel, 1, NULL, &total_work_items, &group_work_items, 0, NULL, &kernelevent);
		if (errcode < 0)
		{
			std::cerr << "Enqueue of kernel failed" << std::endl;
			return 1;
		}

		cl_event readevent;
		errcode = clEnqueueReadBuffer(queue, mult_buffer, CL_FALSE, 0, sizeof(float) * cur_batch_size, result.data() + pos, 0, NULL, &readevent);
		clWaitForEvents(1, &readevent);

		if (errcode < 0)
		{
			std::cerr << "Buffer read failed" << std::endl;
			return 1;
		}

		passes++;
	}

	//////
	auto total_end = std::chrono::high_resolution_clock::now();
	int64_t total_time = std::chrono::duration_cast<std::chrono::milliseconds>(total_end - total_start).count();
	std::cout << "Calculation: " << total_time << " milliseconds, " << passes << " pass(es), " <<
		(double)data1.size() / 1000000 / total_time << " GFLOPS\n";

	for (uint64_t i = 0; i < data1.size(); i++)
	{
		if (data1[i] * data2[i] != result[i])
		{
			std::cerr << "The result doesn't match at " << i << "; expected " << data1[i] * data2[i] << ", was " << result[i] << "\n";
			break;
		}
	}
	/////

	clReleaseKernel(kernel);
	clReleaseMemObject(mult_buffer);
	clReleaseMemObject(input_buffer);
	clReleaseMemObject(input_buffer2);
	clReleaseCommandQueue(queue);
	clReleaseProgram(program);
	clReleaseContext(context);

	return 0;
}