#pragma OPENCL EXTENSION cl_khr_local_int32_base_atomics : enable
#pragma OPENCL EXTENSION cl_khr_global_int32_base_atomics : enable
#pragma OPENCL EXTENSION cl_khr_local_int32_extended_atomics : enable
#pragma OPENCL EXTENSION cl_khr_global_int32_extended_atomics : enable

#define PI 3.1415926535f
#define PI_2 (3.1415926535f * 2.0)
#define EPSILON 1e-5f
#define SAMPLES_NUMS 16

struct ScatteringParams
{
    float sunRadius;
	float sunRadiance;

	float mieG;
	float mieHeight;

	float rayleighHeight;

	float3 waveLambdaMie;
	float3 waveLambdaOzone;
	float3 waveLambdaRayleigh;

	float earthRadius;
	float earthAtmTopRadius;
	float3 earthCenter;
};

float3 ComputeSphereNormal(float2 coord, float phiStart, float phiLength, float thetaStart, float thetaLength)
{
	float3 normal;
	normal.x = -sin(thetaStart + coord.y * thetaLength) * sin(phiStart + coord.x * phiLength);
	normal.y = -cos(thetaStart + coord.y * thetaLength);
	normal.z = -sin(thetaStart + coord.y * thetaLength) * cos(phiStart + coord.x * phiLength);
	return normalize(normal);
}

float2 ComputeRaySphereIntersection(float3 position, float3 dir, float3 center, float radius)
{
	float3 origin = position - center;
	float B = dot(origin, dir);
	float C = dot(origin, origin) - radius * radius;
	float D = B * B - C;

	float2 minimaxIntersections;
	if (D < 0.0f)
	{
		minimaxIntersections = (float2)(-1.0f, -1.0f);
	}
	else
	{
		D = sqrt(D);
		minimaxIntersections = (float2)(-B - D, -B + D);
	}

	return minimaxIntersections;
}

float3 ComputeWaveLambdaRayleigh(float3 lambda)
{
	const float n = 1.0003f;
	const float N = 2.545E25f;
	const float pn = 0.035f;
	const float n2 = n * n;
	const float pi3 = 3.1415f * 3.1415f * 3.1415f;
	const float rayleighConst = (8.0f * pi3 * pow(n2 - 1.0f, 2.0f)) / (3.0f * N) * ((6.0f + 3.0f * pn) / (6.0f - 7.0f * pn));
	return rayleighConst / (lambda * lambda * lambda * lambda);
}

float ComputePhaseMie(float theta, float g)
{
	float g2 = g * g;
	return (1.0f - g2) / pow(1.0f + g2 - 2.0f * g * clamp(theta, 0.0f, 1.0f), 1.5f) / (4.0f * 3.1415f);
}

float ComputePhaseRayleigh(float theta)
{
	float theta2 = theta * theta;
	return (theta2 * 0.75f + 0.75f) / (4.0f * 3.1415f);
}

float ChapmanApproximation(float X, float h, float cosZenith)
{
	float c = sqrt(X + h);
	float c_exp_h = c * exp(-h);

	if (cosZenith >= 0.0f)
	{
		return c_exp_h / (c * cosZenith + 1.0f);
	}
	else
	{
		float x0 = sqrt(1.0f - cosZenith * cosZenith) * (X + h);
		float c0 = sqrt(x0);

		return 2.0f * c0 * exp(X - x0) - c_exp_h / (1.0f - c * cosZenith);
	}
}

float GetOpticalDepthSchueler(float h, float H, float earthRadius, float cosZenith)
{
	return H * ChapmanApproximation(earthRadius / H, h / H, cosZenith);
}

float3 GetTransmittance(struct ScatteringParams setting, float3 L, float3 V)
{
	float ch = GetOpticalDepthSchueler(L.y, setting.rayleighHeight, setting.earthRadius, V.y);
	return exp(-(setting.waveLambdaMie + setting.waveLambdaRayleigh) * ch);
}

float2 ComputeOpticalDepth(struct ScatteringParams setting, float3 samplePoint, float3 V, float3 L, float neg)
{
	float rl = length(samplePoint);
	float h = rl - setting.earthRadius;
	float3 r = samplePoint / rl;

	float cos_chi_sun = dot(r, L);
	float cos_chi_ray = dot(r, V * neg);

	float opticalDepthSun = GetOpticalDepthSchueler(h, setting.rayleighHeight, setting.earthRadius, cos_chi_sun);
	float opticalDepthCamera = GetOpticalDepthSchueler(h, setting.rayleighHeight, setting.earthRadius, cos_chi_ray) * neg;

	return (float2)(opticalDepthSun, opticalDepthCamera);
}

