#version 430 core

const float PI  = 3.141592653589793238462643383;
const float TAU = 6.283185307179586476925286766;
const float EPSILON = 0.000001;

in vec3 frag_position;
in vec3 frag_normal;
in vec3 frag_tangent;
in vec3 frag_bitangent;
in vec2 frag_texture_coord;
in vec4 view_position;

out vec4 FragColor;

uniform int   has_albedo_texture;
uniform sampler2D albedo_texture;
uniform vec4      albedo_factor;
uniform int   has_metallic_texture;
uniform sampler2D metallic_texture;
uniform float     metallic_factor;
uniform int   has_roughness_texture;
uniform sampler2D roughness_texture;
uniform float     roughness_factor;
uniform int   has_normal_texture;
uniform sampler2D normal_texture;
uniform int   has_emissive_texture;
uniform sampler2D emissive_texture;
uniform vec4      emissive_factor;

uniform mat4 view_matrix;
uniform mat4 view_matrix_inverse;
uniform mat4 model_matrix;

uniform int has_shadow_map;
uniform sampler2DShadow shadow_map;
uniform mat4 shadow_matrix;

uniform int has_environment_map;
uniform samplerCube environment_map;


struct SunLight
{
	vec3 direction;
	float _padding0;
	vec3 color;
	float intensity;
};
struct PointLight
{
	vec3 position;
	float _padding0;
	vec3 color;
	float intensity;
};
struct SpotLight
{
	vec3 position;
	float inner_radius;
	vec3 color;
	float intensity;
	vec3 direction;
	float outer_radius;
};

#define MAX_SUN_LIGHTS 4
#define MAX_POINT_LIGHTS 128
#define MAX_SPOT_LIGHTS 128
layout (std140) uniform lights
{
	uint sun_light_count;
	uint point_light_count;
	uint spot_light_count;
	float ambient_light;
	SunLight sun_lights[MAX_SUN_LIGHTS];
	PointLight point_lights[MAX_POINT_LIGHTS];
	SpotLight spot_lights[MAX_SPOT_LIGHTS];
};

struct MaterialInfo
{
	vec4  Albedo;
	float Metallic;
	float Roughness;
	vec4  Emissive;
};

uniform float time;
uniform float width;
uniform float height;


float clamped_dot(vec3 v, vec3 w)
{
	return max( dot(v, w) , 0.0);
}

vec3 DiffuseBRDF(MaterialInfo material, vec3 L, vec3 V)
{
	// Lambertian
	return material.Albedo.xyz / PI;
}

float SpecularNDF(MaterialInfo material, vec3 N, vec3 H)
{
	// GGX
	float a  = material.Roughness * material.Roughness;
	float a2 = a * a;

	float n_dot_h  = clamped_dot(N, H);
	float n_dot_h2 = n_dot_h * n_dot_h;

	float d = n_dot_h2 * (a2 - 1.0) + 1.0;
	return a2 / (PI * d*d + EPSILON);
}

float SpecularGeometricAttenuation(MaterialInfo material, vec3 N, vec3 L, vec3 V, vec3 H)
{
	// Schlick
	float r = material.Roughness + 1;
	float k = r*r / 8.0;

	float n_dot_v = clamped_dot(N, V);
	float n_dot_l = clamped_dot(N, L);

	float denom_v = n_dot_v * (1.0 - k) + k;
	float denom_l = n_dot_l * (1.0 - k) + k;
	return (n_dot_v / denom_v) * (n_dot_l / denom_l);
}

vec3 SpecularFresnel(vec3 F0, vec3 V, vec3 H)
{
	// Schlick, Epic paper
	float v_dot_h = clamped_dot(V, H);
	float exp = (-5.55473 * v_dot_h - 6.98316) * v_dot_h;
	return F0 + (1.0 - F0) * pow(2, exp);
}

vec3 BRDF(MaterialInfo material, vec3 radiance, vec3 N, vec3 L, vec3 V, vec3 H)
{
	// Cook-Torrance
	vec3 F0 = mix(vec3(0.04), material.Albedo.xyz, material.Metallic);

	float D = SpecularNDF(material, N, H);
	vec3  F = SpecularFresnel(F0, N, H);
	float G = SpecularGeometricAttenuation(material, N, L, V, H);

	float n_dot_v = clamped_dot(N, V);
	float n_dot_l = clamped_dot(N, L);

	float specular_denom = (4.0 * n_dot_l * n_dot_v) + EPSILON;
	vec3 specular = D * F * G / specular_denom;
	vec3 diffuse  = (1.0 - F) * DiffuseBRDF(material, L, V);

	return (diffuse + specular) * radiance * n_dot_l;
}



