#version 440 core

in vec3 pos;
in vec2 uv;
in vec3 normal;

out vec4 color;

uniform sampler2D textureSampler;
uniform usampler3D octreeSampler;

layout(location = $width$) uniform int width;
layout(location = $height$) uniform int height;
layout(location = $angle$) uniform float angle;
layout(location = $aspect$) uniform float aspect;

layout(location = $lightDirection$) uniform vec3 lightDirection;
layout(location = $cameraPosition$) uniform vec3 cameraPosition;

ivec3 texSize;

//************************************
// Intersection test between ray and cube, also gives intersection points as p0 + tmin * ray and p0 + tmax * ray
//************************************
bool rayCube(vec3 p0, vec3 ray, vec3 center, float extent, out float tmin, out float tmax, out uint minFace, out uint maxFace) {

	// Translate ray origin based on cube center
	p0 -= center;

	// Get t from ray and cube's plane equations and use it to get the intersection coordinates
	float t1 = -(dot(p0, vec3(1., 0., 0.)) - extent) / dot(ray, vec3(1., 0., 0.)); vec3 test1 = p0 + t1*ray;
	float t2 = -(dot(p0, vec3(-1., 0., 0.)) - extent) / dot(ray, vec3(-1., 0., 0.)); vec3 test2 = p0 + t2*ray;
	float t3 = -(dot(p0, vec3(0., 1., 0.)) - extent) / dot(ray, vec3(0., 1., 0.)); vec3 test3 = p0 + t3*ray;
	float t4 = -(dot(p0, vec3(0., -1., 0.)) - extent) / dot(ray, vec3(0., -1., 0.)); vec3 test4 = p0 + t4*ray;
	float t5 = -(dot(p0, vec3(0., 0., 1.)) - extent) / dot(ray, vec3(0., 0., 1.)); vec3 test5 = p0 + t5*ray;
	float t6 = -(dot(p0, vec3(0., 0., -1.)) - extent) / dot(ray, vec3(0., 0., -1.)); vec3 test6 = p0 + t6*ray;

	// Check if t was not negative and that the ray-plane intersection falls within the cube face
	if (t1 < 0. || any(greaterThan(test1.yz, vec2(extent))) || any(lessThan(test1.yz, vec2(-extent))))
		t1 = 0.;
	if (t2 < 0. || any(greaterThan(test2.yz, vec2(extent))) || any(lessThan(test2.yz, vec2(-extent))))
		t2 = 0.;
	if (t3 < 0. || any(greaterThan(test3.xz, vec2(extent))) || any(lessThan(test3.xz, vec2(-extent))))
		t3 = 0.;
	if (t4 < 0. || any(greaterThan(test4.xz, vec2(extent))) || any(lessThan(test4.xz, vec2(-extent))))
		t4 = 0.;
	if (t5 < 0. || any(greaterThan(test5.xy, vec2(extent))) || any(lessThan(test5.xy, vec2(-extent))))
		t5 = 0.;
	if (t6 < 0. || any(greaterThan(test6.xy, vec2(extent))) || any(lessThan(test6.xy, vec2(-extent))))
		t6 = 0.;

	// Initialize tmin and tmax values that define the two intersection points
	tmin = 9999999999.;
	tmax = 0.;

	// Use the lowest value of t that is not 0 for tmin
	if (t1 > 0. && t1 < tmin) { tmin = t1; minFace = 1; }
	if (t2 > 0. && t2 < tmin) { tmin = t2; minFace = 2; }
	if (t3 > 0. && t3 < tmin) { tmin = t3; minFace = 3; }
	if (t4 > 0. && t4 < tmin) { tmin = t4; minFace = 4; }
	if (t5 > 0. && t5 < tmin) { tmin = t5; minFace = 5; }
	if (t6 > 0. && t6 < tmin) { tmin = t6; minFace = 6; }

	// Use the highest value of t that is not 0 for tmax
	if (t1 > 0. && t1 > tmax) { tmax = t1; maxFace = 1; }
	if (t2 > 0. && t2 > tmax) { tmax = t2; maxFace = 2; }
	if (t3 > 0. && t3 > tmax) { tmax = t3; maxFace = 3; }
	if (t4 > 0. && t4 > tmax) { tmax = t4; maxFace = 4; }
	if (t5 > 0. && t5 > tmax) { tmax = t5; maxFace = 5; }
	if (t6 > 0. && t6 > tmax) { tmax = t6; maxFace = 6; }

	// If tmin = tmax, the ray origin is within the cube, so set tmin to 0
	if (tmin == tmax)
		tmin = 0.;

	// If tmax is not 0, an intersection was found
	return tmax > 0.;
}

uvec3 getUloc(vec3 p0, vec3 ray, float t, uint hitFace)
{
	vec3 samplePoint = p0 + t * ray;
	vec3 raySign = clamp(sign(ray), -1, 0);
	// Calculate the location of the first leaf node that the ray enters
	uvec3 uLoc = uvec3(floor(samplePoint));
	// If the ray is moving in a negative direction, then the uLoc will be one position above the desired position.
	switch(hitFace)
	{
		case 1: case 2: uLoc.x = uint(int(round(samplePoint.x + raySign.x))); break;
		case 3: case 4: uLoc.y = uint(int(round(samplePoint.y + raySign.y))); break;
		case 5: case 6: uLoc.z = uint(int(round(samplePoint.z + raySign.z))); break;
	}
	return uLoc;
}