void AerialPerspective(struct ScatteringParams setting, float3 start, float3 end, float3 V, float3 L, bool infinite, float3* transmittance, float3* insctrMie, float3* insctrRayleigh)
{
	float inf_neg = infinite ? 1.0f : -1.0f;

	float3 sampleStep = (end - start) / (float)SAMPLES_NUMS;
	float3 samplePoint = end - sampleStep;
	float3 sampleLambda = setting.waveLambdaMie + setting.waveLambdaRayleigh + setting.waveLambdaOzone;

	float sampleLength = length(sampleStep);

	float3 scattering = (float3)(0.0f, 0.0f, 0.0f);
	float2 lastOpticalDepth = ComputeOpticalDepth(setting, end, V, L, inf_neg);

	for (int i = 1; i < SAMPLES_NUMS; i++, samplePoint -= sampleStep)
	{
		float2 opticalDepth = ComputeOpticalDepth(setting, samplePoint, V, L, inf_neg);

		float3 segment_s = exp(-sampleLambda * (opticalDepth.x + lastOpticalDepth.x));
		float3 segment_t = exp(-sampleLambda * (opticalDepth.y - lastOpticalDepth.y));
		
		*transmittance *= segment_t;
		
		scattering = scattering * segment_t;
		scattering += exp(-(length(samplePoint) - setting.earthRadius) / setting.rayleighHeight) * segment_s;

		lastOpticalDepth = opticalDepth;
	}

	*insctrMie = scattering * setting.waveLambdaMie * sampleLength;
	*insctrRayleigh = scattering * setting.waveLambdaRayleigh * sampleLength;
}

float ComputeSkyboxChapman(struct ScatteringParams setting, float3 eye, float3 V, float3 L, float3* transmittance, float3* insctrMie, float3* insctrRayleigh)
{
	bool neg = true;

	float2 outerIntersections = ComputeRaySphereIntersection(eye, V, setting.earthCenter, setting.earthAtmTopRadius);
	if (outerIntersections.y < 0.0f) return 0.0f;

	float2 innerIntersections = ComputeRaySphereIntersection(eye, V, setting.earthCenter, setting.earthRadius);
	if (innerIntersections.x > 0.0f)
	{
		neg = false;
		outerIntersections.y = innerIntersections.x;
	}

	eye -= setting.earthCenter;

	float3 start = eye + V * max(0.0f, outerIntersections.x);
	float3 end = eye + V * outerIntersections.y;

	AerialPerspective(setting, start, end, V, L, neg, transmittance, insctrMie, insctrRayleigh);

	bool intersectionTest = innerIntersections.x < 0.0f && innerIntersections.y < 0.0f;
	return intersectionTest ? 1.0f : 0.0f;
}

float4 ComputeSkyInscattering(struct ScatteringParams setting, float3 eye, float3 V, float3 L)
{
	float3 insctrMie = (float3)(0.0f, 0.0f, 0.0f);
	float3 insctrRayleigh = (float3)(0.0f, 0.0f, 0.0f);
	float3 insctrOpticalLength = (float3)(1.0f, 1.0f, 1.0f);
	float intersectionTest = ComputeSkyboxChapman(setting, eye, V, L, &insctrOpticalLength, &insctrMie, &insctrRayleigh);

	float phaseTheta = dot(V, L);
	float phaseMie = ComputePhaseMie(phaseTheta, setting.mieG);
	float phaseRayleigh = ComputePhaseRayleigh(phaseTheta);
	float phaseNight = 1.0f - clamp(insctrOpticalLength.x * EPSILON, 0.0f, 1.0f);

	float3 insctrTotalMie = insctrMie * phaseMie;
	float3 insctrTotalRayleigh = insctrRayleigh * phaseRayleigh;

	float3 sky = (insctrTotalMie + insctrTotalRayleigh) * setting.sunRadiance;

	float angle = clamp((1.0f - phaseTheta) * setting.sunRadius, 0.0f, 1.0f);
	float cosAngle = cos(angle * PI * 0.5f);
	float edge = ((angle >= 0.9f) ? smoothstep(0.9f, 1.0f, angle) : 0.0f);
                         
	float3 limbDarkening = GetTransmittance(setting, -L, V);
	limbDarkening *= pow((float3)(cosAngle, cosAngle, cosAngle), (float3)(0.420f, 0.503f, 0.652f)) * mix((float3)(1.0f, 1.0f, 1.0f), (float3)(1.2f, 0.9f, 0.5f), edge) * intersectionTest;

	sky += limbDarkening;

	return (float4)(sky.x, sky.y, sky.z, phaseNight * intersectionTest);
}

float3 TonemapACES(float3 x)
{
	const float A = 2.51f;
	const float B = 0.03f;
	const float C = 2.43f;
	const float D = 0.59f;
	const float E = 0.14f;
	return (x * (A * x + B)) / (x * (C * x + D) + E);
}

float noise(float2 uv)
{
	float x = dot(sin(uv.xyx * uv.xyy * 1024.0f), (float3)(341896.483f, 891618.637f, 602649.7031f));
	return (x - floor(x));
}

