#version 450

#define PI 3.1415926535f

float GGXPartialGeometry(float NV, float roughness)
{
    float k = (roughness + 1.0f) * (roughness + 1.0f) * 0.125f;
    return NV / (NV * (1.0f - k) + k);
}

float GGXDistribution(float NH, float roughness)
{
    float roughness2 = roughness * roughness;
    float alpha2 = roughness2 * roughness2;
    float distr = (NH * NH) * (alpha2 - 1.0f) + 1.0f;
    float distr2 = distr * distr;
    float totalDistr = alpha2 / (PI * distr2);
    return totalDistr;
}

float GGXSmith(float NV, float NL, float roughness)
{
    return GGXPartialGeometry(NV, roughness)* GGXPartialGeometry(NL, roughness);
}

vec3 fresnelSchlick(vec3 F0, float HV)
{
    float p = (-5.55473f * HV - 6.98136f) * HV;
    return F0 + (vec3(1.0f) - F0) * pow(2.0f, p);
}

mat3 computeSampleTransform(vec3 normal)
{
    vec3 up = abs(normal.z) < 0.999f ? vec3(0.0f, 0.0f, 1.0f) : vec3(1.0f, 0.0f, 0.0f);

    mat3 w;
    w[0] = normalize(cross(normal, up));
    w[1] = cross(normal, w[0]);
    w[2] = normal;
    return w;
}

vec3 GGXImportanceSample(vec2 Xi, float roughness, vec3 normal)
{
    mat3 transform = computeSampleTransform(normal);

    float alpha = roughness * roughness;
    float phi = 2.0f * PI * Xi.x;
    float cosTheta = sqrt((1.0f - Xi.y) / (1.0f + (alpha * alpha - 1.0f) * Xi.y));
    float sinTheta = sqrt(1.0f - cosTheta * cosTheta);
    vec3 H = vec3(sinTheta * cos(phi), sinTheta * sin(phi), cosTheta);
    return transform * H;
}

float computeA(samplerCube tex, int sampleCount)
{
    vec2 sizes = textureSize(tex, 0);
    float A = 0.5f * log2(sizes.x * sizes.y / float(sampleCount));
    return A;
}

float computeLOD(float A, float pdf, vec3 lightDirection)
{
    float du = 2.0f * 1.2f * abs(lightDirection.y + 1.0f);
    return max(0.0f, A - 0.5f * log2(pdf * du * du) + 1.0f);
}

float radicalInverseVDC(uint bits)
{
    bits = (bits << 16u) | (bits >> 16u);
    bits = ((bits & 0x55555555u) << 1u) | ((bits & 0xAAAAAAAAu) >> 1u);
    bits = ((bits & 0x33333333u) << 2u) | ((bits & 0xCCCCCCCCu) >> 2u);
    bits = ((bits & 0x0F0F0F0Fu) << 4u) | ((bits & 0xF0F0F0F0u) >> 4u);
    bits = ((bits & 0x00FF00FFu) << 8u) | ((bits & 0xFF00FF00u) >> 8u);
    return float(bits) * 2.3283064365386963e-10; // / 0x100000000
}

vec2 sampleHammersley(uint i, float invSampleCount)
{
    return vec2(i * invSampleCount, radicalInverseVDC(i));
}

vec3 GGXCookTorranceSampled(vec3 normal, vec3 lightDirection, vec3 viewDirection, float roughness, float metallic, vec3 albedo, out vec3 FK, out float pdf)
{
    vec3 H = normalize(viewDirection + lightDirection);
    float NV = dot(normal, viewDirection);
    float NL = dot(normal, lightDirection);
    float NH = dot(normal, H);
    float HV = dot(H, viewDirection);

    if (NV < 0.0f) return vec3(0.0f);
    if (NL < 0.0f) return vec3(0.0f);

    vec3 F0 = mix(vec3(0.04f), albedo, metallic);

    float G = GGXSmith(NV, NL, roughness);
    float D = GGXDistribution(NH, roughness);
    vec3 F = fresnelSchlick(F0, HV);

    FK = F;
    pdf = D * NH / (4.0f * HV);

    vec3 specular = G * F * HV / (NV * NH);
    return max(vec3(0.0f), specular);
}vec3 reconstructWorldPosition(float depth, vec2 texcoord, mat4 invViewProjMatrix)
{
	vec4 normPosition = vec4(2.0f * texcoord - vec2(1.0f), depth, 1.0f);
	vec4 worldPosition = invViewProjMatrix * normPosition;
	worldPosition /= worldPosition.w;
	return worldPosition.xyz;
}

