#version 430 core
#extension GL_ARB_compute_shader : enable
#extension GL_ARB_shader_image_load_store : enable

// output of compute shader is a TEXTURE
layout(rgba32f, binding = 0) uniform image2D outTexture;


// receive as inputs, rays from the centre of the camera to the frustum corners
layout( location = 0 ) uniform vec4 ray00; // top left in world space
layout( location = 1 ) uniform vec4 ray01; // top right in world space
layout( location = 2 ) uniform vec4 ray10; // bottom left in world space
layout( location = 3 ) uniform vec4 ray11; // bottom right in world space

// receive the camera position
layout( location = 4 ) uniform vec3 camPos; // camera position

// receive how many primitives of each type
// geomCount.x = number of spheres
// geomCount.y = number of triangles
// geomCount.z = number of planes
// geomCount.w = number of obbs
layout( location = 5 ) uniform ivec4 geomCount; // sphere, triangle, planes, obbs

// how many lights to process (how many elements in the array "light")
layout( location = 6 ) uniform int lightCount;

// receive an ARRAY of light_type,
struct light_type {
	 vec4 position;
	 vec4 colour;
};
layout(std430, binding = 2) buffer light
{
	light_type lights_data[];
};

struct sphere_type {
	vec4 centre_radius;
	vec4 colour;
};
// 8 floats per sphere: c cy cz rad r g b a
layout(std430, binding = 3) buffer sphere
{
	sphere_type spheres_data[];
};

struct plane_type {
	vec4 normal_d;
	vec4 colour;
};
// 8 floats per plane: nx ny nz d r g b a
layout(std430, binding = 4) buffer plane
{
	plane_type planes_data[];
};

struct triangle_type {
	vec4 vtx0, vtx1, vtx2;
	vec4 colour;
};
layout(std430, binding = 5) buffer triangle
{
	triangle_type triangles_data[];
};

// 20 floats per obb: centre4, u3, hu, v3, hv, w3, hw, rgba4
struct obb_type {
	vec4 centre, u_hu, v_hv, w_hw, colour;
};
layout(std430, binding = 6) buffer obb
{
	obb_type obb_data[];
};

layout (local_size_x = 20, local_size_y = 20) in;

#define SPHERE_TYPE 0
#define LIGHT_TYPE 1
#define PLANE_TYPE 2
#define TRIANGLE_TYPE 3
#define OBB_TYPE 4

// primitives tests forward declarations
float sphereTest(vec3 rayDir, vec3 rayOrigin, vec4 centre_radius);
float planeTest(vec3 rayDir, vec3 rayOrigin, vec4 plane_info);
float triangleTest(vec3 rayDir, vec3 rayOrigin, triangle_type tri);
float obbTest(vec3 rayDir, vec3 rayOrigin, obb_type obb);

// entire scene test forward declaration
void sceneTest(in vec4 ray_dir, inout float lastT, inout int objIndex, inout int objType);

// shade a point of a surface (for any primitive) forward declaration
vec4 shade(in vec3 pointOnSurface, in vec3 normal, in vec3 colour);

vec4 castRay(ivec2 pos)
{
	// normalise values from [ 0,width; 0,height ] to [ 0,1; 0,1]
	vec2 interp = vec2(pos) / imageSize(outTexture);

	// compute ray as interpolation of input rays (corners)
	vec4 ray_dir = normalize(
		mix(
			mix(ray00,ray01,interp.y), // left corners together
			mix(ray10,ray11,interp.y), // right corners together
			interp.x // join new interpolated rays into a final ray
		)
	);

	// lastT is the last intersection found (will keep the closest value)
	float lastT=-1;
	// will remember which object index was hit if any.
	int objIndex = -1;
	// will remember which object type was hit if any.
	int objType = -1;

	// do ray versus scene test
	sceneTest(ray_dir, lastT, objIndex, objType);

	vec3 pointOnSurface = ray_dir.xyz * lastT;
	// set pixel to BACKGROUND colour
	vec4 finalPixelOut = vec4(0.2,0.2,0.2, 1.0);

	if (objIndex >= 0) {

		// did we hit something?

		// IMPLEMENT HERE.
		// FIND POINT IN SPACE WHERE THE INTERSECTION HAPPENED.
		vec3 pointOnSurface;
		if (objType==SPHERE_TYPE)
			{
				// IMPLEMENT HERE
				// COMPUTE FINAL PIXEL COLOUR USING SHADE() FUNCTION
				finalPixelOut = shade(pointOnSurface, spheres_data[objIndex].centre_radius.xyz,finalPixelOut.xyz);
			}
		else if (objType==LIGHT_TYPE)
		{
			// IMPLEMENT HERE.
			// LIGHTS ARE DRAWN AS SPHERES, JUST GIVE A CONSTANT COLOUR TO IT.
			// DO NOT SHADE WITH SHADE() FUNCTION.
			finalPixelOut = vec4(1.0,1.0,1.0,1.0);
		}
		else if (objType == PLANE_TYPE)
		{
			// IMPLEMENT HERE
			// COMPUTE FINAL PIXEL COLOUR USING SHADE() FUNCTION
			finalPixelOut = shade(pointOnSurface, planes_data[objIndex].normal_d.xyz,finalPixelOut.xyz);
		}
		else if (objType == TRIANGLE_TYPE)
		{

			// IMPLEMENT HERE
			// COMPUTE FINAL PIXEL COLOUR USING SHADE() FUNCTION
		vec3 normal = cross(triangles_data[objIndex].vtx1.xyz - triangles_data[objIndex].vtx0.xyz, triangles_data[objIndex].vtx2.xyz - triangles_data[objIndex].vtx0.xyz);
		finalPixelOut = shade(pointOnSurface, normal,finalPixelOut.xyz);
		}
		else if (objType == OBB_TYPE)
        {
            // IMPLEMENT HERE
            // COMPUTE FINAL PIXEL COLOUR USING SHADE() FUNCTION
            float epsilon = 0.0001;
            obb_type obb = obb_data[objIndex];
            vec3 normal;
            bool normalFound = false;

            vec4 vecArr[3];
            vecArr[0] = obb.u_hu;
            vecArr[1] = obb.v_hv;
            vecArr[2] = obb.w_hw;

            for(int i = 0; i < 3 && normalFound == false; i++)
            {
                vec3 n = normalize(vecArr[i].xyz);
                vec3 centreToFace =  n * vecArr[i].w;
                vec3 pointOnFace = obb.centre.xyz + centreToFace;

                if(abs(dot(pointOnSurface - pointOnFace, n)) < epsilon)
                {
                    normalFound = true;
                    normal = n;
                }
            }

            finalPixelOut = shade(pointOnSurface, normal, finalPixelOut.xyz);
        }
    }
    return finalPixelOut;
}