void main()
{
	vec4 final_color = vec4(0,0,0,1);

	// PBR Parameters
	MaterialInfo material;
	material.Albedo = albedo_factor;
	if(has_albedo_texture != 0)
		material.Albedo = albedo_factor * texture(albedo_texture, frag_texture_coord);

	material.Metallic = metallic_factor;
	if(has_metallic_texture != 0)
		material.Metallic = metallic_factor * texture(metallic_texture, frag_texture_coord).b;

	material.Roughness = roughness_factor;
	if(has_roughness_texture != 0)
		material.Roughness = roughness_factor * texture(roughness_texture, frag_texture_coord).g;

	material.Emissive = emissive_factor;
	if(has_emissive_texture != 0)
		material.Emissive = emissive_factor * texture(emissive_texture, frag_texture_coord);

	vec3 Normal = normalize(frag_normal);
	if(has_normal_texture != 0)
	{
		vec3 normal_map = normalize(texture(normal_texture, frag_texture_coord).rgb * 2.0 - 1.0);
		mat3 TBN = mat3(normalize(frag_tangent), normalize(frag_bitangent), normalize(frag_normal));
		Normal = normalize(TBN * normal_map);
	}


	vec4 camera_position_4 = view_matrix_inverse * vec4(view_position.xy / view_position.w, 0, 1);
	vec3 camera_position = camera_position_4.xyz / camera_position_4.w;
	// vec4 pixel_position_4 = view_matrix_inverse * view_position;
	// vec3 pixel_position = pixel_position_4.xyz / pixel_position_4.w;
	// vec3 view_direction = normalize(pixel_position - camera_position);
	vec3 view_direction = normalize(frag_position - camera_position);
	vec3 V = -view_direction;
	vec3 N = Normal;



	// Sun lights
	for(uint i = 0; i < sun_light_count; i++)
	{
		SunLight light = sun_lights[i];

		float shadow = 1.0;
		if(has_shadow_map != 0)
		{
			shadow = 0;
			vec2 texel_size = 1.0 / textureSize(shadow_map, 0);
			float bias = 0.0001;

			vec4 from_light_view = shadow_matrix * vec4(frag_position, 1.0);
			from_light_view /= from_light_view.w;

			from_light_view = from_light_view * 0.5 + 0.5; // [-1,1] => [0,1] uv coords
			from_light_view.z = from_light_view.z - bias; // Bias to attenuate z fighting

			for(int x = -1; x <= 1; x++)
			for(int y = -1; y <= 1; y++)
			{
				float s = texture(shadow_map, from_light_view.xyz + vec3(x * texel_size.x, y * texel_size.y, 0));
				s = clamp(s, 0.0, 1.0);
				shadow += s;
			}
			shadow /= 9;
		}

		vec3 L = -light.direction;
		vec3 H = normalize(V + L);

		vec3 radiance = light.color * light.intensity * shadow;
		vec3 color = BRDF(material, radiance, N, L, V, H);
		final_color.xyz += color;
	}

	// Point lights
	for(uint i = 0; i < point_light_count; i++)
	{
		PointLight light = point_lights[i];

		vec3  diff  = light.position - frag_position;
		float dist2 = abs(dot(diff, diff));
		float attenuation = 1.0 / (dist2 + EPSILON);
		float intensity = light.intensity * attenuation;

		vec3 L = normalize(light.position - frag_position);
		vec3 H = normalize(V + L);

		vec3 radiance = light.color * intensity;
		vec3 color = BRDF(material, radiance, N, L, V, H);
		final_color.xyz += color;
	}

	// Spot lights
	for(uint i = 0; i < spot_light_count; i++)
	{
		SpotLight light = spot_lights[i];

		vec3 light_direction = normalize(frag_position - light.position);

		float dist2 = abs(dot(light.position, frag_position));
		float attenuation = 1.0 / (dist2 + EPSILON);
		float intensity = light.intensity * attenuation;

		float angle = acos(dot(light.direction, light_direction));
		float spot_factor = (angle - light.outer_radius) / (light.inner_radius - light.outer_radius);
		spot_factor = clamp(spot_factor, 0.0, 1.0);

		intensity *= spot_factor;

		vec3 L = normalize(light.position - frag_position);
		vec3 H = normalize(V + L);

		vec3 radiance = light.color * intensity;
		vec3 color = BRDF(material, radiance, N, L, V, H);
		final_color.xyz += color;
	}

	// Environment map light
	if(has_environment_map != 0)
	{
		vec3 L = normalize(reflect(-V, N));
		vec3 H = N;//normalize(V + L);

		float levels = textureQueryLevels(environment_map);
		vec3 radiance = textureLod(environment_map, L.xzy, material.Roughness * levels).rgb;

		//vec3 color = BRDF(material, radiance, N, L, V, H);
		vec3 F0 = mix(vec3(0.04), material.Albedo.xyz, material.Metallic);

		float D = SpecularNDF(material, N, H);
		vec3  F = SpecularFresnel(F0, N, H);
		float G = SpecularGeometricAttenuation(material, N, L, V, H);

		float n_dot_v = clamped_dot(N, V);
		float n_dot_l = clamped_dot(N, L);

		float specular_denom = (4.0 * n_dot_l * n_dot_v) + EPSILON;
		vec3 specular = D * F * G / specular_denom;
		vec3 diffuse  = (1.0 - F) * DiffuseBRDF(material, L, V);

		vec3 color = (diffuse /*+ specular*/) * radiance * n_dot_l;

		final_color.xyz += color;
	}

	// Ambient light
	final_color.xyz += ambient_light * material.Albedo.xyz;

	final_color.a = material.Albedo.a;

	// Emissive color
	final_color.xyz += material.Emissive.xyz * material.Emissive.a;

	FragColor = final_color;
	//FragColor.xyz = Normal;
}