uniform float u_elapsed_frames;
uniform vec2 u_scene_seed;
uniform vec2 u_camera_shift;

layout(location = 0) out vec4 o_color;

const float pi = 3.1415926535897932384626433832795;
const uint number_of_bounces = 8;
const float epsilon = 0.001;
const float max_distance = 10000.0;
const vec3 x_axis = vec3(1.0, 0.0, 0.0);
const vec3 y_axis = vec3(0.0, 1.0, 0.0);
const vec3 z_axis = vec3(0.0, 0.0, 1.0);
const vec3 origin = vec3(0.0);
const vec3 miss = vec3(-1.0);
const vec3 black = vec3(0.0);
const vec3 white = vec3(1.0);

struct sphere
{
	float r;
	vec3 c;
	vec3 albedo;
};

struct plane
{
	vec3 n;
	vec3 c;
};

float hash(vec2 seed)
{
	return fract(sin(dot(seed, vec2(12.9898, 78.233))) * 43758.5453);
}

float rand(vec2 seed) 
{
    return fract(sin(dot(seed, vec2(12.9898, 78.233) + u_elapsed_frames)) * 43758.5453);
}

vec3 palette(in float t, in vec3 a, in vec3 b, in vec3 c, in vec3 d)
{
    return a + b*cos(2.0 * pi * (c * t + d));
}

vec3 sphere_color(in float t)
{
	return palette(t * 0.5, vec3(0.5, 0.5, 0.5), vec3(0.5, 0.5, 0.5), vec3(1.0, 1.0, 1.0), vec3(0.00, 0.10, 0.20));
}

const uint number_of_spheres = 8;
const uint number_of_planes = 1;
const float radius = 0.5;

vec3 position_on_circle(int index)
{
	float theta = (float(index) / float(number_of_spheres - 2)) * 2.0 * pi;
	const float r = 1.45;
	float x = cos(theta) * r;
	float z = sin(theta) * r;

	return vec3(x, -0.5, z);
}

sphere spheres[] = {
	// Ground
	sphere(100.0, vec3(0.0, -100.9, 0.0), vec3(1.0, 0.8, 0.8)), 
	
	sphere(radius, position_on_circle(0), sphere_color(1.0 / float(number_of_spheres))),
	sphere(radius, position_on_circle(1), sphere_color(2.0 / float(number_of_spheres))),
	sphere(radius, position_on_circle(2), sphere_color(3.0 / float(number_of_spheres))),
	sphere(radius, position_on_circle(3), sphere_color(4.0 / float(number_of_spheres))),
	sphere(radius, position_on_circle(4), sphere_color(5.0 / float(number_of_spheres))),
	sphere(radius, position_on_circle(5), sphere_color(6.0 / float(number_of_spheres))),

	// Top
	sphere(0.95, vec3(0.0, 1.0, 0.0), vec3(0.8)) 
};

plane planes[] = {
	plane(y_axis, origin)
};

bool intersect_sphere(in sphere s, in vec3 ro, in vec3 rd, out vec3 hit)
{
	vec3 v = ro - s.c;
	float b = 2.0 * dot(v, rd);
	float c = dot(v, v) - s.r * s.r;

	float discriminant = b * b - 4.0 * c;

	if (discriminant < 0.0)
	{
		hit = miss;
		return false;
	}

	discriminant = sqrt(discriminant);
	float t0 = -b + discriminant;
	float t1 = -b - discriminant;

	float tf;
	if (t1 > epsilon)
	{
		tf = t1;
	} 
	else if (t0 > epsilon)
	{
		tf = t0;
	} 
	else 
	{	
		hit = miss;
		return false;
	}

	hit = ro + rd * tf * 0.5;

	return true;
}

bool intersect_plane(in plane p, in vec3 ro, in vec3 rd, out vec3 hit)
{
	float denom = dot(rd, p.n);

	if (abs(denom) > epsilon)
	{
		vec3 temp = p.c - ro;
		float t = dot(temp, p.n) / denom;

		hit = ro + rd * t;

		if (t >= epsilon && hit.y < 1.0)
		{
			return true;
		}
	}
	else
	{
		return false;
	}
}