float calcShadowFactor2D(vec4 fragPosLight, sampler2D depthMap, float bias, int blurIterations)
{
	vec3 projCoords = fragPosLight.xyz / fragPosLight.w;
	if (projCoords.z > 0.99f) return 1.0f; // do not handle corner cases, assume now shadows
	float currentDepth = projCoords.z - bias;
	float shadowFactor = 0.0f;
	vec2 texelSize = 1.0f / textureSize(depthMap, 0);
	for (int x = -blurIterations; x <= blurIterations; x++)
	{
		for (int y = -blurIterations; y <= blurIterations; y++)
		{
			float pcfDepth = textureLod(depthMap, projCoords.xy + vec2(x, y) * 1.5f * texelSize, 0).r;
			shadowFactor += (currentDepth > pcfDepth) ? 0.0f : 1.0f;
		}
	}
	int iterations = 2 * blurIterations + 1;
	shadowFactor /= float(iterations * iterations);
	return shadowFactor;
}

const int POINT_LIGHT_SAMPLES = 20;
vec3 sampleOffsetDirections[POINT_LIGHT_SAMPLES] = vec3[]
(
	vec3(1, 1, 1), vec3(1, -1, 1), vec3(-1, -1, 1), vec3(-1, 1, 1),
	vec3(1, 1, -1), vec3(1, -1, -1), vec3(-1, -1, -1), vec3(-1, 1, -1),
	vec3(1, 1, 0), vec3(1, -1, 0), vec3(-1, -1, 0), vec3(-1, 1, 0),
	vec3(1, 0, 1), vec3(-1, 0, 1), vec3(1, 0, -1), vec3(-1, 0, -1),
	vec3(0, 1, 1), vec3(0, -1, 1), vec3(0, -1, -1), vec3(0, 1, -1)
);

float CalcShadowFactor3D(vec3 fragToLightRay, vec3 viewDist, float zfar, float bias, samplerCube depthMap)
{
	float invZfar = 1.0f / zfar;
	float currentDepth = length(fragToLightRay);
	currentDepth = (currentDepth - bias) * invZfar;
	float diskRadius = (1.0f + invZfar) * 0.04f;
	float shadowFactor = 0.0f;

	for (int i = 0; i < POINT_LIGHT_SAMPLES; i++)
	{
		float closestDepth = textureLod(depthMap, sampleOffsetDirections[i] * diskRadius - fragToLightRay, 0).r;
		shadowFactor += (currentDepth > closestDepth) ? 0.0f : 1.0f;
	}
	shadowFactor /= float(POINT_LIGHT_SAMPLES);
	return shadowFactor;
}

vec3 calcReflectionColor(samplerCube reflectionMap, mat3 reflectionMapTransform, vec3 viewDir, vec3 normal)
{
	vec3 I = -viewDir;
	vec3 reflectionRay = reflect(I, normal);
	reflectionRay = reflectionMapTransform * reflectionRay;
	vec3 color = texture(reflectionMap, reflectionRay).rgb;
	return color;
}

vec4 worldToFragSpace(vec3 v, mat4 viewProj)
{
	vec4 proj = viewProj * vec4(v, 1.0f);
	proj.xyz /= proj.w;
	proj.xy = proj.xy * 0.5f + vec2(0.5f);
	return proj;
}

struct FragmentInfo
{
	vec3 albedo;
	float ambientOcclusion;
	float emmisionFactor;
	float roughnessFactor;
	float metallicFactor;
	float depth;
	vec3 normal;
	vec3 position;
};

struct EnvironmentInfo
{
	samplerCube skybox;
	samplerCube irradiance;
	mat3 skyboxRotation;
	float intensity;
};

FragmentInfo getFragmentInfo(vec2 texCoord, sampler2D albedoTexture, sampler2D normalTexture, sampler2D materialTexture, sampler2D depthTexture, mat4 invViewProjMatrix)
{
	FragmentInfo fragment;

	fragment.normal = normalize(texture(normalTexture, texCoord).rgb - vec3(0.5f));
	vec4 albedo = texture(albedoTexture, texCoord).rgba;
	vec4 material = texture(materialTexture, texCoord).rgba;
	fragment.depth = texture(depthTexture, texCoord).r;

	fragment.albedo = albedo.rgb;
	fragment.ambientOcclusion = albedo.a;
	fragment.emmisionFactor = material.r / (1.0f - material.r);
	fragment.roughnessFactor = material.g;
	fragment.metallicFactor = material.b;

	fragment.position = reconstructWorldPosition(fragment.depth, texCoord, invViewProjMatrix);

	return fragment;
}


