#ifndef __BRDF_INC
#define __BRDF_INC

#include "pema99.cginc"
#include "pbr.cginc"
#include "lighting.cginc"
#include "lysenko.cginc"
#include "math.cginc"

float pow5(float x) {
  float x2 = x * x;
  return x2 * x2 * x;
}

float Fd_Lambertian(float NoL) {
  return NoL;
}

// Schlick "An Inexpensive BRDF Model for Physically-based Rendering".
// Equation 24.
// f0: Reflectance at normal incidence. Typically around 0.04.
// f90: Reflectance at grazing incidence. Typically around 1.0.
float3 F_Schlick(float LoH, float3 f0, float f90) {
  float term5 = pow5(1.0f - LoH);
  float3 f90v = float3(f90, f90, f90);
  return f0 + (f90v - f0) * term5;
}

// Walter "Microfacet Models for Refraction through Rough Surfaces"
// Equation 33.
// In the paper:
//  - m = microsurface normal
//  - n = macrosurface normal
//  - theta_m = angle between micro- & macrosurface normals
//  - alpha = roughness
//  - cos(theta_m) = NoH
// Per sohcahtoa:
//  tan(theta) = sin(theta) / cos(theta)
//  tan^2(theta) = sin^2(theta) / cos^2(theta)
//               = (1 - cos^2(theta)) / cos^2(theta)
//               = -1 + 1 / cos^2(theta)
float D_GGX(float roughness, float NoH) {
  float r2 = roughness * roughness;
  float NoH2 = NoH * NoH;
  float NoH4 = NoH2 * NoH2;

  float k = rcp(NoH2) - 1;
  float r2_plus_k = r2 + k;
  float denom = NoH4 * r2_plus_k * r2_plus_k;

  //return min(4096, r2 / denom);
  return r2 / denom;
}

float D_Estevez(float roughness, float NoH) {
  float r_rcp = rcp(roughness);
  float sin_theta = sqrt(1 - NoH * NoH);
  float D = (2 + r_rcp) * pow(sin_theta, r_rcp) / TAU;

  return D;
}

// Hammon "PBR Diffuse Lighting for GGX+Smith Microsurfaces"
// Slide 84. Note that we remove the (4 * NoL * NoV) from the
// denominator of the specular lobe because of some cancellations.
// The original, un-optimized equation is:
//  2 * NoL * NoV / lerp(2 * NoL * NoV, NoL + NoV, roughness)
float G_GGXSmith(float roughness, float NoL, float NoV) {
  float denom = 2.0f * lerp(2.0f * NoL * NoV, NoL + NoV, roughness);
  return rcp(denom);
}

float L_Estevez(float r, float x) {
  // Recover constants according to Table 1.
  float interpolator = 1 - r * r;
  float one_minus_i = 1 - interpolator;
  float a = interpolator * 25.3245 + one_minus_i * 21.5473;
  float b = interpolator * 3.32435 + one_minus_i * 3.82987;
  float c = interpolator * 0.16801 + one_minus_i * 0.19823;
  float d = interpolator * -1.27393 + one_minus_i * -1.97760;
  float e = interpolator * -4.85967 + one_minus_i * -4.32054;

  return a / (1 + b*pow(x, c)) + d*x + e;
}

float Lambda_Estevez(float cos_theta, float roughness) {
  // Equation 3
  float lambda = cos_theta < 0.5
    ? exp(L_Estevez(roughness, cos_theta))
    : exp(2 * L_Estevez(roughness, 0.5) - L_Estevez(roughness, 1 - cos_theta));
  // Equation 4
  return pow(lambda, 1 + 2 * pow(1 - cos_theta, 8));
}

// Estevez & Kulla "Production Friendly Microfacet Sheen BRDF"
// Height-correlated Smith: G2 / (4 * NoL * NoV)
float G_Estevez(float roughness, float NoL, float NoV) {
  float lambda_l = Lambda_Estevez(NoL, roughness);
  float lambda_v = Lambda_Estevez(NoV, roughness);
  return 1.0 / ((1.0 + lambda_l + lambda_v) * 4.0 * NoL * NoV);
}

