struct Geometry {
			glm::vec4 signature;
			glm::vec4 vertex1_or_sphere_center;
			glm::vec4 vertex2_or_sphere_radius;
			glm::vec4 vertex3_or_nothing;
		};

		struct GeometryCollection {
			Geometry geometry[1 << expOfTwoOfMaxGeometryElementsInCollection]; // geometry slots
		};

		struct AABB {
			glm::vec4 position; // this is the point of the aabb with min x, y and z components
			glm::vec4 dimensions; // this is the amount to be added to aabbVertex to obtain the point with max x, y and z in the aabb
		};

		struct alignas(16) NodeData {
			AABB aabb;
			alignas(sizeof(glm::uint32)) glm::uint32 left; // this can also be the content if right is zero
			alignas(sizeof(glm::uint32))glm::uint32 right; // this is zero if the node is a leaf
		};

		struct BLAS {
			glm::mat4 ModelMatrix;

			GeometryCollection collection[1 << expOfTwoOfMaxCollectionElementsInBLAS];

			NodeData tree[((1 << expOfTwoOfMaxCollectionElementsInBLAS) * 2) - 1];
		};

		struct TLAS {
			glm::mat4 ViewMatrix;
			NodeData tree[((1 << expOfTwoOfMaxBLASElementsInTLAS) * 2) - 1];
			BLAS blas[1 << expOfTwoOfMaxBLASElementsInTLAS];
		};


/******************************************************************************************/

#version 450

#define COMP_SIZE 32

#define expOfTwoOfMaxGeometryElementsInCollection 3
#define expOfTwoOfMaxCollectionElementsInBLAS 5
#define expOfTwoOfMaxBLASElementsInTLAS 5

const float PI = 3.14159265359;
const float infinity = 1. / 0.;
const mat4 identityTransform = mat4(1);

// This gets adjusted on each call by the OpenGL core
layout(local_size_x = COMP_SIZE, local_size_y = COMP_SIZE, local_size_z = 1) in;

layout(rgba32f, binding = 0) uniform image2D img_output;

uniform uint width; // Render target surface width
uniform uint height; // Render target surface height

struct Geometry {
	vec4 signature;
	vec4 vertex1_or_sphere_center;
	vec4 vertex2_or_sphere_radius;
	vec4 vertex3_or_nothing;
};

struct GeometryCollection {
	Geometry geometry[1 << expOfTwoOfMaxGeometryElementsInCollection]; // geometry slots
};

struct AABB {
	vec4 position; // this is the point of the aabb with min x, y and z components
	vec4 dimensions; // this is the amount to be added to aabbVertex to obtain the point with max x, y and z in the aabb
};

struct NodeData {
	AABB aabb;
	uint left; // this can also be the content if right is zero
	uint right; // this is zero if the node is a leaf
};

struct BLAS {
	mat4 ModelMatrix;

	GeometryCollection collection[1 << expOfTwoOfMaxCollectionElementsInBLAS];
	
	NodeData tree[((1 << expOfTwoOfMaxCollectionElementsInBLAS) * 2) - 1];
};

struct TLAS {
	mat4 ViewMatrix;
	NodeData tree[((1 << expOfTwoOfMaxBLASElementsInTLAS) * 2) - 1];
	BLAS blas[1 << expOfTwoOfMaxBLASElementsInTLAS];
};

// Output buffer object to be written
layout(std430, binding = 3) buffer scene {
	TLAS tlas[];
};

uniform uint shadingAlgorithm;

uniform vec3 mOrigin;
uniform vec3 mLowerLeftCorner;
uniform vec3 mHorizontal;
uniform vec3 mVertical;

/*************************************************************************************************************************
 *                                                  Algorithms                                                           *
 *************************************************************************************************************************/

struct RayGeometryIntersection {
	float dist;
	vec4 point;
	vec4 normal;
};

struct Ray {
	vec4 origin;
	vec4 direction;
};

// Empty geometry is a sphere with radius equals to 0.0 such geometry CANNOT be intersected
const Geometry emptyGeometry = Geometry(
	vec4(0, 0, 0, 0),
	vec4(0, 0, 0, 0),
	vec4(0, 0, 0, 0),
	vec4(0, 0, 0, 0)
);