float4 AtmosphericScattering(float4 V, float4 L)
{
	struct ScatteringParams setting;
	setting.sunRadius = 500.0f;
	setting.sunRadiance = 20.0f;
	setting.mieG = 0.76f;
	setting.mieHeight = 1200.0f;
	setting.rayleighHeight = 8000.0f;
	setting.earthRadius = 6360000.0f;
	setting.earthAtmTopRadius = 6420000.0f;
	setting.earthCenter = (float3)(0.0f, -setting.earthRadius, 0.0f);
	setting.waveLambdaMie = (float3)(2e-7f, 2e-7f, 2e-7f);
    setting.waveLambdaRayleigh = ComputeWaveLambdaRayleigh((float3)(680e-9f, 550e-9f, 450e-9f));
	setting.waveLambdaOzone = (float3)(1.36820899679147f, 3.31405330400124f, 0.13601728252538f) * 0.6e-6f * 2.504f;
	
    float3 eye = (float3)(0.0f, 1000.0f, 0.0f);
   	float4 sky = ComputeSkyInscattering(setting, eye, V.xyz, L.xyz);
    sky.xyz = TonemapACES(sky.xyz * 2.0f);
    //sky.xyz = pow(sky.xyz, float3(1.0f / 2.2f));

	return (float4)(sky.x, sky.y, sky.z, 1.0f);
}

__kernel void TraceNaive(__global float4* triangles,
	__global float4* rays,
	__global float4* results,
	int trianglesCount,
	int raysCount)
{
	int i = get_global_id(0);
	if (i >= raysCount)
	{
		return;
	}

	float4 o = rays[i * 4 + 0];
	float4 d = rays[i * 4 + 1];

	bool hit = true;
	int id = -1;
	float bu = 0.0f, bv = 0.0f;
	float dist = d.w;
	
	for (int n = 0; n < trianglesCount; n++)
	{
		float4 v = triangles[n * 3];
		float4 p = triangles[n * 3 + 1];
		float4 q = triangles[n * 3 + 2];

		float o_z = v.w - o.x * v.x - o.y * v.y - o.z * v.z;
		float i_z = 1.0f / (d.x * v.x + d.y * v.y + d.z * v.z);
		float t = o_z * i_z;

		if (t > o.w && t < dist)
		{
			float o_x = p.w + o.x * p.x + o.y * p.y + o.z * p.z;
			float d_x = d.x * p.x + d.y * p.y + d.z * p.z;
			float u = o_x + t * d_x;

			if (u >= 0.0f && u <= 1.0f)
			{
				float o_y = q.w + o.x * q.x + o.y * q.y + o.z * q.z;
				float d_y = d.x * q.x + d.y * q.y + d.z * q.z;
				float v = o_y + t * d_y;

				if (v >= 0.0f && u + v <= 1.0f)
				{
					dist = t;
					bu = u;
					bv = v;
					id = n;
				}
			}
		}
	}

	results[i] = (float4)(bu, bv, dist, as_float(id));
}

// Uncomment to render statistics
//#define RENDER_STATISTICS
//#define MEMORY_STATISTICS

struct KDNode
{
	union
	{
		float split;
		unsigned int prim_offset;
	};

	union
	{
		unsigned int flags;
		unsigned int prim_count;
		unsigned int above_child;
	};
};

#define SPATIAL_STACK_SIZE 32
struct KDStackNode
{
	unsigned int node;
	float near, far;
	float pad;
};

