#include "NBody.hpp"

extern int numBodies;	

/////////////////////////////////// NBody::setupCLKernels Func \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\
//
// This function only sets the Kernel arguments
//
/////////////////////////////////// NBody::setupCLKernels Func \\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\\

int NBody::setupCLKernels()
{
    cl_int   status;

	/* get a kernel object handle for a kernel with the given name */
    kernel = clCreateKernel(
                 program,
                 kernelFunctionName,
                 &status);
	
	if(status != CL_SUCCESS)
    {
        std::cout << "clCreateKernel failed." << std::endl;
        return 1;
    }

	//================================================
	// Arguments specific to all kernels:

    /*** Set appropriate arguments to the kernel ***/

    /* Particle positions */
    status = clSetKernelArg(
                 kernel,
                 0,					// Argument one;
                 sizeof(cl_mem),
                 (void *)&updatedPos);

	if(status != CL_SUCCESS)
    {
        std::cout << "clSetKernelArg failed. (updatedPos)" << std::endl;
        return 1;
    }

    /* Particle velocity */
    status = clSetKernelArg(
                 kernel,
                 1,					// Argument two;
                 sizeof(cl_mem),
                 (void *)&updatedVel);
	
	if(status != CL_SUCCESS)
    {
        std::cout << "clSetKernelArg failed. (updatedVel)" << std::endl;
        return 1;
    }
 
    /* numBodies */
    status = clSetKernelArg(
                 kernel,
                 2,					// Argument three;
                 sizeof(cl_int),
                 (void *)&numBodies);

	if(status != CL_SUCCESS)
    {
        std::cout << "clSetKernelArg failed. (numBodies)" << std::endl;
        return 1;
    }

    /* time step */
    status = clSetKernelArg(
                 kernel,
                 3,					// Argument four;
                 sizeof(cl_float),
                 (void *)&delT);
	if(status != CL_SUCCESS)
    {
        std::cout << "clSetKernelArg failed. (delT)" << std::endl;
        return 1;
    }

    /* upward Pseudoprobability */
    status = clSetKernelArg(
                 kernel,
                 4,					// Argument five;
                 sizeof(cl_float),
                 (void *)&espSqr);

	if(status != CL_SUCCESS)
    {
        std::cout << "clSetKernelArg failed. (espSqr)" << std::endl;
        return 1;
    }

    /* local memory */
    status = clSetKernelArg(
                 kernel,
                 5,					// Argument six;
                 GROUP_SIZE * 4 * sizeof(float),
                 NULL);

	if(status != CL_SUCCESS)
    {
        std::cout << "clSetKernelArg failed. (localPos)" << std::endl;
        return 1;
    }

	//================================================
	// Arguments specific to Leapfrog kernel:

	if (integrator == LEAPFROG){

		/* local memory */
		status = clSetKernelArg(
					 kernel,
					 6,					// Argument seven;
					 GROUP_SIZE * 4 * sizeof(float),
					 NULL);

		if(status != CL_SUCCESS)
		{
			std::cout << "clSetKernelArg failed. (localVel)" << std::endl;
			return 1;
		}

		/* Particle acceleration */
		status = clSetKernelArg(
					 kernel,
					 7,					// Argument eight;
					 sizeof(cl_mem),
					 (void *)&updatedAcc);

		if(status != CL_SUCCESS)
		{
			std::cout << "clSetKernelArg failed. (updatedAcc)" << std::endl;
			return 1;
		}

		/* If it is the first run and we do not have the acc computed yet */
		status = clSetKernelArg(
                 kernel,
                 8,					// Argument nine;
                 sizeof(cl_int),
                 (void *)&initialAccComputation);

		if(status != CL_SUCCESS)
		{
			std::cout << "clSetKernelArg failed. (initialAccComputation)" << std::endl;
			return 1;
		}

		/* If it is the first run and we do not have the acc computed yet */
		status = clSetKernelArg(
                 kernel,
                 9,					// Argument 10;
                 sizeof(cl_mem),
                 (void *)&updatedCollTime);

		if(status != CL_SUCCESS)
		{
			std::cout << "clSetKernelArg failed. (updatedCollTime)" << std::endl;
			return 1;
		}
	}

	//================================================
	// Try to see if there is enough local memory
	// on the device for the above parameter
	// declared to be __local:
    status = clGetKernelWorkGroupInfo(kernel,
                                      devices[0],
                                      CL_KERNEL_LOCAL_MEM_SIZE,
                                      sizeof(cl_ulong),
                                      &usedLocalMemory,
                                      NULL);
	if(status != CL_SUCCESS)
    {
        std::cout << "clGetKernelWorkGroupInfo CL_KERNEL_LOCAL_MEM_SIZE failed." << std::endl;
        return 1;
    }

    if(usedLocalMemory > totalLocalMemory)
    {
        std::cout << "Unsupported: Insufficient local memory on device." << std::endl;
        return 1;
    }

    return 0;
}