const RayGeometryIntersection miss = RayGeometryIntersection(infinity, vec4(0, 0, 0, 1), vec4(0, 0, 0, 0));

bool isLeaf(NodeData node) {
	return node.right == 0;
}

bool hasMissed(RayGeometryIntersection test) {
	return isinf(test.dist) || isnan(test.dist);
}

vec3 rayPointAt(Ray ray, float coeff) {
	return vec3(ray.origin) + coeff * vec3(ray.direction);
}

vec3 applyShading(RayGeometryIntersection isect, uint algorithm) {
	if (algorithm == 0) { // hit or miss
		return hasMissed(isect) ? vec3(0,0,0) : vec3(1,1,1);
	} else if (algorithm == 1) {
		// Here the camera position is vec3(0, 0, 0) because all of this is in camera space
		return vec3(max(0, dot(isect.normal, normalize(/*mPosition*/ -isect.point))));
	}

	return vec3(0.5, 0.5, 0.5); // this is to get attention as this should never verify
}

vec4 getInvDirection(Ray ray) {
	return vec4(1.0/ray.direction.x, 1.0/ray.direction.y, 1.0/ray.direction.z, 0);
}

AABB transformAABB(AABB starting, mat4 transformMatrix) {
	vec4 v[8] = {
		vec4(starting.position.x, starting.position.y, starting.position.z, 1),
		vec4(starting.position.x, starting.position.y, starting.position.z, 1),
		vec4(starting.position.x, starting.position.y, starting.position.z, 1),
		vec4(starting.position.x, starting.position.y, starting.position.z, 1),
		vec4(starting.position.x, starting.position.y, starting.position.z, 1),
		vec4(starting.position.x, starting.position.y, starting.position.z, 1),
		vec4(starting.position.x, starting.position.y, starting.position.z, 1),
		vec4(starting.position.x, starting.position.y, starting.position.z, 1),
	};

	float maxX = 1.0/0.0, maxY = 1.0/0.0, maxZ=1.0/0.0, minX=-1.0/0.0, minY=-1.0/0.0, minZ=-1.0/0.0;

	for (uint i = 0; i < 8; ++i) {
		maxX = (v[i].x > maxX) ? v[i].x : maxX;
		maxY = (v[i].y > maxY) ? v[i].y : maxY;
		maxZ = (v[i].z > maxZ) ? v[i].z : maxZ;
		minX = (v[i].x < minX) ? v[i].x : minX;
		minY = (v[i].y < minY) ? v[i].y : minY;
		minZ = (v[i].z < minZ) ? v[i].z : minZ;
	}

	return AABB(vec4(minX, minY, minZ, 1), vec4(maxX-minX, maxY-minY, maxZ-minZ, 0));
}

struct raySigns {
	bool[3] signs;
};

bool intersectAABB(Ray ray, AABB aabb, mat4 transformMatrix) {
	AABB transformedAABB = /*(transformMatrix != identityTransform) ? */transformAABB(aabb, transformMatrix) /*: *this*/;

	float tmin, tmax, tymin, tymax, tzmin, tzmax;

	vec3 orig = ray.origin.xyz;
	vec3 invdir = getInvDirection(ray).xyz;

	vec3 bounds[2] = {
		transformedAABB.position.xyz,
		(transformedAABB.position.xyz) + (transformedAABB.dimensions.xyz)
	};

	const int[3] raySigns = int[3](
		(invdir.x < 0) ? 1 : 0,
		(invdir.y < 0) ? 1 : 0,
		(invdir.z < 0) ? 1 : 0
	);

	tmin = (bounds[raySigns[0]].x - orig.x) * invdir.x;
	tmax = (bounds[1 - raySigns[0]].x - orig.x) * invdir.x;
	tymin = (bounds[raySigns[1]].y - orig.y) * invdir.y;
	tymax = (bounds[1 - raySigns[1]].y - orig.y) * invdir.y;

	if ((tmin > tymax) || (tymin > tmax))
		return false;
	if (tymin > tmin)
		tmin = tymin;
	if (tymax < tmax)
		tmax = tymax;

	tzmin = (bounds[raySigns[2]].z - orig.z) * invdir.z;
	tzmax = (bounds[1 - raySigns[2]].z - orig.z) * invdir.z;

	if ((tmin > tzmax) || (tzmin > tmax))
		return false;
	if (tzmin > tmin)
		tmin = tzmin;
	if (tzmax < tmax)
		tmax = tzmax;

	return true;
}

