ior_const.frag

#version 430 core

layout (location = 0) out vec4 OutColor;

in vec3 WorldPos;
in vec2 TexCoords;
in vec3 Normal;

layout (location = 0) uniform vec3 uEmissive;
layout (location = 1) uniform vec3 uAlbedo;
layout (location = 2) uniform float uRoughness1;
layout (location = 3) uniform float uRoughness2;
layout (location = 4) uniform vec3 uRealIoR;
layout (location = 5) uniform float uMetallic;
layout (location = 6) uniform float uHeightFactor;
layout (location = 7) uniform vec3 uLightColor;
layout (location = 8) uniform vec3 uLightPos;
layout (location = 9) uniform vec3 uEyePos;
layout (location = 10) uniform float uLightIntensity;
layout (location = 11) uniform vec3 uImagIoR;
layout (binding = 0) uniform sampler2D uAlbedoMap;
layout (binding = 1) uniform sampler2D uRoughnessMap;
layout (binding = 2) uniform sampler2D uIoRMap;
layout (binding = 3) uniform sampler2D uMetallicMap;
layout (binding = 4) uniform sampler2D uBumpMap;
layout (binding = 5) uniform sampler2D uAOMap;

float getRoughness(in vec2 texCoords);
float getMetallic(in vec2 texCoords);
vec3 getReflectance(in vec2 texCoords);
float getAO(in vec2 texCoords);
vec3 getAlbedo(in vec2 texCoords);
float IoRfromF0(float F0);
float ReIoRfromF0andImIoR(float F0, float ImIoR);

#define FLT_MIN 1.175494351e-38
#define FLT_MAX 3.402823466e+38
#define FLT_INF (1.0/0.0)

#define REFLECTANCE_SCALE_FACTOR 0.16
#ifndef _BRDF_DIFFUSE_OREN_NAYAR_INCLUDED_
#define _BRDF_DIFFUSE_OREN_NAYAR_INCLUDED_

float oren_nayar(in float _a2, in vec3 N, in vec3 L, in vec3 V, in float NdotL, in float NdotV)
{
    // theta - polar angles
    // phi - azimuth angles
    float a2 = _a2*0.5; //todo read about this
    vec2 AB = vec2(1.0, 0.0) + vec2(-0.5, 0.45) * vec2(a2, a2)/vec2(a2+0.33, a2+0.09);
    vec2 cos_theta = vec2(NdotL, NdotV);
    vec2 cos_theta2 = cos_theta*cos_theta;
    float sin_theta = sqrt((1.0 - cos_theta2.x) * (1.0 - cos_theta2.y)); //this is actually equal to (sin(theta_i) * sin(theta_r))
    float C = sin_theta / max(cos_theta.x, cos_theta.y);
    
    vec3 light_plane = normalize(L - cos_theta.x*N);
    vec3 view_plane = normalize(V - cos_theta.y*N);
    float cos_phi = max(0.0, dot(light_plane, view_plane));//not sure about this

    return (AB.x + AB.y * cos_phi * C) / 3.14159265359;
}

#endif
#ifndef _BRDF_SPECULAR_NDF_GGX_TROWBRIDGE_REITZ_INCLUDED_
#define _BRDF_SPECULAR_NDF_GGX_TROWBRIDGE_REITZ_INCLUDED_

float GGXTrowbridgeReitz(in float a2, in float NdotH)
{
    float denom = NdotH*NdotH * (a2 - 1.0) + 1.0;
    return a2 / (3.14159265359 * denom*denom);
}

#endif
#ifndef _BRDF_SPECULAR_GEOM_GGX_SMITH_INCLUDED_
#define _BRDF_SPECULAR_GEOM_GGX_SMITH_INCLUDED_

float _GGXSmith_G1_(in float a2, in float NdotX)
{
    return (2.0*NdotX) / (NdotX + sqrt(a2 + (1.0 - a2)*NdotX*NdotX));
}
float _GGXSmith_G1_wo_numerator(in float a2, in float NdotX)
{
    return 1.0 / (NdotX + sqrt(a2 + (1.0 - a2)*NdotX*NdotX));
}

