#include "../shared.h"
#include "randomnumbers.h"
#include "program.h"

#define BALLCOUNT		4096
#define GRIDX			(800 / 10)
#define GRIDY			(480 / 10)
#define BALLSPERCELL	8

__kernel void clear( write_only image2d_t outimg )
{
	int column = get_global_id( 0 );
	int line = get_global_id( 1 );
	if ((column >= 800) || (line >= 480)) return;
	write_imagef( outimg, (int2)(column, line), 0 );
}

__kernel void update( global float* balls )
{
	int idx = get_global_id( 0 );
	balls[idx * 6 + 0] += balls[idx * 6 + 2];
	balls[idx * 6 + 1] += balls[idx * 6 + 3];
}

__kernel void clearGrid( global unsigned int* grid )
{
	unsigned int idx = get_global_id( 0 );
	unsigned int baseIdx = idx * (BALLSPERCELL + 1);
	grid[baseIdx] = 0;
}

__kernel void fillGrid( global float* balls, global unsigned int* grid )
{
	unsigned int idx = get_global_id( 0 );
	unsigned int gx = balls[idx * 6 + 0] * GRIDX;
	unsigned int gy = balls[idx * 6 + 1] * GRIDY;
	if ((gx < 0) || (gy < 0) || (gx >= GRIDX) || (gy >= GRIDY)) return;
	unsigned int baseIdx = (gx + gy * GRIDX) * (BALLSPERCELL + 1);
	if (grid[baseIdx] < BALLSPERCELL)
	{ 
		unsigned int count = atomic_inc( grid + baseIdx );
		if (count < BALLSPERCELL) grid[baseIdx + count + 1] = idx;
	}
	else balls[idx * 6 + 1] = balls[idx * 6 + 5] = 0.1;
}

__kernel void render( write_only image2d_t outimg, global float* balls, global unsigned int* grid )
{
	int column = get_global_id( 0 );
	int line = get_global_id( 1 );
	if ((column >= 800) || (line >= 480)) return;
	float2 uv = { (float)column / 800.0, (float)line / 480.0 };
	// draw balls using grid
	int gx = uv.x * GRIDX;
	int gy = uv.y * GRIDY;
	int gx1 = max( 0, gx - 1 ), gx2 = min( GRIDX - 1, gx + 1 );
	int gy1 = max( 0, gy - 1 ), gy2 = min( GRIDY - 1, gy + 1 );
	for( int y = gy1; y <= gy2; y++ ) for( int x = gx1; x <= gx2; x++ )
	{
		unsigned int baseIdx = (x + y * GRIDX) * (BALLSPERCELL + 1);
		unsigned int count = min( (unsigned int)BALLSPERCELL, grid[baseIdx] );
		for( int i = 0; i < count; i++ )
		{ 
			unsigned int idx = grid[baseIdx + i + 1];
			float2 pos = { balls[idx * 6], balls[idx * 6 + 1] };
		float dist = length( pos - uv );
			if (dist < 0.01f)
		{
			write_imagef( outimg, (int2)(column, 479 - line), (float4)(1,0,0,1) );
			break;
		}
	}
	}
}