RayGeometryIntersection intersect(Ray ray, Geometry geometry, mat4 transformMatrix, float minDistance, float maxDistance) {
	if (geometry.signature.w == 0.0) { // test ray-sphere intersection
		const vec3 center = vec3(transformMatrix * geometry.vertex1_or_sphere_center);
		const float radius = geometry.vertex2_or_sphere_radius.x;

		if (radius == 0) return miss;

		const vec3 origin = vec3(ray.origin);
		const vec3 direction = vec3(ray.direction);

		vec3 oc = origin - center;
	
		const float a = dot(direction, direction);
		const float b = dot(oc, direction);
		const float c = dot(oc, oc) - radius * radius;
		const float discriminant = b * b - a * c;

		if (discriminant > 0.0) { // if delta > 0 then we have two intersections (one for each side of the sphere)
			const float squareRoot = sqrt(discriminant);

			// x0 and x1 are the distances to the origin of the ray
			float x0 = (-b - squareRoot) / a, x1 = (-b + squareRoot) / a;

			// Use x0 and x1 to calculate intersection points
			vec3 point_x0 = rayPointAt(ray, x0), point_x1 = rayPointAt(ray, x1); // so if I use that distance as a coefficient I obtain the intersection point
		
			// Use intersecting points to calculate the surface normal at that point
			vec3 normal_x0 = (point_x0 - center) / radius, normal_x1 = (point_x1 - center) / radius; // and the normal for a point p is (point - center)/radius

			// No valid intersection? return a miss.
			if (((x0 <= minDistance) || (x0 >= maxDistance)) && ((x1 <= minDistance) || (x1 >= maxDistance))) {
				return miss;
			}

			// Choose the closest intersection point
			return (((x0 > minDistance) && (x0 < maxDistance)) && (x0 <= x1)) || ((x1 <= minDistance) || (x1 >= maxDistance)) ?
				RayGeometryIntersection(x0, vec4(point_x0, 1), vec4(normal_x0, 0)) : RayGeometryIntersection(x1, vec4(point_x1, 1), vec4(normal_x1, 0));
		}
	} else if (geometry.signature.w == 1.0) { // test ray-triangle intersection

	}

	return miss;
}

RayGeometryIntersection intersect(Ray ray, GeometryCollection collection, mat4 transformMatrix, float minDistance, float maxDistance) {
	RayGeometryIntersection bestHit = miss;

	for (uint i = 0; i < collection.geometry.length(); ++i) {
		// Execute the ray-geometry intersection algorithm
		const RayGeometryIntersection currentIntersectionInfo = intersect(ray, collection.geometry[i], transformMatrix, minDistance, maxDistance);

		// Check if this is a better hit than the former one
		if (currentIntersectionInfo.dist < bestHit.dist) {
			bestHit = currentIntersectionInfo;
		}
	}

	return bestHit;
}

RayGeometryIntersection intersect(Ray ray, BLAS blas, mat4 transformMatrix, float minDistance, float maxDistance) {
	// Adjust transformation matrix to consider the BLAS model matrix
	transformMatrix = transformMatrix * blas.ModelMatrix;

	RayGeometryIntersection bestHit = miss;

	int currentDepth = 0;
	uint currentPath = 0;

	while ((currentDepth < 32) && (currentDepth >= 0)) {
		uint currentNodeIndex = 0;

		// This for cycles is used to reach the interesting node
		for (uint i = 0; i < currentDepth; ++i) {
			currentNodeIndex = ((currentPath & (1 << currentDepth)) == 0) ? blas.tree[currentNodeIndex].left : blas.tree[currentNodeIndex].left;
		}

		if ((!isLeaf(blas.tree[currentNodeIndex])) && (intersectAABB(ray, blas.tree[currentNodeIndex].aabb, transformMatrix))) {
			// Go one level deeper
			currentDepth += 1;
		} else {
			// This is a leaf that should be tested...
			if ((isLeaf(blas.tree[currentNodeIndex])) && (intersectAABB(ray, blas.tree[currentNodeIndex].aabb, transformMatrix))) {
				RayGeometryIntersection currentHit = intersect(ray, blas.collection[blas.tree[currentNodeIndex].left], transformMatrix, minDistance, maxDistance);

				if (currentHit.dist < bestHit.dist) bestHit = currentHit;
			}

			// in either case: subtree discarded or visited leaf the algorithm should reduce the depth and change the path
			currentDepth -= 1;

			while (((currentPath & (1 << currentDepth)) == 1) && (currentDepth >= 0)) {
				//set the bit at currentDepth position to 0
				currentPath &= ~(1 << currentDepth);

				// and decrease currentDepth
				currentDepth -= 1;
			}

			if (currentDepth >= 0) {
				// set the bit at currentDepth position to 1 (the next visit will choose another path)
				currentPath |= (1 << currentDepth);
			} else {
				// the visited node was the very last one (the right-most) and the iteration should come to an end
				break;
			}
		}
	}

	return bestHit;
}