layout (local_size_x = 20, local_size_y = 20) in;
void main()
{
	// pixel coordinate, x in [0,width-1], y in [0,height-1]
	ivec2 pos = ivec2(gl_GlobalInvocationID.xy);
	vec4 finalPixel = castRay(pos);
	imageStore(outTexture, pos, finalPixel);
	return;
};

void sceneTest(in vec4 ray_dir,inout float lastT, inout int objIndex, inout int objType)
{

	// use camPos to know where is the ray origin.
	// example:
	float testT;
	// test all spheres (in geomCount.x)
	for (int i=0;i< geomCount.x;i++)
	{
		sphere_type sphere = spheres_data[i];
		testT = sphereTest(ray_dir.xyz, camPos, sphere.centre_radius);
		if(testT != -1){
			if(testT < lastT || lastT == -1){
				lastT = testT;
				objType = SPHERE_TYPE;
				objIndex = i;
			}
		}
		// IMPLEMENT HERE.
		// TEST SPHERE, UPDATE lastT, objIndex and objType if necessary.
	}
	// lights (we render lights as spheres)
	for (int i=0;i<lightCount;i++)
	{
		light_type light = lights_data[i];

		testT = sphereTest(ray_dir.xyz, camPos, light.position);
		if(testT != -1){
			if(testT < lastT || lastT == -1){
				lastT = testT;
				objType = SPHERE_TYPE;
				objIndex = i;
			}
		}
		// IMPLEMENT HERE.
		// TEST SPHERE (LIGHT), UPDATE lastT, objIndex and objType if necessary.
	}
	// planes
	for (int i = 0; i<geomCount.y; i++)
	{
		plane_type plane = planes_data[i];
		testT = planeTest(ray_dir.xyz, camPos, plane.normal_d);
		if(testT != -1){
			if(testT < lastT || lastT == -1){
				lastT = testT;
				objType = PLANE_TYPE;
				objIndex = i;
				}
		}
		// IMPLEMENT HERE.
		// TEST PLANE, UPDATE lastT, objIndex and objType if necessary.
	}
	// triangles
	for (int i = 0; i<geomCount.z; i++)
	{
		triangle_type triangle = triangles_data[i];
		// IMPLEMENT HERE.
		testT = triangleTest(ray_dir.xyz, camPos, triangle);
		if(testT != -1){
			if(testT < lastT || lastT == -1){
				lastT = testT;
				objType = TRIANGLE_TYPE;
				objIndex = i;
			}
		}

		// TEST TRIANGLE, UPDATE lastT,  objIndex and objType if necessary.
	}
	// OBBs
	for (int i = 0; i < geomCount.w; i++)
	{
		obb_type obb = obb_data[i];

		testT = obbTest(ray_dir.xyz, camPos, obb);
		if(testT != -1){
			if(testT < lastT || lastT == -1){
				lastT = testT;
				objType = OBB_TYPE;
				objIndex = i;
			}
		}
		// IMPLEMENT HERE.
		// TEST OBB, UPDATE lastT, objIndex and objType if necessary.
	}
}

