path: root/shaders/cook_torrance.frag

precision highp float;

uniform vec4  BaseColorFactor;
uniform float MetallicFactor;
uniform float RoughnessFactor;
uniform vec3  EmissiveFactor;

#ifdef HAS_ALBEDO_MAP
uniform sampler2D BaseColorTexture;
#endif
#ifdef HAS_METALLIC_ROUGHNESS_MAP
uniform sampler2D MetallicRoughnessTexture;
#endif
#ifdef HAS_EMISSIVE_MAP
uniform sampler2D EmissiveTexture;
#endif
#ifdef HAS_OCCLUSION_MAP
uniform sampler2D AmbientOcclusionTexture;
#endif
#ifdef HAS_NORMAL_MAP
uniform sampler2D NormalMap;
#endif

// TODO: Handle case in which there is no sky. Pass a boolean.
uniform samplerCube Sky;
uniform samplerCube IrradianceMap;
uniform samplerCube PrefilteredEnvironmentMap;
uniform sampler2D   BRDFIntegrationMap;
uniform float       MaxReflectionLOD;

uniform vec3 CameraPosition; // World space.

// World-space position, normal and tangent.
in vec3 Position;
#ifdef HAS_NORMALS
in vec3 Normal;
#endif
#ifdef HAS_TANGENTS
in vec4 Tangent;
#endif
#ifdef HAS_TEXCOORDS
in vec2 Texcoord;
#endif

layout (location = 0) out vec4 Colour;

#define PI 3.1415926535897932384626433832795
#define INV_PI 0.3183098861837907

/// Transform a normal map sample into world space.
///
/// |normalWs| is the surface normal in world space.
/// |normalMapSample| is the normal map sample, not necessarily normalized.
///
/// TODO: Move to "normal.h"
#if defined(HAS_NORMAL_MAP) && (defined(HAS_TANGENTS) || defined(HAS_TEXCOORDS))
vec3 get_ws_normal(vec3 normalWs, vec3 normalMapSample) {
  vec3 N = normalize(Normal);
#ifdef HAS_TANGENTS
  //vec3 T = normalize(tangent.xyz - dot(tangent.xyz, N) * N);
  vec3 T = Tangent.xyz;
  vec3 B = Tangent.w * cross(N, T);
#elif HAS_TEXCOORDS // No tangents, but must have texcoords.
  vec3 pos_dx = dFdx(Position);
  vec3 pos_dy = dFdy(Position);
  // vec3 uv_dx = vec3(dFdx(Texcoord), 0.0);
  // vec3 uv_dy = vec3(dFdy(Texcoord), 0.0);
  vec3 uv_dx = dFdx(vec3(Texcoord, 0.0));
  vec3 uv_dy = dFdy(vec3(Texcoord, 0.0));
  vec3 T = (uv_dy.t * pos_dx - uv_dx.t * pos_dy) /
           (uv_dx.s * uv_dy.t - uv_dy.s * uv_dx.t);
  // vec3 T = pos_dx * uv_dy.t - pos_dy * uv_dx.t;
  T = normalize(T - dot(T, N) * N);
  vec3 B = normalize(cross(N, T));
#endif

  if (gl_FrontFacing == false) {
    T = -T;
    B = -B;
    N = -N;
  }

  vec3 s = normalMapSample;
  //return normalize(s.x * T + s.y * B + s.z * N);
  return normalize(mat3(T,B,N) * s);
}
#endif // HAS_TANGENTS || HAS_TEXCOORDS

float trowbridge_reitz_GGX(float roughness, float NdotH) {
  float a = roughness * roughness;
  float a2 = a * a;
  float d = NdotH * NdotH * (a2 - 1.0) + 1.0;
  return a2 / (PI * d * d);
}

float geometry_schlick_GGX(float k, float NdotV) {
  return NdotV / (NdotV * (1.0 - k) + k);
}

float geometry_smith(float roughness, float NdotL, float NdotV) {
  float k = roughness * roughness / 2.0; // IBL
  return geometry_schlick_GGX(k, NdotV) * geometry_schlick_GGX(k, NdotL);
}

vec3 fresnel_schlick(vec3 F0, float HdotV) {
  return F0 + (1.0 - F0) * pow(clamp(1.0 - HdotV, 0.0, 1.0), 5.0);
}

vec3 fresnel_schlick_roughness(vec3 F0, float NdotV, float roughness) {
  return F0
    + (max(vec3(1.0 - roughness), F0) - F0)
    * pow(clamp(1.0 - NdotV, 0.0, 1.0), 5.0);
}

// Cook-Torrance BRDF for a single light direction.
vec3 cook_torrance(
    vec3 albedo, float metallic, float roughness,
    float NdotL, float NdotV, float NdotH, float HdotV) {
  vec3 F0 = mix(vec3(0.04), albedo, metallic);
  float D = trowbridge_reitz_GGX(roughness, NdotH);
  vec3  F = fresnel_schlick(F0, HdotV);
  float G = geometry_smith(roughness, NdotL, NdotV);
  vec3 Kd = mix(vec3(1.0) - F, vec3(0.0), metallic);
  vec3 diffuse = Kd*albedo*INV_PI;
  // Take a max to prevent division by 0 when either dot product is 0.
  vec3 specular = (D*F*G) / max(4.0 * NdotV * NdotL, 0.0001);
  return diffuse + specular;
}