RayGeometryIntersection intersect(Ray ray, TLAS tlas, mat4 transformMatrix, float minDistance, float maxDistance) {
	// Adjust transformation matrix to consider the BLAS model matrix
	transformMatrix = transformMatrix * tlas.ViewMatrix;

	RayGeometryIntersection bestHit = miss;

	int currentDepth = 0;
	uint currentPath = 0;

	while ((currentDepth < 32) && (currentDepth >= 0)) {
		uint currentNodeIndex = 0;

		// This for cycles is used to reach the interesting node
		for (uint i = 0; i < currentDepth; ++i) {
			currentNodeIndex = ((currentPath & (1 << currentDepth)) == 0) ? tlas.tree[currentNodeIndex].left : tlas.tree[currentNodeIndex].left;
		}

		if ((!isLeaf(tlas.tree[currentNodeIndex])) && (intersectAABB(ray, tlas.tree[currentNodeIndex].aabb, transformMatrix))) {
			// Go one level deeper
			currentDepth += 1;
		} else {
			// This is a leaf that should be tested...
			if ((isLeaf(tlas.tree[currentNodeIndex])) && (intersectAABB(ray, tlas.tree[currentNodeIndex].aabb, transformMatrix))) {
				RayGeometryIntersection currentHit = intersect(ray, tlas.blas[tlas.tree[currentNodeIndex].left], transformMatrix, minDistance, maxDistance);

				if (currentHit.dist < bestHit.dist) bestHit = currentHit;
			}

			// in either case: subtree discarded or visited leaf the algorithm should reduce the depth and change the path
			currentDepth -= 1;

			while (((currentPath & (1 << currentDepth)) == 1) && (currentDepth >= 0)) {
				//set the bit at currentDepth position to 0
				currentPath &= ~(1 << currentDepth);

				// and decrease currentDepth
				currentDepth -= 1;
			}

			if (currentDepth >= 0) {
				// set the bit at currentDepth position to 1 (the next visit will choose another path)
				currentPath |= (1 << currentDepth);
			} else {
				// the visited node was the very last one (the right-most) and the iteration should come to an end
				break;
			}
		}
	}

	return bestHit;
}

Ray generateCameraRay(float s, float t) {
	return Ray(vec4(mOrigin, 1), vec4(mLowerLeftCorner + s * mHorizontal + t * mVertical - mOrigin, 0));
}

void main() {
	// base pixel colour for image
	vec4 pixel = vec4(0, 0, 0, 0);
	// get index in global work group i.e x,y position
	ivec2 pixel_coords = ivec2(gl_GlobalInvocationID.xy);

	// Avoid calculating useless pixels
	if ((gl_GlobalInvocationID.x >= width) || (gl_GlobalInvocationID.y >= height)) {
		return;
	}
	
	const float u = float(gl_GlobalInvocationID.x) / float(width);
	const float v = float(gl_GlobalInvocationID.y) / float(height);

	// Generate camera ray
	Ray cameraRay = generateCameraRay(u, v);

	pixel = vec4(applyShading(intersect(cameraRay, tlas[0], identityTransform, 0.001, 1000.0), shadingAlgorithm), 1.0);
  
	// output to a specific pixel in the image
	imageStore(img_output, pixel_coords, pixel);
}