float GGXSmith(in float a2, in float NdotL, in float NdotV)
{
    return _GGXSmith_G1_(a2, NdotL) * _GGXSmith_G1_(a2, NdotV);
}
float GGXSmith_wo_numerator(in float a2, in float NdotL, in float NdotV)
{
    return _GGXSmith_G1_wo_numerator(a2, NdotL) * _GGXSmith_G1_wo_numerator(a2, NdotV);
}

#endif
#ifndef _BRDF_SPECULAR_FRESNEL_FRESNEL_INCLUDED_
#define _BRDF_SPECULAR_FRESNEL_FRESNEL_INCLUDED_

vec3 FresnelSchlick(in vec3 F0, in float VdotH)
{
    float x = 1.0 - VdotH;
    return F0 + (1.0 - F0) * x*x*x*x*x;
}

// code from https://seblagarde.wordpress.com/2013/04/29/memo-on-fresnel-equations/
vec3 Fresnel_conductor(vec3 Eta, vec3 Etak, float CosTheta)
{  
   float CosTheta2 = CosTheta * CosTheta;
   float SinTheta2 = 1.0 - CosTheta2;
   vec3 Eta2 = Eta * Eta;
   vec3 Etak2 = Etak * Etak;

   vec3 t0 = Eta2 - Etak2 - SinTheta2;
   vec3 a2plusb2 = sqrt(t0 * t0 + 4 * Eta2 * Etak2);
   vec3 t1 = a2plusb2 + CosTheta2;
   vec3 a = sqrt(0.5 * (a2plusb2 + t0));
   vec3 t2 = 2 * a * CosTheta;
   vec3 Rs = (t1 - t2) / (t1 + t2);

   vec3 t3 = CosTheta2 * a2plusb2 + SinTheta2 * SinTheta2;
   vec3 t4 = t2 * SinTheta2;   
   vec3 Rp = Rs * (t3 - t4) / (t3 + t4);

   return 0.5 * (Rp + Rs);
}
float Fresnel_dielectric(in float Eta, in float CosTheta)
{
   float SinTheta2 = 1.0 - CosTheta * CosTheta;

   float t0 = sqrt(1.0 - (SinTheta2 / (Eta * Eta)));
   float t1 = Eta * t0;
   float t2 = Eta * CosTheta;

   float rs = (CosTheta - t1) / (CosTheta + t1);
   float rp = (t0 - t2) / (t0 + t2);

   return 0.5 * (rs * rs + rp * rp);
}

#endif

float diffuse(in float a2, in vec3 N, in vec3 L, in vec3 V, in float NdotL, in float NdotV);
vec3 specular(in float a2, in float NdotL, in float NdotV, in float NdotH, in float VdotH, in mat2x3 ior, in float metallic, out vec3 out_fresnel);
vec3 Fresnel_combined(in mat2x3 ior, in float cosTheta, in float metallic);

#define derivScaleFactor (1000.0)

vec3 calculateSurfaceGradient(in vec3 normal, in vec3 dpdx, in vec3 dpdy, in float dhdx, in float dhdy)
{
    vec3 r1 = cross(dpdy, normal);
    vec3 r2 = cross(normal, dpdx);
 
    return (r1*dhdx + r2*dhdy) / dot(dpdx, r1);
}

vec3 perturbNormal(in vec3 normal, in vec3 dpdx, in vec3 dpdy, in float dhdx, in float dhdy)
{
    return normalize(normal - calculateSurfaceGradient(normal, dpdx, dpdy, dhdx, dhdy));
}

float applyChainRule(in vec2 h_gradient, in vec2 dUVd_)
{
    return dot(h_gradient, dUVd_);
}

// Calculate the surface normal using the uv-space gradient (dhdu, dhdv)
vec3 calculateSurfaceNormal(in vec3 position, in vec2 uv, in vec3 normal, in vec2 h_gradient)
{
    vec3 dpdx = dFdx(position);
    vec3 dpdy = dFdy(position);
 
	vec2 dUVdx = dFdx(uv);
	vec2 dUVdy = dFdy(uv);
 
    float dhdx = applyChainRule(h_gradient, dUVdx);
    float dhdy = applyChainRule(h_gradient, dUVdy);
 
    return perturbNormal(normal, dpdx, dpdy, dhdx, dhdy);
}