float4 brdf(Pbr pbr, LightData data, out BrdfData bd) {
  bd = (BrdfData)0;
  float3 specular = 0;
  float3 diffuse = 0;

//#define FURNACE_TEST_DIRECT
#if defined(FURNACE_TEST_DIRECT)
  // Create the conditions for the standard BRDF furnace test.
  // Only applies to the direct lighting stage. The only variable left over is
  // NoV.
  f0 = 1;
  data.direct.color = 1;
  data.direct.NoL = 1;
  data.direct.NoH = 1;
  data.direct.LoH = 1;
#endif

  // TODO parameterize
  float f0 = 0.04f;
  const float f90 = 1.0f;
  float2 dfg_uv = float2(data.common.NoV, pbr.roughness_perceptual);
  [branch]
  if (textureExists(_DFG_LUT)) {
    bd.ibl_dfg = _DFG_LUT.SampleLevel(bilinear_clamp_s, dfg_uv, 0).rgb;
  } else {
    bd.ibl_dfg = float3(1, 1, 1);
  }
  float3 f0_color = lerp(f0, pbr.albedo.xyz, pbr.metallic);
  float3 energy_comp = 1.0f + f0_color * (1.0f / (bd.ibl_dfg.xxx + bd.ibl_dfg.yyy) - 1.0f);

#if defined(_CLEARCOAT)
  const float cc_f0 = 0.04f;
  float2 cc_dfg_uv = float2(data.common.NoV_cc, pbr.cc_roughness_perceptual);
  [branch]
  if (textureExists(_DFG_LUT)) {
    bd.ibl_dfg_cc = _DFG_LUT.SampleLevel(bilinear_clamp_s, cc_dfg_uv, 0).rgb;
  } else {
    bd.ibl_dfg_cc = float3(1, 1, 1);
  }
  float3 cc_f0_color = lerp(cc_f0, pbr.albedo.xyz, pbr.metallic);
  float3 cc_energy_comp = 1.0f + cc_f0_color * (1.0f / (bd.ibl_dfg_cc.xxx + bd.ibl_dfg_cc.yyy) - 1.0f);
#endif

  // Direct
  {
    float3 remainder = 1.0f;

#if defined(_CLEARCOAT)
    bd.direct_f_cc = F_Schlick(data.direct.LoH, cc_f0, f90);
    bd.direct_d_cc = D_GGX(pbr.cc_roughness, data.direct.NoH_cc);
    bd.direct_g_cc = G_GGXSmith(pbr.cc_roughness, data.direct.NoL_cc, data.common.NoV_cc);
    float DFGcc = bd.direct_f_cc * bd.direct_d_cc * bd.direct_g_cc;
    float3 direct_specular_cc = DFGcc * data.direct.color * data.direct.NoL_cc * pbr.cc_strength;
    direct_specular_cc *= cc_energy_comp;
    direct_specular_cc *= remainder;
    direct_specular_cc = max(0, direct_specular_cc);
    specular += direct_specular_cc;
    remainder *= saturate(1.0f - bd.direct_f_cc * pbr.cc_strength);
#endif

#if defined(_CLOTH)
    float3 cloth_f0 = _Cloth_Sheen.rgb;
    bd.direct_f = F_Schlick(data.direct.LoH, cloth_f0, f90);
    bd.direct_d = D_Estevez(pbr.roughness, data.direct.NoH);
    bd.direct_g = G_Estevez(pbr.roughness, data.direct.NoL, data.common.NoV);

    float3 direct_specular = (bd.direct_d * bd.direct_g) * bd.direct_f;
    direct_specular *= data.direct.color * data.direct.NoL;
    direct_specular *= remainder;
    specular += direct_specular;

    float Fd = Fd_Lambertian(data.direct.NoL) / PI;
    float3 direct_diffuse = Fd * pbr.albedo.xyz * data.direct.color;
    direct_diffuse *= remainder;
    direct_diffuse = max(0, direct_diffuse);
    diffuse += direct_diffuse;
#else
    bd.direct_f = F_Schlick(data.direct.LoH, f0_color, f90);
    bd.direct_d = D_GGX(pbr.roughness, data.direct.NoH);
    bd.direct_g = G_GGXSmith(pbr.roughness, data.direct.NoL, data.common.NoV);

    float3 direct_specular = (bd.direct_d * bd.direct_g) * bd.direct_f;
    direct_specular *= data.direct.color * data.direct.NoL;
    direct_specular *= energy_comp;
    direct_specular *= remainder;
    specular += direct_specular;

#if defined(F_OREN_NAYAR)
    float Fd = Fd_OrenNayar(pbr.roughness, data.common.NoV, data.direct.NoL, data.direct.LoV);
#else
    float Fd = Fd_Lambertian(data.direct.NoL);
#endif
    float3 direct_diffuse = Fd * (1.0f - pbr.metallic) * pbr.albedo.xyz * data.direct.color;
    direct_diffuse *= remainder;
    direct_diffuse = max(0, direct_diffuse);
    diffuse += direct_diffuse;
#endif
  }

  // Indirect
#if !defined(FURNACE_TEST_DIRECT) && (defined(FORWARD_BASE_PASS) || defined(OUTLINES_PASS))
  {
    float3 remainder = 1.0f;
#if defined(_CLEARCOAT)
    float3 cc_specular_dfg = bd.ibl_dfg_cc.xxx * cc_f0_color + bd.ibl_dfg_cc.yyy;  // filament 5.3.4.6
    float3 cc_indirect_specular = data.indirect.specular_cc * cc_specular_dfg;
    cc_indirect_specular *= cc_energy_comp;
    specular += cc_indirect_specular;
    remainder -= cc_specular_dfg;
#endif

#if defined(_CLOTH)
    float3 specular_dfg = _Cloth_Sheen.rgb * bd.ibl_dfg.zzz;
    float3 indirect_specular = data.indirect.specular * specular_dfg;
    specular += indirect_specular * remainder;

    float3 indirect_diffuse = pbr.albedo.xyz * data.indirect.diffuse;
    diffuse  += indirect_diffuse * remainder;
#else
    float3 specular_dfg = bd.ibl_dfg.xxx * f0_color + bd.ibl_dfg.yyy;  // filament 5.3.4.6
    float3 indirect_specular = data.indirect.specular * specular_dfg;

    indirect_specular *= energy_comp;
    specular += indirect_specular * remainder;

    float3 indirect_diffuse = pbr.albedo.xyz * data.indirect.diffuse * (1.0 - pbr.metallic);
    diffuse  += indirect_diffuse * remainder;
#endif
  }
#endif

#if defined(FORWARD_BASE_PASS)
  {
    [branch]
    if (_UdonLightVolumeEnabled) {
      specular += LightVolumeSpecular(pbr.albedo.xyz, pbr.smoothness,
          pbr.metallic, pbr.normal, data.common.V, data.indirect.L00,
          data.indirect.L01r, data.indirect.L01g, data.indirect.L01b);
    }
  }
#endif

  diffuse  *= data.common.ao;
  specular *= data.common.spec_ao;

#if (defined(_EMISSIONS) || defined(_LETTER_GRID)) && defined(FORWARD_BASE_PASS)
  float3 emission = pbr.emission;
#else
  float3 emission = 0;
#endif
  float4 lit = float4(diffuse + specular + emission, pbr.albedo.a);
  // Scale albedo by alpha.
  return float4(lit.rgb * lit.a, lit.a);
}

#endif  // __BRDF_INC