float sphereTest(in vec3 rayDir, in vec3 rayOrigin, in vec4 centre_radius)
{
	vec3 oc = rayOrigin - centre_radius.xyz;
	float p = dot(rayDir, oc);
	float q = dot(oc,oc) - (centre_radius.w * centre_radius.w);
	if(p*p - q < 0){
	return -1.0f;
	}else
		return min(-p+sqrt(p*p - q), -p - sqrt(p*p -q));
}

float planeTest(vec3 rayDir, vec3 rayOrigin, vec4 plane_info)
{
	// IMPLEMENT HERE
	vec3 oc = rayOrigin - plane_info.xyz;
	if(dot(plane_info.xyz,rayDir) == 0)
	{
		return -1;
	}
	else
	{
	float t = (-plane_info.w- dot(plane_info.xyz,rayOrigin)) / dot(plane_info.xyz,rayDir);
	if( t >= 0)
	{
		return t;
	}
	}

		return -1;

}

float triangleTest(vec3 rayDir, vec3 rayOrigin, triangle_type tri)
{
	vec3 v1,v2, s, tuv, temp;
    float epsilon = 0.0001;
	v1 = tri.vtx1.xyz - tri.vtx0.xyz;
	v2 = tri.vtx2.xyz - tri.vtx0.xyz;
	s = rayOrigin  - tri.vtx0.xyz;
	vec3 q = cross(rayDir, v2);
	float a = dot(v1, q);
    float tReturn = -1;
	temp.x = determinant(mat3(s,v1,v2));
	temp.y = determinant(mat3(-rayDir,s,v2));
	temp.z = determinant(mat3(-rayDir,v1,s));

	if(determinant(mat3(-rayDir,v1,v2)) > 0.0f){
	tuv = (1 / determinant(mat3(-rayDir,v1,v2)))*temp;

	if(tuv.x >= 0.0){
	tReturn = tuv.x;
	}

        if((a > -epsilon && a < epsilon) || (tuv.y < 0.0)|| ((tuv.z < 0.0) || (tuv.y + tuv.z > 1.0)))
        {
            tReturn = -1;
        }
	}
		return tReturn;
}

float obbTest(vec3 rayDir, vec3 rayOrigin, obb_type o)
{
// FIXAAAAAAAAAAAAAAAAAAA
   float epsilon = 0.0001;
    float tReturn = -1;
    float tMin = -1000000;
    float tMax = 1000000;
    float t1;
    float t2;
    float e;
    float f;
    vec3 p = o.centre.xyz - rayOrigin;
    bool behindCamera = false;
    bool miss = false;
    vec4 obbDataArr[3];
    obbDataArr[0] = o.u_hu;
    obbDataArr[1] = o.v_hv;
    obbDataArr[2] = o.w_hw;

    for(int i = 0; i < 3 && miss == false
    && behindCamera == false; i++)
    {
        e = dot(obbDataArr[i].xyz, p);
        f = dot(obbDataArr[i].xyz, rayDir);

        if(abs(f) > epsilon)
        {
            float t1 = (e + obbDataArr[i].w)/f;
            float t2 = (e - obbDataArr[i].w)/f;
            if(t1 > t2)
            {
                float temp = t1;
                t1 = t2;
                t2 = temp;
            }
            if(t1 > tMin)
                tMin = t1;
            if(t2 < tMax)
                tMax = t2;
            if(tMin > tMax)
            {
                tReturn = -1;
                miss = true;
            }
            if(tMax < 0)
            {
                tReturn = -1; //behind camera
                behindCamera = true;
            }
        }
        else if(((-e) - obbDataArr[i].w) > 0.0 || ((-e) + obbDataArr[i].w) < 0.0)
        {
            tReturn = -1;
            miss = true;
        }
    }

    if(miss == false && behindCamera == false)
    {
        if(tMin > 0.0)
            tReturn = tMin;
        else
            tReturn = tMax;
    }


    return tReturn;
};

// everythin in World Space
vec4 shade(in vec3 pointOnSurface, in vec3 normal, in vec3 colour)
{
	// ambient (combine with material color!)
	// 0.2 is arbitrary, that is the value used in the model solution.
	vec4 final_colour = vec4(0.2,0.2,0.2,1);
	// diffuse, no attenuation.
	for (int i = 0; i < lightCount; i++)
	{
		vec4 light_pos = lights_data[i].position;
		vec4 light_colour = lights_data[i].colour;
		vec3 lightVector = normalize(light_pos.xyz - pointOnSurface);
		vec3 diffuseFactor = colour * dot(normal,lightVector);
		vec3 colorObject =	lights_data[i].colour.xyz * colour * diffuseFactor * 1/abs(light_pos.xyz - pointOnSurface);
		final_colour += vec4(colorObject,1);
		// IMPLEMENT HERE DIFFUSE SHADING

	}
	if(final_colour.x > 1)
		final_colour.x = 1;
	if(final_colour.y > 1)
		final_colour.y = 1;
	if(final_colour.z > 1)
		final_colour.z = 1;
	if(final_colour.w > 1)
		final_colour.w = 1;
	// UPDATE THIS LINE TO ACCOUNT FOR SATURATION (PIXEL COLOUR CANNOT GO OVER 1.0)
	return final_colour;
}