// Copyright 2009 Intel Corporation
// SPDX-License-Identifier: Apache-2.0

#pragma once

#include "BSDF.ih"
#include "Fresnel.ih"
#include "GGXDistribution.ih"
#include "MicrofacetAlbedoTables.ih"

OSPRAY_BEGIN_ISPC_NAMESPACE

// Microfacet conductor BRDF with the Smith microsurface model and approximate
// multiple scattering. [Kulla and Conty, 2017, "Revisiting Physically Based
// Shading at Imageworks"] [Jakob et al., 2014, "A Comprehensive Framework for
// Rendering Layered Materials", Extended Technical Report]
struct MicrofacetConductor
{
  BSDF super;
  Fresnel *uniform fresnel;

  GGXDistribution microfacet;
  float roughness;
  const MicrofacetAlbedoTables *uniform microfacetAlbedoTables;

  // Energy compensation [Kulla and Conty, 2017]
  float Eavg;
  vec3f fmsScale;
};

inline BSDF_EvalRes MicrofacetConductor_eval(
    const varying MicrofacetConductor *uniform self,
    const vec3f &wo,
    const vec3f &wi)
{
  BSDF_EvalRes res;

  float cosThetaO = dot(wo, getN(&self->super));
  float cosThetaI = dot(wi, getN(&self->super));
  if (cosThetaO <= 0.f || cosThetaI <= 0.f)
    return make_BSDF_EvalRes_zero();

  // Compute the microfacet normal
  vec3f wh = normalize(wi + wo);
  float cosThetaOH = dot(wo, wh);
  float cosThetaIH = dot(wi, wh);

  linear3f toLocal = transposed(getFrame(&self->super));
  vec3f wo0 = toLocal * wo;
  vec3f wi0 = toLocal * wi;
  vec3f wh0 = toLocal * wh;

  vec3f F = Fresnel_dispatch_eval(self->fresnel, cosThetaOH);
  float whPdf;
  float D = evalVisible(self->microfacet, wh0, wo0, cosThetaOH, whPdf);
  float G = evalG2(self->microfacet, wo0, wi0, cosThetaOH, cosThetaIH);

  // Energy compensation
  float Eo = MicrofacetAlbedoTable_eval(
      self->microfacetAlbedoTables, cosThetaO, self->roughness);
  float Ei = MicrofacetAlbedoTable_eval(
      self->microfacetAlbedoTables, cosThetaI, self->roughness);
  vec3f fms = self->fmsScale
      * ((1.f - Eo) * (1.f - Ei) * rcp((float)pi * (1.f - self->Eavg))
          * cosThetaI);

  res.pdf = whPdf * rcp(4.f * abs(cosThetaOH));
  res.value = F * (D * G * rcp(4.f * cosThetaO)) + fms;
  return res;
}

inline BSDF_SampleRes MicrofacetConductor_sample(
    const varying MicrofacetConductor *uniform self,
    const vec3f &wo,
    const vec2f &s,
    float)
{
  BSDF_SampleRes res;

  float cosThetaO = dot(wo, getN(&self->super));
  if (cosThetaO <= 0.f)
    return make_BSDF_SampleRes_zero();

  linear3f toGlobal = getFrame(&self->super);
  linear3f toLocal = transposed(getFrame(&self->super));
  vec3f wo0 = toLocal * wo;

  // Sample the microfacet normal
  float whPdf;
  vec3f wh = toGlobal * sampleVisible(self->microfacet, wo0, whPdf, s);

  res.wi = reflect(wo, wh);
  float cosThetaI = dot(res.wi, getN(&self->super));
  if (cosThetaI <= 0.f)
    return make_BSDF_SampleRes_zero();
  float cosThetaOH = dot(wo, wh);
  float cosThetaIH = dot(res.wi, wh);
  vec3f wi0 = toLocal * res.wi;

  vec3f F = Fresnel_dispatch_eval(self->fresnel, cosThetaOH);
  float G = evalG2(self->microfacet, wo0, wi0, cosThetaOH, cosThetaIH);

  // Energy compensation
  float Eo = MicrofacetAlbedoTable_eval(
      self->microfacetAlbedoTables, cosThetaO, self->roughness);
  float Ei = MicrofacetAlbedoTable_eval(
      self->microfacetAlbedoTables, cosThetaI, self->roughness);
  vec3f fms = self->fmsScale
      * ((1.f - Eo) * (1.f - Ei) * rcp((float)pi * (1.f - self->Eavg))
          * cosThetaI);

  res.type = BSDF_GLOSSY_REFLECTION;
  res.pdf = whPdf * rcp(4.f * abs(cosThetaOH));
  res.weight = F * (G * rcp_safe(evalG1(self->microfacet, wo0, cosThetaOH)))
      + (fms * rcp(res.pdf));

  return res;
}

inline void MicrofacetConductor_Constructor(
    varying MicrofacetConductor *uniform self,
    MicrofacetAlbedoTables *uniform microfacetAlbedoTables,
    const varying linear3f *uniform frame,
    Fresnel *uniform fresnel,
    float roughness,
    float anisotropy)
{
  // Energy compensation
  self->Eavg = MicrofacetAlbedoTable_evalAvg(microfacetAlbedoTables, roughness);
  vec3f Favg = Fresnel_dispatch_evalAvg(fresnel);
  self->fmsScale = sqr(Favg) * self->Eavg
      / (1.f - Favg * (1.f - self->Eavg)); // Stephen Hill's tweak

  BSDF_Constructor(&self->super,
      Favg * self->Eavg, // TODO better estimate
      BSDF_GLOSSY_REFLECTION,
      BSDF_TYPE_MICROFACET_CONDUCTOR,
      frame);
  self->fresnel = fresnel;
  self->microfacet =
      make_GGXDistribution(roughnessToAlpha(roughness, anisotropy));
  self->roughness = roughness;
  self->microfacetAlbedoTables = microfacetAlbedoTables;
}

inline varying BSDF *uniform MicrofacetConductor_create(
    uniform ShadingContext *uniform ctx,
    MicrofacetAlbedoTables *uniform microfacetAlbedoTables,
    const varying linear3f *uniform frame,
    Fresnel *uniform fresnel,
    float roughness,
    float anisotropy)
{
  varying MicrofacetConductor *uniform self =
      (varying MicrofacetConductor * uniform)
          ShadingContext_alloc(ctx, sizeof(MicrofacetConductor));
  MicrofacetConductor_Constructor(
      self, microfacetAlbedoTables, frame, fresnel, roughness, anisotropy);
  return &self->super;
}

OSPRAY_END_ISPC_NAMESPACE
