/*******************************************************************************************
*
*   rPBR [shader] - Physically based rendering fragment shader
*
*   Copyright (c) 2017 Victor Fisac
*
**********************************************************************************************/

#version 330

#define     MAX_REFLECTION_LOD      4.0
#define     MAX_DEPTH_LAYER         20
#define     MIN_DEPTH_LAYER         10

#define     MAX_LIGHTS              4
#define     LIGHT_DIRECTIONAL       0
#define     LIGHT_POINT             1

struct MaterialProperty {
    vec3 color;
    int useSampler;
    sampler2D sampler;
};

struct Light {
    int enabled;
    int type;
    vec3 position;
    vec3 target;
    vec4 color;
};

// Input vertex attributes (from vertex shader)
in vec3 fragPosition;
in vec2 fragTexCoord;
in vec3 fragNormal;
in vec3 fragTangent;
in vec3 fragBinormal;

// Input material values
uniform MaterialProperty albedo;
uniform MaterialProperty normals;
uniform MaterialProperty metalness;
uniform MaterialProperty roughness;
uniform MaterialProperty occlusion;
uniform MaterialProperty emission;
uniform MaterialProperty height;

// Input uniform values
uniform samplerCube irradianceMap;
uniform samplerCube prefilterMap;
uniform sampler2D brdfLUT;

// Input lighting values
uniform Light lights[MAX_LIGHTS];

// Other uniform values
uniform int renderMode;
uniform vec3 viewPos;
vec2 texCoord;

// Constant values
const float PI = 3.14159265359;

// Output fragment color
out vec4 finalColor;

vec3 ComputeMaterialProperty(MaterialProperty property);
float DistributionGGX(vec3 N, vec3 H, float roughness);
float GeometrySchlickGGX(float NdotV, float roughness);
float GeometrySmith(vec3 N, vec3 V, vec3 L, float roughness);
vec3 fresnelSchlick(float cosTheta, vec3 F0);
vec3 fresnelSchlickRoughness(float cosTheta, vec3 F0, float roughness);
vec2 ParallaxMapping(vec2 texCoords, vec3 viewDir);

// WARNING: There is some weird behaviour with this function, always returns black!
// Yes, I even tried: return texture(property.sampler, texCoord).rgb;
vec3 ComputeMaterialProperty(MaterialProperty property)
{
    vec3 result = vec3(0.0, 0.0, 0.0);

    if (property.useSampler == 1) result = texture(property.sampler, texCoord).rgb;
    else result = property.color;

    return result;
}

float DistributionGGX(vec3 N, vec3 H, float roughness)
{
    float a = roughness*roughness;
    float a2 = a*a;
    float NdotH = max(dot(N, H), 0.0);
    float NdotH2 = NdotH*NdotH;

    float nom = a2;
    float denom = (NdotH2*(a2 - 1.0) + 1.0);
    denom = PI*denom*denom;

    return nom/denom;
}

float GeometrySchlickGGX(float NdotV, float roughness)
{
    float r = (roughness + 1.0);
    float k = r*r/8.0;

    float nom = NdotV;
    float denom = NdotV*(1.0 - k) + k;

    return nom/denom;
}
float GeometrySmith(vec3 N, vec3 V, vec3 L, float roughness)
{
    float NdotV = max(dot(N, V), 0.0);
    float NdotL = max(dot(N, L), 0.0);
    float ggx2 = GeometrySchlickGGX(NdotV, roughness);
    float ggx1 = GeometrySchlickGGX(NdotL, roughness);

    return ggx1*ggx2;
}

vec3 fresnelSchlick(float cosTheta, vec3 F0)
{
    return F0 + (1.0 - F0)*pow(1.0 - cosTheta, 5.0);
}

vec3 fresnelSchlickRoughness(float cosTheta, vec3 F0, float roughness)
{
    return F0 + (max(vec3(1.0 - roughness), F0) - F0)*pow(1.0 - cosTheta, 5.0);
}