vec3 rand_sphere(in vec3 normal, in vec2 seed) 
{
    float z = rand(seed);
    float r = sqrt(1.0 - z * z);
    float phi = 2.0 * pi * rand(seed + 100.0);
    float x = cos(phi) * r;
    float y = sin(phi) * r;
 
    vec3 major_axis;
    const float isqrt = 1.0 / sqrt(3.0);
    if (abs(normal.x) < isqrt) 
    { 
        major_axis = x_axis;
    } 
    else if (abs(normal.y) < isqrt) 
    { 
        major_axis = y_axis;
    } 
    else 
    {
        major_axis = z_axis;
    }

    vec3 u = normalize(cross(major_axis, normal));
    vec3 v = cross(normal, u);
    vec3 w = normal;

    return normalize(u * x + v * y + w * z);
}

vec3 shade_sky(in vec3 rd)
{
	float pct = 0.5 * (rd.y + 1.0);
    const vec3 sky = vec3(0.5, 0.7, 1.0);	

    return mix(white, sky, pct);
}

bool intersect_scene(in vec3 ro, in vec3 rd, out vec3 closest_point)
{
	float closest_t = max_distance;
	uint closest_index = 0;
	bool hit_anything = false;

	for(uint j = 0; j < number_of_spheres; ++j)
	{	
		vec3 hit;
		if (intersect_sphere(spheres[j], ro, rd, hit))
		{
			hit_anything = true;

			float t = length(ro - hit);
			if (t < closest_t)
			{	
				closest_t = t;
				closest_point = hit;
				closest_index = j;
			}
		} 
	}

	for(uint j = 0; j < number_of_planes; ++j)
	{
		vec3 hit;
		if (intersect_plane(planes[j], ro, rd, hit))
		{
			hit_anything = true;

			float t = length(ro - hit);
			if (t < closest_t)
			{
				closest_t = t;
				closest_point = hit;
				closest_index = j;
			}
		} 
	}

	return hit_anything;
}

mat3 lookat(in vec3 t, in vec3 p)
{
	vec3 k = normalize(t - p);
	vec3 i = cross(k, vec3(0.0, 1.0, 0.0));
	vec3 j = cross(i, k);
	return mat3(i, j, k);
}

vec3 trace()
{
	// Sample texture that contains random noise
	vec2 jitter = texture(sTD2DInputs[0], vUV.st).rg * 2.0 - 1.0;

	vec2 uv = vUV.st;
	uv = uv * 2.0 - 1.0;
	uv += (jitter / 1024.0) * 2.0;

	vec3 camera_position = vec3(0.0, -0.25, -4.0);
	vec3 ro = camera_position;
	ro.xy += u_camera_shift;

	vec3 rd = normalize(lookat(origin, ro) * vec3(uv, 1.0));
	
	vec3 attenuation = white;

	bool hit_anything = false;

	for (uint i = 0; i < number_of_bounces; ++i)
	{	
		float closest_t = max_distance;
		vec3 closest_point;
		uint closest_index = 0;
		
		bool valid_intersection = false;

		for(uint j = 0; j < number_of_spheres; ++j)
		{	
			vec3 hit;
			if (intersect_sphere(spheres[j], ro, rd, hit))
			{
				valid_intersection = true;
				hit_anything = true;

				float t = length(ro - hit);
				if (t < closest_t)
				{
					closest_t = t;
					closest_point = hit;
					closest_index = j;
				}
			} 
		}

		if (i == (number_of_bounces - 1))
		{
			attenuation *= black;
			break;
		}

		if (valid_intersection)
		{
			attenuation *= spheres[closest_index].albedo;
			vec3 n = normalize(closest_point - spheres[closest_index].c);

			vec3 rand_dir = rand_sphere(n, gl_FragCoord.xy + rd.z + i);

			float mat = hash(vec2(closest_index));
			if (closest_index == 0) mat = 0.8;

			ro = closest_point + n * epsilon;
			rd = normalize(mix(reflect(rd, n), rand_dir, mat));
		} 
		else 
		{
        	attenuation *= shade_sky(rd); 
        	break;
		}
	}

	if (!hit_anything)
	{
		return shade_sky(rd);
	}
	else 
	{
		return attenuation;
	}
}

void main()
{	
	// Path tracing routine
	vec3 trace_color = trace();

	// Apply gamma correction
	trace_color = pow(trace_color, vec3(1.0 / 2.2));

    o_color = vec4(trace_color, 1.0);
}