#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: cx 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);

	// set pixel to BACKGROUND colour
	vec4 finalPixelOut = vec4(0.2,0.,0.2, 1.0);

	if (objIndex >= 0) {
		// did we hit something?

		// IMPLEMENT HERE.
		// FIND POINT IN SPACE WHERE THE INTERSECTION HAPPENED.
		vec3 pointOnSurface = ray_dir.xyz * lastT + camPos;

		if (objType==SPHERE_TYPE)
			{
				// IMPLEMENT HERE
				// COMPUTE FINAL PIXEL COLOUR USING SHADE() FUNCTION
				vec3 n = normalize(pointOnSurface - spheres_data[objIndex].centre_radius.xyz);
				finalPixelOut = shade (pointOnSurface, n, spheres_data[objIndex].colour.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, planes_data[objIndex].colour.xyz);
		}
		else if (objType == TRIANGLE_TYPE)
		{
			// IMPLEMENT HERE
			// COMPUTE FINAL PIXEL COLOUR USING SHADE() FUNCTION
			vec3 u = triangles_data[objIndex].vtx1.xyz - triangles_data[objIndex].vtx0.xyz, v = triangles_data[objIndex].vtx2.xyz -  triangles_data[objIndex].vtx0.xyz;
			vec3 n = cross(u , v);
			finalPixelOut = shade(pointOnSurface, n, triangles_data[objIndex].colour.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 nor;
			bool found = 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 && found == false; i++)
			{
			    vec3 n = normalize(vecArr[i].xyz);
                vec3 centreToFace =  n * vecArr[i].w;
                vec3 pointOnFace = obb.centre.xyz + centreToFace;
				if(abs(dot(pointOnSurface - centreToFace, vecArr[i].xyz)) < epsilon){
					found = true;
					nor = n;
				}
			}
			finalPixelOut = shade(pointOnSurface, nor, obb.colour.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:
	vec3 rayOrigin = camPos;
	float t;
	// test all spheres (in geomCount.x)
	for (int i=0;i< geomCount.x;i++)
	{
		sphere_type sphere = spheres_data[i];
		// IMPLEMENT HERE.
		// TEST SPHERE, UPDATE lastT, objIndex and objType if necessary.
		t = sphereTest(ray_dir.xyz, rayOrigin, sphere.centre_radius);
		if (t != -1 ){
			if(t < lastT || lastT == -1){
			lastT = t;
			objType = SPHERE_TYPE;
			objIndex = i;
			}
		}
	}
	// lights (we render lights as spheres)
	for (int i=0;i<lightCount;i++)
	{
		light_type light = lights_data[i];

		t = sphereTest(ray_dir.xyz, rayOrigin, light.position);
		if(t != -1 ){
			if(t < lastT || lastT == -1){
			lastT = t;
			objType = LIGHT_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];
		// IMPLEMENT HERE.
		// TEST PLANE, UPDATE lastT, objIndex and objType if necessary.
		t = planeTest(ray_dir.xyz, rayOrigin, plane.normal_d);
		if( t != -1){
			if(t < lastT || lastT == -1){
			lastT = t;
			objType = PLANE_TYPE;
			objIndex = i;
			}
		}

	}
	// triangles
	for (int i = 0; i<geomCount.z; i++)
	{
		triangle_type triangle = triangles_data[i];
		// IMPLEMENT HERE.
		// TEST TRIANGLE, UPDATE lastT, objIndex and objType if necessary.
		t = triangleTest(ray_dir.xyz, rayOrigin, triangle);
		if(t != -1){
			if(t < lastT || lastT == -1){
			lastT = t;
			objType = TRIANGLE_TYPE;
			objIndex = i;
			}
		}
	}
	// OBBss
	for (int i = 0; i < geomCount.w; i++)
	{
		obb_type obb = obb_data[i];
		// IMPLEMENT HERE.
		// TEST OBB, UPDATE lastT, objIndex and objType if necessary.
		t = obbTest(ray_dir.xyz, rayOrigin, obb);
		if(t != -1){
			if(t < lastT || lastT == -1){
			lastT = t;
			objType = OBB_TYPE;
			objIndex = i;
			}
		}
	}
}

float sphereTest(in vec3 rayDir, in vec3 rayOrigin, in vec4 centre_radius)
{
	// IMPLEMENT HERE
	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{
		float t = min(-p + sqrt(p*p - q), -p - sqrt(p*p - q));
		if(t > 0)
			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;
	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;
        }
	}
	if(tReturn >= 0){
		return tReturn;
	}
	else{

	return -1;
	}
}

float obbTest(vec3 rayDir, vec3 rayOrigin, obb_type o)
{
	float epsilon = 0.0000001;
	float t_min= -1000000,
		t_max = 1000000 ;
	vec3 p = o.centre.xyz - rayOrigin;

	float t,e,f,w,t1,t2;
    vec4 obbArr[3];
    obbArr[0] = o.u_hu;
    obbArr[1] = o.v_hv;
    obbArr[2] = o.w_hw;

    for(int i = 0; i < 3; i++)
    {
        e = dot(obbArr[i].xyz, p);
        f = dot(obbArr[i].xyz, rayDir);
        if(abs(f) > epsilon)
        {
            w = obbArr[i].w;
            float t1 = (e+w)/f;
            float t2 = (e-w)/f;
            if(t1>t2){
            t=t1;
            t1=t2;
            t2=t;
        }
            if(t1 > t_min)
                t_min = t1;

            if(t2 < t_max)
                t_max = t2;

            if(t_min > t_max)
                return  -1;

            if(t_max < 0)
                return  -1;
        }
        else if(-e-obbArr[i].w > 0.0 || -e+obbArr[i].w < 0.0)
            return  -1;
    }
        if(t_min > 0.0)
            return t_min;

        else
            return t_max;

};

// 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.
	 // 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) * vec4(colour,1.0); //* vec4(colour,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 lightVec = normalize(light_pos.xyz - pointOnSurface);
        float diffuseFactor = dot(normalize(normal),lightVec);
        float lenght = distance(pointOnSurface, light_pos.xyz);
        vec3 colorObject = diffuseFactor*(light_colour.xyz*colour)*light_pos.w/lenght; //* (1.0/  abs(light_pos.xyz - pointOnSurface));
        final_colour += vec4(colorObject,1.0);

    }
    final_colour = min( final_colour, vec4(1.0,1.0,1.0,1.0));

    // UPDATE THIS LINE TO ACCOUNT FOR SATURATION (PIXEL COLOUR CANNOT GO OVER 1.0)
    return final_colour;
}