vec2 ParallaxMapping(vec2 texCoords, vec3 viewDir)
{
    // Calculate the number of depth layers and calculate the size of each layer
    float numLayers = mix(MAX_DEPTH_LAYER, MIN_DEPTH_LAYER, abs(dot(vec3(0.0, 0.0, 1.0), viewDir)));  
    float layerDepth = 1.0/numLayers;

    // Calculate depth of current layer
    float currentLayerDepth = 0.0;

    // Calculate the amount to shift the texture coordinates per layer (from vector P)
    // Note: height amount is stored in height material attribute color R channel (sampler use is independent)
    vec2 P = viewDir.xy*height.color.r; 
    vec2 deltaTexCoords = P/numLayers;

    // Store initial texture coordinates and depth values
    vec2 currentTexCoords = texCoords;
    float currentDepthMapValue = texture(height.sampler, currentTexCoords).r;

    while (currentLayerDepth < currentDepthMapValue)
    {
        // Shift texture coordinates along direction of P
        currentTexCoords -= deltaTexCoords;

        // Get depth map value at current texture coordinates
        currentDepthMapValue = texture(height.sampler, currentTexCoords).r;

        // Get depth of next layer
        currentLayerDepth += layerDepth;  
    }

    // Get texture coordinates before collision (reverse operations)
    vec2 prevTexCoords = currentTexCoords + deltaTexCoords;

    // Get depth after and before collision for linear interpolation
    float afterDepth = currentDepthMapValue - currentLayerDepth;
    float beforeDepth = texture(height.sampler, prevTexCoords).r - currentLayerDepth + layerDepth;

    // Interpolation of texture coordinates
    float weight = afterDepth/(afterDepth - beforeDepth);
    vec2 finalTexCoords = prevTexCoords*weight + currentTexCoords*(1.0 - weight);

    return finalTexCoords;
}

