#version 330
#extension GL_ARB_shading_language_420pack: enable
layout(binding=0) uniform usamplerBuffer octree;
uniform uint root;
uniform vec3 p;

in vec3 dir;
layout(location=0) out vec4 color;

const uint LEAF_FLAG = 0x80000000u;
const uint LEAF_MASK = LEAF_FLAG - 1u;
const uint TYPE_MASK = 0x00000fffu;

bool isAir(uint leaf) {
	return (leaf & TYPE_MASK) == 0u;
}

void main(void) {
	vec3 d = normalize(dir);

	// Algorithm based on "Efficient Sparse Voxel Octrees - Analysis, Extensions, and Implementations"
	const int s_max = 23; // Maximum scale: number of mantissa bits
	const float epsilon = exp2(float(-s_max));

	uint stack[s_max + 1];

	// Prevent division by zero
	vec3 d2 = vec3(epsilon);
	if (d.x < 0) d2.x *= -1;
	if (d.y < 0) d2.y *= -1;
	if (d.z < 0) d2.z *= -1;
	d = mix(d, d2, lessThan(abs(d), vec3(epsilon)));

	// Precompute coefficients. Octree lies in [1, 2].
	vec3 t_coef = 1 / -abs(d);
	vec3 t_bias = t_coef * p;

	// Mirror the coordinate system so that the ray is negative in each dimension
	// TODO: figure out why the paper had this initialised to 0x7.
	// NOTE: If you uncomment this, remember to flip the comparisions when
	// you reverse the mirroring.
	//int octant_mask = 0x7;
	bvec3 d_gt = greaterThan(d, vec3(0));
	vec3 t_bias_reversed = 3 * t_coef - t_bias;
	//if (d_gt.x) t_bias.x = t_bias_reversed.x;
	//if (d_gt.y) t_bias.y = t_bias_reversed.y;
	//if (d_gt.z) t_bias.z = t_bias_reversed.z;
	// XXX This is ~0.14ms slower.
	t_bias = mix(t_bias, t_bias_reversed, d_gt);
	uint octant_mask = uint(d_gt.x) & 0x1u | uint(d_gt.y) & 0x2u | uint(d_gt.z) & 0x4u;

	// Calculate initial active span
	vec3 t_mins = 2 * t_coef - t_bias;
	float t_min = max(max(t_mins.x, t_mins.y), t_mins.z);
	vec3 t_maxs = t_coef - t_bias;
	float h = min(min(t_maxs.x, t_maxs.y), t_maxs.z);
	t_min = max(t_min, 0);

	// Initialize current voxel
	uint parent = root;
	int scale = s_max - 1;
	float scale_exp2 = 0.5;
	uint child = 0xffffffffu;

	vec3 t_init = 1.5 * t_coef - t_bias;
	bvec3 t_gt = greaterThan(t_init, vec3(t_min));
	uint idx = uint(t_gt.x) & 0x1u | uint(t_gt.y) & 0x2u | uint(t_gt.z) & 0x4u;
	vec3 pos = mix(vec3(1), vec3(1.5), t_gt);

	// Traverse!
	while (scale < s_max) {
		vec3 t_corner = pos * t_coef - t_bias;
		float tc_max = min(min(t_corner.x, t_corner.y), t_corner.z);
		
		// Process the voxel if it is non-empty
		uint child_idx = idx ^ octant_mask;
		//ifdef __ENDIAN_LITTLE__
		child = texelFetch(octree, int(parent*8u + child_idx)).x;
		//else
		//child = as_uint(as_uchar4(parent->children[child_idx]).wzyx);
		//endif
		bool is_leaf = (child & LEAF_FLAG) != 0u;
		if (!is_leaf || !isAir(child)) {
			// Check if we've hit a leaf
			if (is_leaf) {
				break;
			}

			// Intersect
			vec3 t_center = (scale_exp2 * 0.5) * t_coef + t_corner;

			// PUSH
			if (tc_max < h) {
				stack[scale] = parent;
			}
			h = tc_max;
			parent = child;

			// Select first child
			scale--;
			scale_exp2 *= 0.5;

			bvec3 tc_gt = greaterThan(t_center, vec3(t_min));
			idx = uint(tc_gt.x) & 0x1u | uint(tc_gt.y) & 0x2u | uint(tc_gt.z) & 0x4u;
			pos = mix(pos, pos + scale_exp2, tc_gt);

			continue;
		}

		// ADVANCE;
		bvec3 tc_le = lessThanEqual(t_corner, vec3(tc_max));
		uint step_mask = uint(tc_le.x) & 0x1u | uint(tc_le.y) & 0x2u | uint(tc_le.z) & 0x4u;
		pos = mix(pos, pos - scale_exp2, tc_le);

		// Update t-span and child index
		t_min = tc_max;
		idx ^= step_mask;

		// Pop if the bit flips disagree with the ray direction
		if ((idx & step_mask) != 0u) {
			// POP
			// Find the highest differing bits
			uint differing_bits = 0u;
			uvec3 diff_bits = floatBitsToUint(pos) ^ floatBitsToUint(pos + scale_exp2);
			if ((step_mask & 1u) != 0u) differing_bits |= diff_bits.x;
			if ((step_mask & 2u) != 0u) differing_bits |= diff_bits.y;
			if ((step_mask & 4u) != 0u) differing_bits |= diff_bits.z;
			scale = (floatBitsToInt(float(differing_bits)) >> 23) - 127; // position of the highest bit
			scale_exp2 = intBitsToFloat((scale - s_max + 127) << 23); // exp2(scale - s_max)

			// Restore parent voxel
			parent = stack[scale];

			// Round cube positions and extract child index
			uvec3 sh = floatBitsToUint(pos) >> scale;
			pos = uintBitsToFloat(sh << scale);
			idx = (sh.x & 1u) | ((sh.y & 1u) << 1) | ((sh.z & 1u) << 2);

			// Prevent same parent from being stored again
			h = 0;
		}
	}

	// Indicate miss if we are outside the octree
	if (scale >= s_max) t_min = 2;

	// Undo mirroring
	vec3 pos_reversed = 3 - scale_exp2 - pos;
	pos = mix(pos, pos_reversed, notEqual(octant_mask & uvec3(0x1, 0x2, 0x4), uvec3(0)));

	vec3 hit_pos = min(max(p + t_min * d, pos + epsilon), pos + scale_exp2 - epsilon);
	vec3 cube_pos = mod(hit_pos, scale_exp2) / scale_exp2;

	uint type = child & TYPE_MASK;
	if (t_min < 2) {
		switch (type) {
		case 1u: // Stone
			color = vec4(0.5, 0.5, 0.5, 1);
			break;
		case 2u: // Grass
			color = vec4(0, 0.25, 0, 1);
			break;
		case 3u: // Dirt
			color = vec4(0.5, 0.25, 0.125, 1);
			break;
		case 4u: // Cobble
			color = vec4(0.25, 0.25, 0.25, 1);
			break;
		case 5u: // Wood planks
			color = vec4(0.5, 0.5, 0.125, 1);
			break;
		case 8u: // Water
			color = vec4(0, 0, 1, 1);
			break;
		case 9u: // Stationary water
			color = vec4(0, 0, 0.5, 1);
			break;
		case 12u: // Sand
			color = vec4(0.85, 0.825, 0.625, 1);
			break;
		case 13u: // Gravel
			color = vec4(0.5, 0.485, 0.48, 1);
			break;
		case 17u: // Wood
			color = vec4(0.25, 0.25, 0.125, 1);
			break;
		case 18u: // Leaves
			color = vec4(0, 0.5, 0, 1);
			break;
		case 31u: // Tall grass
			color = vec4(0.1, 0.35, 0.1, 1);
			break;
		case 32u: // Dead bush
			color = vec4(0.5, 0.3, 0.1, 1);
			break;
		case 78u: // Snow
			color = vec4(1, 1, 1, 1);
			break;
		case 81u: // Cactus
			color = vec4(0.125, 0.4, 0.1, 1);
			break;
		case 106u: // Vines
			color = vec4(0, 0.35, 0, 1);
			break;
		case 127u: // Cocoa
			color = vec4(0.6, 0.3, 0.125, 1);
			break;
		default:
			color = vec4((child >> 16) & 0xffu, (child >> 8) & 0xffu, child & 0xffu, 255u) / 255;
		}
		vec3 cpos = abs(cube_pos - 0.5);
		if (cpos.y < cpos.x || cpos.y < cpos.z) {
			if (cpos.x > cpos.z) {
				color *= 0.9;
			} else {
				color *= 0.8;
			}
		}
	} else {
		color = vec4(d, 1);
	}
}