void main() {
    vec3 N = normalize(Normal);

    const vec3 relLightPos = uLightPos - WorldPos;
    float NdotL = dot(N, relLightPos);

	vec3 color = vec3(0.0);
	if (NdotL>FLT_MIN)
	{
		const float relLightPosLen2 = dot(relLightPos, relLightPos);
		NdotL *= inversesqrt(relLightPosLen2);

		// there are better identities to get all of these
		const vec3 V = normalize(uEyePos - WorldPos);
		const vec3 L = normalize(relLightPos);

		float NdotV = dot(N, V);
        float LdotV = dot(L, V);

        // dots with H identities taken from Earl Hammon's PBR Diffuse Lighting GDC17 lecture
        const float LplusV_lenSq = 2.0 + 2.0*LdotV;
        const float LplusV_rcpLen = inversesqrt(LplusV_lenSq);
        const float NdotH = max((NdotL + NdotV) * LplusV_rcpLen, 0.0);
        const float VdotH = max(LplusV_rcpLen + LplusV_rcpLen*LdotV, 0.0);
        NdotV = max(NdotV, 0.0);
		// identity comment end (but also do you need to clamp all of them?)

		const vec2 texCoords = vec2(TexCoords.x, 1.0-TexCoords.y);

		const float a2 = getRoughness(texCoords);
		const float metallic = getMetallic(texCoords);
		const vec3 baseColor = getAlbedo(texCoords);
		const float ao = getAO(texCoords);
		vec3 realIoR = getReflectance(texCoords);
        mat2x3 IoR = mat2x3(realIoR, uImagIoR);

		float diffuse = diffuse(a2, N, L, V, NdotL, NdotV) * (1.0 - metallic);
		vec3 fresnel;
		vec3 spec = specular(a2, NdotL, NdotV, NdotH, VdotH, IoR, metallic, fresnel);

		color += ((diffuse * baseColor * ao * (vec3(1.0) - fresnel)) + spec) * NdotL * uLightIntensity * uLightColor / relLightPosLen2;
	}
	OutColor = vec4(color, 1.0);
}

float IoRfromF0(float F0) {
    return 2.0/(1.0 - sqrt(F0)) - 1.0;
}
float ReIoRfromF0andImIoR(float F0, float ImIoR)
{
   float T0 = 1.0-F0;
   float kT0 = ImIoR*T0;
   return (1.0+F0+sqrt(4.0*F0-kT0*kT0))/T0;
}
float getRoughness(in vec2 texCoords) {
	return uRoughness1;
}
float getMetallic(in vec2 texCoords) {
	return 0.0;
}
vec3 getReflectance(in vec2 texCoords) {
	return uRealIoR;
}
float getAO(in vec2 texCoords) {
	return 1.0;
}
vec3 getAlbedo(in vec2 texCoords) {
	return uAlbedo;
}
float diffuse(in float a2, in vec3 N, in vec3 L, in vec3 V, in float NdotL, in float NdotV) {
	return 1.0/3.14159265359;
}
vec3 specular(in float a2, in float NdotL, in float NdotV, in float NdotH, in float VdotH, in mat2x3 ior, in float metallic, out vec3 out_fresnel) {
	//assert(NdotL>FLT_MIN);
    out_fresnel = Fresnel_combined(ior, VdotH, metallic);
    if (NdotV<FLT_MIN)
        return vec3(0.0);
    if (a2<FLT_MIN)
		return vec3(/*NdotH>=(1.0-FLT_MIN) ? FLT_INF:*/0.0);

    float ndf = GGXTrowbridgeReitz(a2, NdotH);
    float geom = GGXSmith_wo_numerator(a2, NdotL, NdotV); // TODO: Correlated Smith!

    // Note: (4.0*NdotV*NdotL) denominator is cancelled by GGXSmith's numerator, thus the use of GGXSmith_wo_numerator()
    return ndf*geom*out_fresnel;
}

vec3 Fresnel_combined(in mat2x3 ior, in float cosTheta, in float metallic) {
    bvec3 is_inf = isinf(ior[0]);
    return mix(
        Fresnel_conductor(ior[0], ior[1], cosTheta),
        vec3(1.0),
        is_inf
    );
}