void main()
{
    // Calculate TBN and RM matrices
    mat3 TBN = transpose(mat3(fragTangent, fragBinormal, fragNormal));

    // Calculate lighting required attributes
    vec3 normal = normalize(fragNormal);
    vec3 view = normalize(viewPos - fragPosition);
    vec3 refl = reflect(-view, normal);

    // Check if parallax mapping is enabled and calculate texture coordinates to use based on height map
    // NOTE: remember that 'texCoord' variable must be assigned before calling any ComputeMaterialProperty() function
    if (height.useSampler == 1) texCoord = ParallaxMapping(fragTexCoord, view);
    else texCoord = fragTexCoord;   // Use default texture coordinates

    // Fetch material values from texture sampler or color attributes
    vec3 color = texture(albedo.sampler, texCoord).rgb; //ComputeMaterialProperty(albedo);
    vec3 metal = texture(metalness.sampler, texCoord).rgb; //ComputeMaterialProperty(metalness);
    vec3 rough = texture(roughness.sampler, texCoord).rgb; //ComputeMaterialProperty(roughness);
    vec3 emiss = texture(emission.sampler, texCoord).rgb; //ComputeMaterialProperty(emission);
    vec3 ao = texture(occlusion.sampler, texCoord).rgb; //ComputeMaterialProperty(occlusion);

    // Check if normal mapping is enabled
    if (normals.useSampler == 1)
    {
        // Fetch normal map color and transform lighting values to tangent space
        normal = texture(normals.sampler, texCoord).rgb; //ComputeMaterialProperty(normals);
        normal = normalize(normal*2.0 - 1.0);
        normal = normalize(normal*TBN);

        // Convert tangent space normal to world space due to cubemap reflection calculations
        refl = normalize(reflect(-view, normal));
    }

    // Calculate reflectance at normal incidence
    vec3 F0 = vec3(0.04);
    F0 = mix(F0, color, metal.r);

    // Calculate lighting for all lights
    vec3 Lo = vec3(0.0);
    vec3 lightDot = vec3(0.0);

    for (int i = 0; i < MAX_LIGHTS; i++)
    {
        if (lights[i].enabled == 1)
        {
            // Calculate per-light radiance
            vec3 light = vec3(0.0);
            vec3 radiance = lights[i].color.rgb;
            if (lights[i].type == LIGHT_DIRECTIONAL) light = -normalize(lights[i].target - lights[i].position);
            else if (lights[i].type == LIGHT_POINT)
            {
                light = normalize(lights[i].position - fragPosition);
                float distance = length(lights[i].position - fragPosition);
                float attenuation = 1.0/(distance*distance);
                radiance *= attenuation;
            }

            // Cook-torrance BRDF
            vec3 high = normalize(view + light);
            float NDF = DistributionGGX(normal, high, rough.r);
            float G = GeometrySmith(normal, view, light, rough.r);
            vec3 F = fresnelSchlick(max(dot(high, view), 0.0), F0);
            vec3 nominator = NDF*G*F;
            float denominator = 4*max(dot(normal, view), 0.0)*max(dot(normal, light), 0.0) + 0.001;
            vec3 brdf = nominator/denominator;

            // Store to kS the fresnel value and calculate energy conservation
            vec3 kS = F;
            vec3 kD = vec3(1.0) - kS;

            // Multiply kD by the inverse metalness such that only non-metals have diffuse lighting
            kD *= 1.0 - metal.r;

            // Scale light by dot product between normal and light direction
            float NdotL = max(dot(normal, light), 0.0);

            // Add to outgoing radiance Lo
            // Note: BRDF is already multiplied by the Fresnel so it doesn't need to be multiplied again
            Lo += (kD*color/PI + brdf)*radiance*NdotL*lights[i].color.a;
            lightDot += radiance*NdotL + brdf*lights[i].color.a;
        }
    }

    // Calculate ambient lighting using IBL
    vec3 F = fresnelSchlickRoughness(max(dot(normal, view), 0.0), F0, rough.r);
    vec3 kS = F;
    vec3 kD = 1.0 - kS;
    kD *= 1.0 - metal.r;

    // Calculate indirect diffuse
    vec3 irradiance = texture(irradianceMap, fragNormal).rgb;
    vec3 diffuse = color*irradiance;

    // Sample both the prefilter map and the BRDF lut and combine them together as per the Split-Sum approximation
    vec3 prefilterColor = textureLod(prefilterMap, refl, rough.r*MAX_REFLECTION_LOD).rgb;
    vec2 brdf = texture(brdfLUT, vec2(max(dot(normal, view), 0.0), rough.r)).rg;
    vec3 reflection = prefilterColor*(F*brdf.x + brdf.y);

    // Calculate final lighting
    vec3 ambient = (kD*diffuse + reflection)*ao;

    // Calculate fragment color based on render mode
    vec3 fragmentColor = ambient + Lo + emiss;                              // Physically Based Rendering

    if (renderMode == 1) fragmentColor = color;                             // Albedo
    else if (renderMode == 2) fragmentColor = normal;                       // Normals
    else if (renderMode == 3) fragmentColor = metal;                        // Metalness
    else if (renderMode == 4) fragmentColor = rough;                        // Roughness
    else if (renderMode == 5) fragmentColor = ao;                           // Ambient Occlusion
    else if (renderMode == 6) fragmentColor = emiss;                        // Emission
    else if (renderMode == 7) fragmentColor = lightDot;                     // Lighting
    else if (renderMode == 8) fragmentColor = kS;                           // Fresnel
    else if (renderMode == 9) fragmentColor = irradiance;                   // Irradiance
    else if (renderMode == 10) fragmentColor = reflection;                  // Reflection

    // Apply HDR tonemapping
    fragmentColor = fragmentColor/(fragmentColor + vec3(1.0));

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

    // Calculate final fragment color
    finalColor = vec4(fragmentColor, 1.0);
}