ivec3 wrapNodePointer(in ivec3 samplePos)
{
	// Make sure that the indices are wrapped if the pointer doesn't fit in the current texture
	if (samplePos.x >= texSize.x)
	{
		samplePos.x = samplePos.x % texSize.x;
		samplePos.y++;
		if (samplePos.y >= texSize.y)
		{
			samplePos.y = samplePos.y % texSize.y;
			samplePos.z++;
		}
		//return vec4(255, 255, 0, float($max_level$ + 1));
	}
	return samplePos;
}

//************************************
// Given a pointer to a node, and it's childmask, and some child index
// returns the nodePointer to the child at the given index.
// Note that if the child doesn't exist, it still returns some pointer as if it would.
//************************************
uvec3 getNodePointer(uvec3 curNodePointer, uint childMask, uint childIndex)
{
	// Based on the childIndex and the childmask, find out what the index of the childpointer is
	int childPointerIndex = bitCount(childMask << (31 - childIndex));

	// Find sample position in 3D texture corresponding to current node and locat
	ivec3 samplePos = ivec3(curNodePointer);
	samplePos.x += childPointerIndex;

	samplePos = wrapNodePointer(samplePos);

	// Fetch the pointer from the texture
	return texelFetch(octreeSampler, samplePos, 0).rgb;
}

// Fetches the childmask of the node at curNodePointer (note that this is not a mask for leaf nodes!)
uint getChildMask(in uvec3 nodePointer)
{
	return texelFetch(octreeSampler, ivec3(nodePointer), 0).b;
}

bool hasChild(in uint childMask, in uint childIndex)
{
	uint mask = 1 << childIndex;
	return (childMask & mask) == mask;
}

uint getChildIndex(uvec3 uLoc, uint level)
{
	uint range = 1 << ($max_level$ - level);
	return ((uLoc.x & range) == range ? 1 : 0)
			+ ((uLoc.y & range) == range ? 2 : 0)
			+ ((uLoc.z & range) == range ? 4 : 0);
}

uvec4 sampleOctree(in uvec3 uLoc)
{
	uvec3 result = uvec3(0);
	uint childMask; uint childIndex;
	// Traverse octree
	uint i = 0;
	for (; i <= $max_level$; ++i) {
		// Read the childmask (stored in "blue"):
		childMask = getChildMask(result);

		// Find the index of the child in which the sampleposition is located
		childIndex = getChildIndex(uLoc, i);

		if (!hasChild(childMask, childIndex))
			break;

		result = getNodePointer(result, childMask, childIndex);
	}
	// Return level reached
	return uvec4(result, i);
}

uvec3 getColor(uvec3 nodePointer)
{
	return texelFetch(octreeSampler, ivec3(nodePointer), 0).rgb;
}

void main() {
	uint maxLoop = 300;
	texSize = textureSize(octreeSampler, 0);

	// Set color using textures
	color = texture(textureSampler, uv);
	
	// Scale the grid so that all leaf cells have a width of exactly 1:
	float extent = $extent$;
	uint range = 1 << ($max_level$ + 1);
	float rangeF = float(range);
	float scale = rangeF / (extent * 2.);
	float offset = extent + 0.5 / scale;

	// Get ray and sample point
	vec3 ray = normalize(pos - cameraPosition);
	uvec3 raySignMask = uvec3(clamp(sign(-ray), 0, 1));
	vec3 p0 = (cameraPosition + offset) * scale;

	// Calculate the new position of the root
	vec3 rootCenter = vec3(range) * 0.5;
	float rootExtent = float(range) * 0.5;

	// Initialize some variables
	float tmin, tmax, tminRoot, tmaxRoot;
	uint minFace, maxFace;

	float leafExtent = 0.5;
	
	// Start the ray at the edge of the root cube:
	bool collision = rayCube(p0, ray, rootCenter, rootExtent, tminRoot, tmaxRoot, minFace, maxFace);
	if (!collision) // Ray doesn't intersect with root.
	{
		//color.rgb = vec3(0,0,1);
		return;
	}

	// Find cell through which the ray enters
	uvec3 uLoc = getUloc(p0, ray, tminRoot, minFace);

	// Keep testing ray box intersection, but prevent endless looping
	int loop = 0;
	while (!any(greaterThanEqual(uLoc, uvec3(range))) && !any(lessThan(uLoc, uvec3(0))) && loop < maxLoop) {
		// Check if the current voxel contains geometry, and if so draw it
		uvec4 result = sampleOctree(uLoc);
		if (result.w >= $max_level$ + 1)
		{
			//color.rgb = vec3(result) / 255.;
			color.rgb = vec3(float(loop) / 20.);
			break;
		}

		//uint reachedLevelRange = 1 << ($max_level$ - result.w + 1);
		//float reachedLevelExtent = float(reachedLevelRange) * 0.5;
		//reachedLevelExtent = 0.5;

		// Otherwise, advance to the next voxel
		vec3 curCubeCenter = uLoc + leafExtent;
		if (!rayCube(p0, ray, curCubeCenter, leafExtent, tmin, tmax, minFace, maxFace)) {
			//color.rgb = vec3(1, float(maxFace) / 6., 0);
			color.rgb = uLoc / range;
			break;
		}

		uLoc = getUloc(p0, ray, tmax, maxFace);

		++loop;
	}
	if (loop == maxLoop)
		color.rgb = vec3(1, 0, 0);

	// Discard pixels that are transparent
	if (color.a < 0.5)
		discard;
}