// Cook-Torrance BRDF for IBL.
vec3 cook_torrance_IBL(
    vec3 albedo, float metallic, float roughness, float occlusion,
    float NdotV,
    vec3 irradiance, vec3 prefiltered_env, vec2 BRDF_env) {
  vec3 F0 = mix(vec3(0.04), albedo, metallic);
  vec3  F = fresnel_schlick_roughness(F0, NdotV, roughness);
  vec3 Kd = mix(vec3(1.0) - F, vec3(0.0), metallic);
  vec3 diffuse = Kd * albedo * INV_PI * irradiance;
  vec3 specular = prefiltered_env * (F * BRDF_env.x + BRDF_env.y);
  return occlusion * (diffuse + specular);
}

void main()
{
  // TODO: Also use the specular F0 map from the model, and emissive. Make sure
  // to use all maps.
  // https://sketchfab.com/models/b81008d513954189a063ff901f7abfe4
#ifdef HAS_NORMAL_MAP
  vec3 normalMapSample = texture(NormalMap, Texcoord).xyz * 2.0 - 1.0;
  vec3 N = get_ws_normal(Normal, normalMapSample);
#elif HAS_NORMALS
  vec3 N = normalize(Normal);
#endif
  vec3 V = normalize(CameraPosition - Position);
  vec3 R = reflect(-V, N);
  // Not needed for IBL.
  //vec3 L = N;
  //vec3 H = normalize(L + V);

  float NdotV = max(0.0, dot(N, V));
  // Not needed for IBL.
  //float NdotL = max(0.0, dot(N,L));
  //float NdotH = max(0.0, dot(N,H));
  //float HdotV = clamp(dot(H,V), 0.0, 1.0);  // Clamp to prevent black spots.

  // TODO: BaseColorFactor and BaseColorTexture are vec4/rgba quantities
  // respectively. Handle the alpha channel.
  // TODO: Other factors.
#ifdef HAS_ALBEDO_MAP
  vec3 albedo = vec3(BaseColorFactor) * texture(BaseColorTexture, Texcoord).rgb;
#else
  vec3 albedo = vec3(BaseColorFactor);
#endif
#ifdef HAS_METALLIC_ROUGHNESS_MAP
  // Spec: "Its green channel contains roughness values and its blue channel
  // contains metalness values."
  vec2 metal_roughness
    = vec2(MetallicFactor, RoughnessFactor)
    * texture(MetallicRoughnessTexture, Texcoord).bg;
#else
  vec2 metal_roughness = vec2(MetallicFactor, RoughnessFactor);
#endif
#ifdef HAS_EMISSIVE_MAP
  vec3 emissive = EmissiveFactor * texture(EmissiveTexture, Texcoord).rgb;
#else
  vec3 emissive = EmissiveFactor;
#endif
#ifdef HAS_OCCLUSION_MAP
  float occlusion = texture(AmbientOcclusionTexture, Texcoord).r;
#else
  float occlusion = 1.0;
#endif
  float metallic  = metal_roughness.x;
  float roughness = metal_roughness.y;

  // For a single light direction:
  // vec3 brdf = cook_torrance(albedo, metallic, roughness, NdotL, NdotV, NdotH, HdotV);
  // vec3 Li = texture(Sky, N).rgb;
  // vec3 colour = brdf * Li * NdotL;

  // For IBL:
  vec3 irradiance = texture(IrradianceMap, N).rgb;
  vec3 prefiltered_env = textureLod(PrefilteredEnvironmentMap, R, roughness * MaxReflectionLOD).rgb;
  vec2 BRDF_env = texture(BRDFIntegrationMap, vec2(NdotV, roughness)).rg;
  vec3 colour = cook_torrance_IBL(
    albedo, metallic, roughness, occlusion, NdotV, irradiance, prefiltered_env, BRDF_env);

  colour += emissive;

  // Reinhard tone mapping.
  colour = colour / (colour + vec3(1.0));

  // Gamma correction.
  colour = pow(colour, vec3(1.0 / 2.2));

  // Debugging.
  //
  // Normal texture.
  //colour = normalMapSample * 0.5 + 0.5;
  //
  // Geometry normal.
  //colour = normalize(Normal) * 0.5 + 0.5;
  //
  // Shading normal.
  //colour = N * 0.5 + 0.5;
  //
  // Tangent and bitangent.
  // {
  //   vec3 pos_dx = dFdx(Position);
  //   vec3 pos_dy = dFdy(Position);
  //   // vec3 uv_dx = vec3(dFdx(Texcoord), 0.0);
  //   // vec3 uv_dy = vec3(dFdy(Texcoord), 0.0);
  //   vec3 uv_dx = dFdx(vec3(Texcoord, 0.0));
  //   vec3 uv_dy = dFdy(vec3(Texcoord, 0.0));
  //   vec3 T = (uv_dy.t * pos_dx - uv_dx.t * pos_dy) /
  //            (uv_dx.s * uv_dy.t - uv_dy.s * uv_dx.t);
  //   // vec3 T = pos_dx * uv_dy.t - pos_dy * uv_dx.t;
  //   vec3 N = normalize(Normal);
  //   T = normalize(T - dot(T, N) * N);
  //   vec3 B = normalize(cross(N, T));

  //   if (gl_FrontFacing == false) {
  //     T = -T;
  //     B = -B;
  //     N = -N;
  //   }

  //   // Tangent.
  //   //colour = T * 0.5 + 0.5;
  //   // Bitangent.
  //   //colour = B * 0.5 + 0.5;
  // }

  Colour = vec4(colour, 1.0);
}