__kernel void TraceSpatial(__global float4* triangles,
	__global float4* rays,
	__global float4* results,
	__global struct KDNode* nodes,
	__global unsigned int* indices,
	float4 boundsMin,
	float4 boundsMax,
	int trianglesCount,
	int raysCount,
	int2 dimensions)
{
	int i = get_global_id(0);
	int j = get_global_id(1);
	if (i >= dimensions.x || j >= dimensions.y)
	{
		return;
	}

	int k = i + j * dimensions.x;

	float4 o = rays[k * 4 + 0];
	float4 d = rays[k * 4 + 1];
	float4 inv = native_recip(d);
	
	struct KDStackNode stack[SPATIAL_STACK_SIZE];
	unsigned int stack_ptr = 0;

	{
		float4 v1 = (boundsMin - o) * inv;
		float4 v2 = (boundsMax - o) * inv;
		float4 near = min(v1, v2);
		float4 far = max(v1, v2);
		float enter = max(near.x, max(near.y, near.z));
		float exit = min(far.x, min(far.y, far.z));

		if (exit > 0.0f && enter < exit)
		{
			stack[stack_ptr].node = 0;
			stack[stack_ptr].near = enter;
			stack[stack_ptr].far = exit;
			stack_ptr++;

		}
	}

	int id = -1;
	float bu = 0.0f, bv = 0.0f;
	float dist = d.w;

#ifdef RENDER_STATISTICS
	int visited = 0;
	int interiors = 0;
	int leaves = 0;
#elif defined MEMORY_STATISTICS
	int globalop = 0;
	int privateop = 0;
#endif

	unsigned int node;

	while (stack_ptr != 0)
	{
		stack_ptr--;
		node = stack[stack_ptr].node;
		float near = stack[stack_ptr].near;

		float far = stack[stack_ptr].far;

		uint axis = (nodes[node].flags & 3);
		uint above_child = (nodes[node].above_child >> 2);
		float split = nodes[node].split;


#ifdef MEMORY_STATISTICS
		globalop += 3;
		privateop += 3;
#endif


		while (axis != 3)
		{
#ifdef RENDER_STATISTICS

			visited++;
			interiors++;
#endif

			float hitpos;

			int below_first;
			unsigned int first, second;

			float orig_tmp;
			float idir_tmp;


			switch (axis)
			{
			case 0:
				orig_tmp = o.x;

				idir_tmp = inv.x;
				break;

			case 1:
				orig_tmp = o.y;

				idir_tmp = inv.y;
				break;

			case 2:


				orig_tmp = o.z;
				idir_tmp = inv.z;
				break;
			}


			hitpos = (split - orig_tmp) * idir_tmp;
			below_first = (orig_tmp < split) || (orig_tmp == split && idir_tmp >= 0.0f);

			if (below_first)
			{

				first = node + 1;
				second = above_child;
			}
			else
			{

				first = above_child;
				second = node + 1;
			}

			if (hitpos > far || hitpos < 0.0f)

			{
				node = first;
			}
			else if (hitpos < near)
			{

				node = second;
			}
			else
			{
				stack[stack_ptr].node = second;

				stack[stack_ptr].near = hitpos;
				stack[stack_ptr].far = far;
				stack_ptr++;
#ifdef MEMORY_STATISTICS
				privateop += 3;
#endif

				node = first;
				far = hitpos;
			}


			axis = (nodes[node].flags & 3);
			above_child = (nodes[node].above_child >> 2);
			split = nodes[node].split;
#ifdef MEMORY_STATISTICS
			globalop += 3;
#endif
		}

		unsigned int prim_offset = nodes[node].prim_offset;

#ifdef MEMORY_STATISTICS
		globalop += 1;
#endif
		unsigned int prims_num = above_child;


#ifdef RENDER_STATISTICS
		visited++;
		leaves += prims_num;
#endif


		if (prims_num > 0)
		{
			__global unsigned int *prims_ids = &indices[prim_offset];

			for (unsigned int n = 0; n < prims_num; n++)
			{
				// Don't trash cache by reading index through it
				int tri_idx = prims_ids[n] * 3;
				float4 r = triangles[tri_idx + 0];
				float4 p = triangles[tri_idx + 1];
				float4 q = triangles[tri_idx + 2];
#ifdef MEMORY_STATISTICS
				globalop += 4;
#endif

				float o_z = r.w - o.x * r.x - o.y * r.y - o.z * r.z;
				float i_z = 1.0f / (d.x * r.x + d.y * r.y + d.z * r.z);
				float t = o_z * i_z;

				if (t > o.w && t < dist)

				{
					float o_x = p.w + o.x * p.x + o.y * p.y + o.z * p.z;
					float d_x = d.x * p.x + d.y * p.y + d.z * p.z;
					float u = o_x + t * d_x;


					if (u >= 0.0f && u <= 1.0f)
					{
						float o_y = q.w + o.x * q.x + o.y * q.y + o.z * q.z;
						float d_y = d.x * q.x + d.y * q.y + d.z * q.z;
						float v = o_y + t * d_y;


						if (v >= 0.0f && u + v <= 1.0f)
						{
							dist = t;
							bu = u;

							bv = v;
							id = tri_idx;
						}
					}
				}

			}

			if (dist < far)
			{

				break;
			}
		}
	}


#ifdef RENDER_STATISTICS
	results[k] = (float4)((float)visited * 0.01f, (float)interiors * 0.01f, (float)leaves * 0.01f, 1.0f);
#elif defined MEMORY_STATISTICS
	results[k] = (float4)((float)globalop * 0.004f, (float)privateop * 0.004f, 0.0f, 1.0f);
#else
	results[k] = (float4)(bu, bv, dist, as_float(id));
#endif
}

struct BVHStackNode
{
	float4 n0xy;
	float4 n1xy;
	float4 nz;
	
	union
	{
		unsigned int prim_offset;
		unsigned int second_child;
	};

	int prim_count;
	unsigned int axis;
	unsigned int pad;
};

#define BVH_STACK_SIZE 64