vec3 calculateLighting(FragmentInfo fragment, vec3 viewDirection, vec3 lightDirection, EnvironmentInfo environment, int GGXSamples, vec3 lightColor, float ambientFactor, float shadowFactor)
{
	vec3 normLightDirection = normalize(lightDirection);

    vec3 specularColor = vec3(0.0f);
    vec3 FKtotal = vec3(0.0f);

    float roughness =  clamp(fragment.roughnessFactor, 0.01f, 0.99f);
    float metallic = clamp(fragment.metallicFactor, 0.01f, 0.99f);

    float invEnvironmentSampleCount = 1.0f / float(GGXSamples);
    float A = computeA(environment.skybox, GGXSamples);
    for (uint i = 0; i < uint(GGXSamples); i++)
    {
        vec2 Xi = sampleHammersley(i, invEnvironmentSampleCount);
        vec3 H = GGXImportanceSample(Xi, roughness, fragment.normal);
        vec3 direction = 2.0f * dot(viewDirection, H) * H - viewDirection;

        vec3 FK;
        float pdf;

        vec3 specularK = GGXCookTorranceSampled(fragment.normal, direction, viewDirection, roughness, fragment.metallicFactor, fragment.albedo, FK, pdf);
        FKtotal += FK;
        float lod = computeLOD(A, pdf, direction);

        vec3 sampleDirection = environment.skyboxRotation * direction;
        vec3 sampledColor = textureLod(environment.skybox, sampleDirection, lod).rgb;
        specularColor += specularK * sampledColor;
    }
    specularColor *= invEnvironmentSampleCount * environment.intensity;
    FKtotal *= invEnvironmentSampleCount;
    vec3 irradianceColor = calcReflectionColor(environment.irradiance, environment.skyboxRotation, viewDirection, fragment.normal);

    vec3 diffuseColor = fragment.albedo * (1.0f - metallic) * (1.0f - FKtotal) * irradianceColor;
    vec3 ambientColor = diffuseColor * ambientFactor;

    float NL = max(dot(normLightDirection, fragment.normal), 0.0f);

    diffuseColor *= NL;
    specularColor *= NL;

    vec3 totalColor = ambientColor + shadowFactor * (diffuseColor + specularColor);
    return totalColor * lightColor * fragment.ambientOcclusion;
}

out vec4 OutColor;

uniform sampler2D albedoTex;
uniform sampler2D normalTex;
uniform sampler2D materialTex;
uniform sampler2D depthTex;

in PointLightInfo
{
	vec3 position;
	float radius;
	vec4 color;
} pointLight;

struct PointLight
{
	vec3 position;
	float radius;
	vec4 color;
};

struct Camera
{
	vec3 position;
	mat4 invViewProjMatrix;
	mat4 viewProjMatrix;
};

uniform samplerCube lightDepthMap;
uniform bool castsShadows;
uniform Camera camera;
uniform int pcfDistance;
uniform int lightSamples;
uniform vec2 viewportSize;

uniform EnvironmentInfo environment;

vec3 calcColorUnderPointLight(FragmentInfo fragment, PointLight light, vec3 viewDir, samplerCube map_shadow, bool computeShadow)
{
	vec3 lightPath = light.position - fragment.position;
	float lightDistance = length(lightPath);

	if (light.radius < lightDistance || isnan(lightDistance))
		discard;

	float shadowFactor = 1.0f;
	if (computeShadow) { shadowFactor = CalcShadowFactor3D(lightPath, viewDir, light.radius, 0.15f, map_shadow); }
	float intensity = clamp(pow(1.0f - pow(lightDistance / light.radius, 4.0f), 2.0f), 0.0f, 1.0f) / (lightDistance * lightDistance + 1.0f);

	return calculateLighting(fragment, viewDir, lightPath, environment, lightSamples, intensity * light.color.rgb, light.color.a, shadowFactor);
}

void main()
{
	vec2 TexCoord = gl_FragCoord.xy / viewportSize;
	FragmentInfo fragment = getFragmentInfo(TexCoord, albedoTex, normalTex, materialTex, depthTex, camera.invViewProjMatrix);

	float fragDistance = length(camera.position - fragment.position);
	vec3 viewDirection = normalize(camera.position - fragment.position);

	PointLight light;
	light.position = pointLight.position;
	light.radius = pointLight.radius;
	light.color = pointLight.color;

	vec3 totalColor = vec3(0.0f);
	totalColor += calcColorUnderPointLight(fragment, light, viewDirection, lightDepthMap, castsShadows);

	OutColor = vec4(totalColor, 1.0f);
}