// Shader is looseley based on the ray tracing coding session by Inigo Quilez (www.iquilezles.org)

#version 450

#extension GL_ARB_separate_shader_objects : enable
#extension GL_ARB_shading_language_420pack : enable

layout (local_size_x =16, local_size_y = 16) in;
layout (binding = 0, rgba8) uniform image2D resultImage;

#define EPSILON 0.0001
#define MAXLEN 1000.0
#define PLANEID 1
#define SPHERECOUNT 3
#define SHADOW 0.5
#define RAYBOUNCES 1
#define REFLECTIONSTRENGTH 0.25

struct Camera {
    vec3 pos;   
    vec3 lookat;
    float fov; 
};

layout (binding = 1) uniform UBO 
{
	vec3 lightPos;
	float aspectRatio;
	vec4 fogColor;
	Camera camera;
	mat4 rotMat;
} ubo;

void reflectRay(inout vec3 rayD, in vec3 mormal)
{
    rayD = rayD + 2.0 * -dot(mormal, rayD) * mormal;
}

// Lighting calculations

float lightDiffuse(vec3 normal, vec3 lightDir) 
{
    return clamp(dot(normal, lightDir), 0.0, 1.0);
}

float lightSpecular(vec3 normal, vec3 lightDir)
{
    vec3 viewVec = normalize(ubo.camera.pos);
    vec3 halfVec = normalize(lightDir + viewVec);
    return pow(clamp(dot(normal, halfVec), 0.0, 1.0), 16.0);
}

// Primitives

// Basic material description
struct Material 
{
	vec3 diffuse;
	vec3 specular;
};

// Sphere
struct Sphere 
{
	int id;
	vec3 pos;
	float r;
	Material material;
} sphere;

Sphere spheres[SPHERECOUNT];

float sphereIntersect(in vec3 rayO, in vec3 rayD, in Sphere sphere)
{
	vec3 oc = rayO - sphere.pos;
	float b = 2.0 * dot(oc, rayD);
	float c = dot(oc, oc) - sphere.r*sphere.r;
	float h = b*b - 4.0*c;
	if (h < 0.0) 
	{
		return -1.0;
	}
	float t = (-b - sqrt(h)) / 2.0;
	return t;
}

vec3 sphereNormal(in vec3 pos, in Sphere sphere)
{
	return (pos - sphere.pos) / sphere.r;
}

// Plane

float planeIntersect(vec3 rayO, vec3 rayD)
{
	return -rayO.y/rayD.y;
}

vec3 planeNormal(in vec3 pos)
{
	return vec3(0.0, 1.0, 0.0);
}
	
int intersect(in vec3 rayO, in vec3 rayD, out float resT)
{
	int id = -1;
	resT = MAXLEN;

	for (int i = 0; i < SPHERECOUNT; i++)
	{
		float tSphere = sphereIntersect(rayO, rayD, spheres[i]);
		if (tSphere > EPSILON)
		{
			id = spheres[i].id;
			resT = tSphere;
			break;
		}
	}	
	
	float tplane = planeIntersect(rayO, rayD);
	if ((tplane > EPSILON) && (tplane < resT))
	{
		id = PLANEID;
		resT = tplane;
	}	
	
	return id;
}

float calcShadow(in vec3 rayO, in vec3 rayD, in int id)
{
	for (int i = 0; i < SPHERECOUNT; i++)
	{
		float tSphere = sphereIntersect(rayO, rayD, spheres[i]);
		if (tSphere > EPSILON)
		{
			return SHADOW;
		}
	}		
	return 1.0;
}

vec3 fog(in float t, in vec3 color)
{
    return mix(color, ubo.fogColor.rgb, clamp(sqrt(t*t)/20.0, 0.0, 1.0));
}

vec3 renderScene(inout vec3 rayO, inout vec3 rayD, inout int id)
{
	vec3 color = vec3(0.0);
	float t = 0.0;

	// Get intersected object ID
	int objectID = intersect(rayO, rayD, t);
	
	vec3 pos = rayO + t * rayD;
	vec3 lightVec = normalize(ubo.lightPos - pos);				
	vec3 normal;
	
	if (objectID == PLANEID)
	{
		normal = planeNormal(pos);
		float diffuse = clamp(dot(normal, lightVec), 0.0, 1.0);		
		color = vec3(1.0, 1.0, 1.0) * diffuse;
	}		
	else
	{
		for (int i = 0; i < SPHERECOUNT; i++)
		{
			if (objectID == spheres[i].id)
			{
				normal = sphereNormal(pos, spheres[i]);	
				float diffuse = lightDiffuse(normal, lightVec);
				float specular = lightSpecular(normal, lightVec);
				color = diffuse * spheres[i].material.diffuse + specular * spheres[i].material.specular;				
			}
		}
	}

	id = objectID;
	
	// Shadows
    color *= calcShadow(pos, lightVec, objectID);
	
	// Fog
	color = fog(t, color);	
	
	// Reflect ray for next render pass
	reflectRay(rayD, normal);
    rayO = pos;	
	
	return color;
}

void main()
{
	// Scene setup
	// todo : from ubo
	spheres[0].id = 2;
	spheres[0].pos = vec3(-2.25, 1.0, 0.0);
	spheres[0].r = 1.0;
	spheres[0].material.diffuse = vec3(1.0, 0.0, 0.0);
	spheres[0].material.specular = vec3(1.0, 1.0, 1.0);

	spheres[1].id = 3;
	spheres[1].pos = vec3(0.0, 2.5, 0.0);
	spheres[1].r = 1.0;
	spheres[1].material.diffuse = vec3(0.0, 0.0, 1.0);
	spheres[1].material.specular = vec3(1.0, 1.0, 1.0);	
	
	spheres[2].id = 4;
	spheres[2].pos = vec3(2.25, 1.0, 0.0);
	spheres[2].r = 1.0;
	spheres[2].material.diffuse = vec3(0.0, 1.0, 0.0);
	spheres[2].material.specular = vec3(1.0, 1.0, 1.0);
	
	ivec2 dim = imageSize(resultImage);
	vec2 uv = vec2(gl_GlobalInvocationID.xy) / dim;

	vec3 rayO = ubo.camera.pos;
	vec3 rayD = normalize(vec3((-1.0 + 2.0 * uv) * vec2(ubo.aspectRatio, 1.0), -1.0));
		
	// Basic color path
	int id = 0;
	vec3 finalColor = renderScene(rayO, rayD, id);
	
	bool reflections = true;
	// Reflection
	if (reflections)
	{
		for (int i = 0; i < RAYBOUNCES; i++)
		{
			vec3 reflectionColor = renderScene(rayO, rayD, id);
			finalColor = (1.0 - REFLECTIONSTRENGTH) * finalColor + REFLECTIONSTRENGTH * mix(reflectionColor, finalColor, 1.0 - REFLECTIONSTRENGTH);			
		}
	}
			
	imageStore(resultImage, ivec2(gl_GlobalInvocationID.xy), vec4(finalColor, 0.0));
}