__kernel void TraceBVH(__global float4* triangles,
	__global float4* rays,
	__global float4* results,
	__global struct BVHStackNode* nodes,
	__global unsigned int* indices,
	int trianglesCount,
	int raysCount,
	int2 dimensions)
{
    int i = get_global_id(0);
	int j = get_global_id(1);
	if (i >= dimensions.x || j >= dimensions.y)
	{
		return;
	}

	int k = i + j * dimensions.x;

	float4 o = rays[k * 4 + 0];
	float4 d = rays[k * 4 + 1];
	float4 inv = native_recip(d);
	float4 oinv = o * inv;
	
	int stack[BVH_STACK_SIZE];
	int stack_ptr = 0;
	stack[stack_ptr] = 0xffffffff;
	int node_id = 0;

	float tmin = 0.0f;
	float tmax = 10000.0f;
	float bU = 0.0f;
	float bV = 0.0f;
	float dist = tmax;

	float leaf = 0.0f;
	int id = -1;

	do
	{
		__global struct BVHStackNode* n = &(nodes[node_id]);
		int prim_count = n->prim_count;

		// Interior node
		if (prim_count == 0)
		{
			// Fetch children bounding boxes
			float4 n0xy = n->n0xy;
			float4 n1xy = n->n1xy;
			float4 nz = n->nz;

			// Test against child AABBs
			float c0lox = n0xy.x * inv.x - oinv.x;
			float c0hix = n0xy.y * inv.x - oinv.x;
			float c0loy = n0xy.z * inv.y - oinv.y;
			float c0hiy = n0xy.w * inv.y - oinv.y;
			float c0loz = nz.x   * inv.z - oinv.z;
			float c0hiz = nz.y   * inv.z - oinv.z;
			float c1loz = nz.z   * inv.z - oinv.z;
			float c1hiz = nz.w   * inv.z - oinv.z;
			float c0min = max(max(min(c0lox, c0hix), min(c0loy, c0hiy)), max(min(c0loz, c0hiz), tmin));
			float c0max = min(min(max(c0lox, c0hix), max(c0loy, c0hiy)), min(max(c0loz, c0hiz), tmax));
			float c1lox = n1xy.x * inv.x - oinv.x;
			float c1hix = n1xy.y * inv.x - oinv.x;
			float c1loy = n1xy.z * inv.y - oinv.y;
			float c1hiy = n1xy.w * inv.y - oinv.y;
			float c1min = max(max(min(c1lox, c1hix), min(c1loy, c1hiy)), max(min(c1loz, c1hiz), tmin));
			float c1max = min(min(max(c1lox, c1hix), max(c1loy, c1hiy)), min(max(c1loz, c1hiz), tmax));

			bool traverseChild0 = (c0max >= c0min);
			bool traverseChild1 = (c1max >= c1min);
			
			// If no children was hit, get node from stack
			if (!traverseChild0 && !traverseChild1)
			{
				node_id = stack[stack_ptr];
				stack_ptr--;
			}
			else
			{
				unsigned int first_child = node_id + 1;
				unsigned int second_child = n->second_child;
				node_id = (traverseChild0) ? first_child : second_child;

                if (traverseChild0 && traverseChild1)
                {
                    if (c1min < c0min)
					{ 
						node_id = second_child;
						stack_ptr++;
						stack[stack_ptr] = first_child; 
					}
					else
					{
						stack_ptr++;
						stack[stack_ptr] = second_child; 
					}
                }
			}
		}
		// Leaf node
		else
		{
			unsigned int prim_offset = n->prim_offset;
			__global unsigned int *prims_ids = &indices[prim_offset];
			
			for (int n = 0; n < prim_count; n++)
			{
				// Don't trash cache by reading index through it
				int tri_idx = prims_ids[n] * 3;
				float4 r = triangles[tri_idx + 0];
				float4 p = triangles[tri_idx + 1];
				float4 q = triangles[tri_idx + 2];

				float o_z = r.w - o.x * r.x - o.y * r.y - o.z * r.z;
				float i_z = 1.0f / (d.x * r.x + d.y * r.y + d.z * r.z);
				float t = o_z * i_z;

				if (t > tmin && t < tmax)
				{
					float o_x = p.w + o.x * p.x + o.y * p.y + o.z * p.z;
					float d_x = d.x * p.x + d.y * p.y + d.z * p.z;
					float u = o_x + t * d_x;

					if (u >= 0.0f && u <= 1.0f)
					{
						float o_y = q.w + o.x * q.x + o.y * q.y + o.z * q.z;
						float d_y = d.x * q.x + d.y * q.y + d.z * q.z;
						float v = o_y + t * d_y;

						if (v >= 0.0f && u + v <= 1.0f)
						{
							tmax = t;
							bU = u;
							bV = v;
							id = tri_idx;
						}
					}
				}
			}

			node_id = stack[stack_ptr];
			stack_ptr--;
		}
	}
	while (node_id != 0xffffffff);

	results[k] = (float4)(bU, bV, tmax, as_float(id));
}

struct TopLevelNode
{
	unsigned int bvh_index;
	unsigned int pad[3];
	float4 mat0;
	float4 mat1;
	float4 mat2;
	float4 mat3;
};

// Transform vector with a 4x4 matrix.
float4 mul4x1(float4 vec, float4 mat0, float4 mat1, float4 mat2, float4 mat3)
{
    float4 out;

	// Assume vec.w = 1 to save a multiplication
    out.x = mat0.x * vec.x + mat1.x * vec.y + mat2.x * vec.z + mat3.x;
    out.y = mat0.y * vec.x + mat1.y * vec.y + mat2.y * vec.z + mat3.y;
    out.z = mat0.z * vec.x + mat1.z * vec.y + mat2.z * vec.z + mat3.z;
	out.w = 1.0f;

    return out;
}

// Transform vector with a 3x3 rotation matrix. 
// This function ignores the 4 row/col of the matrix.
float4 mul3x1(float4 vec, float4 mat0, float4 mat1, float4 mat2)
{
    float4 out;

    out.x = mat0.x * vec.x + mat1.x * vec.y + mat2.x * vec.z;
    out.y = mat0.y * vec.x + mat1.y * vec.y + mat2.y * vec.z;
    out.z = mat0.z * vec.x + mat1.z * vec.y + mat2.z * vec.z;
	out.w = 0.0f;
	
    return out;
}

// Transform vector with a 3x3 rotation matrix. 
// This function ignores the 4 row/col of the matrix.
float4 mul3x1inv(float4 vec, float4 mat0, float4 mat1, float4 mat2)
{
    float4 out;

    out.x = mat0.x * vec.x + mat0.y * vec.y + mat0.z * vec.z;
    out.y = mat1.x * vec.x + mat1.y * vec.y + mat1.z * vec.z;
    out.z = mat2.x * vec.x + mat2.y * vec.y + mat2.z * vec.z;
	out.w = 0.0f;
	
    return out;
}

float4 TraceRay_MultiLevel(float4 ray_origin,
	float4 ray_direction,
	__global float4* triangles,
	__global struct BVHStackNode* nodes,
	__global unsigned int* indices,
	__global struct TopLevelNode* toplevelData,
	int* matrix_idx)
{
	float4 o = ray_origin;
	float4 d = ray_direction;
	float4 inv = native_recip(d);
	float4 oinv = o * inv;
	
	int stack[BVH_STACK_SIZE];
	int stack_ptr = 0;
	int meshbvh_stack_ptr = -1;
	stack[stack_ptr] = 0xffffffff;
	int node_id = 0;

	float tmin = 0.0f;
	float tmax = 10000.0f;
	float bU = 0.0f;
	float bV = 0.0f;
	float dist = tmax;

	float leaf = 0.0f;
	int id = -1;

	int bvh_offset = 0;

	do
	{
		__global struct BVHStackNode* n = &(nodes[node_id]);
		int prim_count = n->prim_count;

		// Interior node
		if (prim_count == 0)
		{
			// Fetch children bounding boxes
			float4 n0xy = n->n0xy;
			float4 n1xy = n->n1xy;
			float4 nz = n->nz;

			// Test against child AABBs
			float c0lox = n0xy.x * inv.x - oinv.x;
			float c0hix = n0xy.y * inv.x - oinv.x;
			float c0loy = n0xy.z * inv.y - oinv.y;
			float c0hiy = n0xy.w * inv.y - oinv.y;
			float c0loz = nz.x   * inv.z - oinv.z;
			float c0hiz = nz.y   * inv.z - oinv.z;
			float c1loz = nz.z   * inv.z - oinv.z;
			float c1hiz = nz.w   * inv.z - oinv.z;
			float c0min = max(max(min(c0lox, c0hix), min(c0loy, c0hiy)), max(min(c0loz, c0hiz), tmin));
			float c0max = min(min(max(c0lox, c0hix), max(c0loy, c0hiy)), min(max(c0loz, c0hiz), tmax));
			float c1lox = n1xy.x * inv.x - oinv.x;
			float c1hix = n1xy.y * inv.x - oinv.x;
			float c1loy = n1xy.z * inv.y - oinv.y;
			float c1hiy = n1xy.w * inv.y - oinv.y;
			float c1min = max(max(min(c1lox, c1hix), min(c1loy, c1hiy)), max(min(c1loz, c1hiz), tmin));
			float c1max = min(min(max(c1lox, c1hix), max(c1loy, c1hiy)), min(max(c1loz, c1hiz), tmax));

			bool traverseChild0 = (c0max >= c0min);
			bool traverseChild1 = (c1max >= c1min);
			
			// If no children was hit, get node from stack
			if (!traverseChild0 && !traverseChild1)
			{
				if (stack_ptr == meshbvh_stack_ptr)
				{
					meshbvh_stack_ptr = -1;
					o = ray_origin;
					d = ray_direction;
					inv = native_recip(d);
					oinv = o * inv;
				}

				node_id = stack[stack_ptr];
				stack_ptr--;
			}
			else if (traverseChild0 || traverseChild1)
			{
				unsigned int first_child = node_id + 1;
				unsigned int second_child = n->second_child;
				node_id = (traverseChild0) ? first_child : second_child;

                if (traverseChild0 && traverseChild1)
                {
                    if (c1min < c0min)
					{ 
						node_id = second_child;
						stack_ptr++;
						stack[stack_ptr] = first_child; 
					}
					else
					{
						stack_ptr++;
						stack[stack_ptr] = second_child; 
					}
                }
			}
		}
		else if (prim_count < 0)
		{
			meshbvh_stack_ptr = stack_ptr;

			unsigned int bottom_level_offset = n->prim_offset;
			__global unsigned int *bottom_level_ids = &indices[bottom_level_offset];
			
			int bottom_levels_count = -prim_count;
			
			bvh_offset = bottom_level_ids[0];
			int bvh_index = toplevelData[bvh_offset].bvh_index;

			node_id = bvh_index;

			o = mul4x1(ray_origin, toplevelData[bvh_offset].mat0, toplevelData[bvh_offset].mat1, toplevelData[bvh_offset].mat2, toplevelData[bvh_offset].mat3);
			d = mul3x1(ray_direction, toplevelData[bvh_offset].mat0, toplevelData[bvh_offset].mat1, toplevelData[bvh_offset].mat2);
			inv = native_recip(d);
			oinv = o * inv;
		}
		// Leaf node
		else
		{
			unsigned int prim_offset = n->prim_offset;
			__global unsigned int *prims_ids = &indices[prim_offset];
			
			for (int n = 0; n < prim_count; n++)
			{
				// Don't trash cache by reading index through it
				int tri_idx = prims_ids[n] * 3;
				float4 r = triangles[tri_idx + 0];
				float4 p = triangles[tri_idx + 1];
				float4 q = triangles[tri_idx + 2];

				float o_z = r.w - o.x * r.x - o.y * r.y - o.z * r.z;
				float i_z = 1.0f / (d.x * r.x + d.y * r.y + d.z * r.z);
				float t = o_z * i_z;

				if (t > tmin && t < tmax)
				{
					float o_x = p.w + o.x * p.x + o.y * p.y + o.z * p.z;
					float d_x = d.x * p.x + d.y * p.y + d.z * p.z;
					float u = o_x + t * d_x;

					if (u >= 0.0f && u <= 1.0f)
					{
						float o_y = q.w + o.x * q.x + o.y * q.y + o.z * q.z;
						float d_y = d.x * q.x + d.y * q.y + d.z * q.z;
						float v = o_y + t * d_y;

						if (v >= 0.0f && u + v <= 1.0f)
						{
							tmax = t;
							bU = u;
							bV = v;
							id = prims_ids[n];

							*matrix_idx = bvh_offset;
						}
					}
				}
			}
			
			if (stack_ptr == meshbvh_stack_ptr)
			{
				meshbvh_stack_ptr = -1;
				o = ray_origin;
				d = ray_direction;
				inv = native_recip(d);
				oinv = o * inv;
			}

			node_id = stack[stack_ptr];
			stack_ptr--;
		}
	}
	while (node_id != 0xffffffff);

	return (float4)(bU, bV, tmax, as_float(id));
}

struct Triangle
{
	float4 a, b, c;
	float4 na, nb, nc;
	float4 ta, tb, tc;
	float4 ba, bb, bc;
	float4 txa, txb, txc;
};

__constant sampler_t sampler = CLK_NORMALIZED_COORDS_TRUE | CLK_ADDRESS_REPEAT | CLK_FILTER_LINEAR;

float4 shadePixel(float4 ray_origin,
	float4 ray_direction,
	__global float4* triangles,
	__global struct BVHStackNode* nodes,
	__global unsigned int* indices,
	float4 traceResult,
	__global struct Triangle* triangleData,
	__read_only image2d_t atlasImage,
	__global float4* atlasCoords,
	__global struct TopLevelNode* toplevelData,
	int matrix_idx,
	__global float* random,
	int* random_idx,
	float4* out_hitpoint,
	float4* out_normal,
	float4* brdf)
{

	float3 barycentric = (float3)(traceResult.x, traceResult.y, 1.0f - traceResult.x - traceResult.y);
	float4 normal = barycentric.x * triangleData[as_uint(traceResult.w)].na + barycentric.y * triangleData[as_uint(traceResult.w)].nb + barycentric.z * triangleData[as_uint(traceResult.w)].nc;
	float2 uv = barycentric.x * triangleData[as_uint(traceResult.w)].txa.xy + barycentric.y * triangleData[as_uint(traceResult.w)].txb.xy + barycentric.z * triangleData[as_uint(traceResult.w)].txc.xy;
	
	float4 tmp = mul3x1inv(normal, toplevelData[matrix_idx].mat0, toplevelData[matrix_idx].mat1, toplevelData[matrix_idx].mat2);
	normal = normalize(tmp);
	*out_normal = normal;

	float2 atlasOffset = atlasCoords[as_uint(triangleData[as_uint(traceResult.w)].txa.z)].xy;
	float2 atlasScale = atlasCoords[as_uint(triangleData[as_uint(traceResult.w)].txa.z)].zw;
	
	uv.x = uv.x - floor(uv.x);
	uv.y = uv.y - floor(uv.y);

	uv *= atlasScale;
	uv += atlasOffset;

	float4 texture = read_imagef(atlasImage, sampler, uv);
	texture.x = pow(texture.x, 2.2f);
	texture.y = pow(texture.y, 2.2f);
	texture.z = pow(texture.z, 2.2f);

	float4 result = (float4)(1.0f, 1.0f, 1.0f, 1.0f) * dot(normal, normalize((float4)(0.15f, 0.37f, 0.9f, 0.0f)));

	// Shadow ray
	float4 origin = ray_origin + ray_direction * traceResult.z + normal * 0.05f;
	*out_hitpoint = origin;
	float4 direction = (float4)(0.15f, 0.37f, 0.9f, 0.0f);
	direction.x += (random[*random_idx] * 2.0f - 1.0f) * 0.01f; *random_idx = (*random_idx + 1) % 4096;
	direction.y += (random[*random_idx] * 2.0f - 1.0f) * 0.01f; *random_idx = (*random_idx + 1) % 4096;
	direction.z += (random[*random_idx] * 2.0f - 1.0f) * 0.01f; *random_idx = (*random_idx + 1) % 4096;
	direction = normalize(direction);
	int shadow_matrix_idx;
	float4 shadowResult = TraceRay_MultiLevel(origin, direction, triangles, nodes, indices, toplevelData, &shadow_matrix_idx);
	if (as_uint(shadowResult.w) != -1)
	{
		result.x *= 0.0f;
		result.y *= 0.0f;
		result.z *= 0.0f;
	}

	*brdf = result;

	return (float4)(texture.x, texture.y, texture.z, texture.w);
}

float4 sample_cosine_hemisphere(float u1, float u2)
{
    const float r = sqrt(u1);
    const float theta = 2.0f * 3.1415f * u2;
 
    const float x = r * cos(theta);
    const float y = r * sin(theta);
 
    return (float4)(x, y, sqrt(max(0.0f, 1.0f - u1)), 0.0f);
}

__kernel void TraceMultiLevel(__global float4* triangles,
	__global float4* rays,
	__global float4* results,
	__global struct BVHStackNode* nodes,
	__global unsigned int* indices,
	__global struct TopLevelNode* toplevelData,
	__global float* random,
	__global struct Triangle* triangleData,
	int trianglesCount,
	int raysCount,
	int2 dimensions,
	__read_only image2d_t atlasImage,
	__global float4* atlasCoords,
	int offset)
{
    int i = get_global_id(0);
	int j = get_global_id(1);
	if (i >= dimensions.x || j >= dimensions.y)
	{
		return;
	}

	int random_idx = (offset + i * (offset + j)) % 4096;

	int k = i + j * dimensions.x;
	
	float4 radiance = (float4)(0.0f, 0.0f, 0.0f, 1.0f);
	
	//for (int samples = 0; samples < 16; samples++)
	{
		float4 origin = rays[k * 4 + 0];
		float4 direction = rays[k * 4 + 1];
		direction.x += (random[random_idx] * 2.0f - 1.0f) * 0.001f; random_idx = (random_idx + 17336 * j + i * 364) % 4096;
		direction.y += (random[random_idx] * 2.0f - 1.0f) * 0.001f; random_idx = (random_idx + 1) % 4096;
		direction.z += (random[random_idx] * 2.0f - 1.0f) * 0.001f; random_idx = (random_idx + 1) % 4096;
		direction.w = 0.0f;
		direction = normalize(direction);

		float4 color = (float4)(0.0f, 0.0f, 0.0f, 1.0f);
		float4 factor = (float4)(1.0f, 1.0f, 1.0f, 1.0f);

		int iters = 0;
		int path_length = 0;
	
		while (true)
		{
			int matrix_idx;
			float4 traceResult = TraceRay_MultiLevel(origin, direction, triangles, nodes, indices, toplevelData, &matrix_idx);

			if (as_uint(traceResult.w) == -1)
			{
				radiance += factor * AtmosphericScattering(direction, normalize((float4)(0.15f, 0.37f, 0.9f, 0.0f)));
				break;
			}

			float4 t;
			float4 w;
			float4 brdf;
			color = shadePixel(origin, direction, triangles, nodes, indices, traceResult, triangleData, atlasImage, atlasCoords, toplevelData, matrix_idx, random, &random_idx, &t, &w, &brdf);
			radiance += factor * pow(color * brdf, 1.0f / 2.2f);
			factor *= pow(color, 1.0f / 2.2f);

			if (color.w > 0.5f)
			{
				float4 u;
				if (w.x > 0.1f) u = (float4)(0.0f, 1.0f, 0.0f, 0.0f);
				else u = (float4)(1.0f, 0.0f, 0.0f, 0.0f);
				u = normalize(cross(u, w));

				float4 v = cross(w, u);
			
				float r1 = random[random_idx] * 2.0f - 1.0f; random_idx = (random_idx + 1) % 4096;
				float r2 = random[random_idx] * 2.0f - 1.0f; random_idx = (random_idx + 1) % 4096;
				float r3 = random[random_idx] * 2.0f - 1.0f; random_idx = (random_idx + 1) % 4096;
				/*float4 d = sample_cosine_hemisphere(r1, r2);

				origin = t;
				direction = normalize(u * d.x + v * d.y + w * d.z);
				direction = direction - 2.0f * dot(-w, direction) * -w;*/
				
				float4 incident = direction;

				origin = origin + direction * traceResult.z + w * 1.0f;
				//direction = normalize(u * d.x + v * d.y + w * d.z);
				direction = (float4)(r1, r2, r3, 0.0f);
				if (dot(direction, w) < 0.0f)
				{
					direction = -direction;
				}
				direction += w;
				direction = normalize(direction);
				
				iters++;
				path_length++;

				if (path_length > 1)
				{
					break;
				}

				continue;

				/*float4 n = normalize(traceResult.x * triangleData[as_uint(traceResult.w)].na + traceResult.y * triangleData[as_uint(traceResult.w)].nb + (1.0f - traceResult.x - traceResult.y) * triangleData[as_uint(traceResult.w)].nc);
			
				origin = origin + direction * traceResult.z + n * 1.0f;
				direction = direction - 2.0f * dot(n, direction) * n;
			
				iters++;
				if (iters > 1)
				{
					break;
				}
				continue;*/
				//break;
			}

			origin = origin + direction * (traceResult.z + 0.01f);

			iters++;
			if (iters > 10)
			{
				break;
			}
		}
	}

	//results[k] = radiance * 0.03125f;
	results[k] = radiance * 0.5f + results[k] * 0.5f;
	//results[k] = radiance;
}