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

#pragma once

#include "MicrofacetAlbedoTables.ih"
#include "MicrofacetDielectric.ih"

OSPRAY_BEGIN_ISPC_NAMESPACE

// Simplified, energy conserving Weidlich-Wilkie microfacet coating BSDF with
// the Smith microsurface model. Refraction is ignored, but absorption is
// computed from the refracted ray lengths.
// [Weidlich and Wilkie, 2007, "Arbitrarily Layered Micro-Facet Surfaces"]
// [Walter et al., 2007, "Microfacet Models for Refraction through Rough
// Surfaces"] [Kulla and Conty, 2017, "Revisiting Physically Based Shading at
// Imageworks"] [Kelemen and Szirmay-Kalos, 2001, "A Microfacet Based Coupled
// Specular-Matte BRDF Model with Importance Sampling"] [Jakob et al., 2014, "A
// Comprehensive Framework for Rendering Layered Materials", Extended Technical
// Report]
struct MicrofacetDielectricLayer
{
  BSDF super;

  float weight;

  float eta;
  vec3f transmittance;
  float thickness;

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

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

inline void MicrofacetDielectricLayer_Constructor(
    varying MicrofacetDielectricLayer *uniform self,
    MicrofacetAlbedoTables *uniform microfacetAlbedoTables,
    const varying linear3f *uniform frame,
    varying BSDF *varying substrate,
    float eta,
    vec3f transmittance,
    float thickness,
    float roughness,
    float anisotropy,
    float weight)
{
  // Energy conservation/compensation
  self->Eavg = MicrofacetAlbedoTable_evalAvg(microfacetAlbedoTables, roughness);
  float Favg = fresnelDielectricAvg(eta);
  self->fmsScale = Favg * (1.f - self->Eavg) * rcp(1.f - Favg * self->Eavg);
  self->Ravg = (MicrofacetDielectricReflectionAlbedoTable_evalAvg(
                    microfacetAlbedoTables, self->eta, roughness)
                   + self->fmsScale * (1.f - self->Eavg))
      * weight; // add missing energy

  // XXX check
  const vec3f albedo = substrate->albedo
      * ((1.0f - weight) + pow(transmittance, thickness) * (1.f - Favg));
  BSDF_Constructor(&self->super,
      albedo,
      BSDF_GLOSSY_REFLECTION | substrate->scatteringType,
      BSDF_TYPE_MICROFACET_DIELECTRIC_LAYER,
      frame);

  self->eta = (eta <= 1.f) ? eta : rcp(eta);
  self->transmittance = transmittance;
  self->thickness = thickness;
  self->microfacet =
      make_GGXDistribution(roughnessToAlpha(roughness, anisotropy));
  self->roughness = roughness;
  self->microfacetAlbedoTables = microfacetAlbedoTables;

  self->weight = weight;
}

inline varying BSDF *uniform MicrofacetDielectricLayer_create(
    uniform ShadingContext *uniform ctx,
    MicrofacetAlbedoTables *uniform microfacetAlbedoTables,
    const varying linear3f *uniform frame,
    varying BSDF *varying substrate,
    float eta,
    vec3f transmittance,
    float thickness,
    float roughness,
    float anisotropy,
    float weight)
{
  varying MicrofacetDielectricLayer *uniform self =
      (varying MicrofacetDielectricLayer * uniform)
          ShadingContext_alloc(ctx, sizeof(MicrofacetDielectricLayer));
  MicrofacetDielectricLayer_Constructor(self,
      microfacetAlbedoTables,
      frame,
      substrate,
      eta,
      transmittance,
      thickness,
      roughness,
      anisotropy,
      weight);
  return &self->super;
}

// Helper function for transparent shadow rays
inline vec3f MicrofacetDielectricLayer_getTransparency(
    MicrofacetAlbedoTables *uniform microfacetAlbedoTables,
    float cosThetaO,
    float eta,
    vec3f transmittance,
    float thickness,
    float roughness,
    float,
    float weight)
{
  if (eta > 1.f)
    eta = rcp(eta);

  // Apply the coating medium absorption
  // Use refracted angles for computing the absorption path length
  float lengthO1 = rcp(refract(cosThetaO, eta)); // rcp(cosThetaO1)
  vec3f value =
      lerp(weight, make_vec3f(1.f), pow(transmittance, thickness * lengthO1));

  // Energy conservation
  float Eavg = MicrofacetAlbedoTable_evalAvg(microfacetAlbedoTables, roughness);
  float Favg = fresnelDielectricAvg(eta);
  float fmsScale = Favg * (1.f - Eavg) * rcp(1.f - Favg * Eavg);
  float Eo =
      MicrofacetAlbedoTable_eval(microfacetAlbedoTables, cosThetaO, roughness);
  float Ro = (MicrofacetDielectricReflectionAlbedoTable_eval(
                  microfacetAlbedoTables, cosThetaO, eta, roughness)
                 + fmsScale * (1.f - Eo))
      * weight; // add missing energy
  float T =
      1.f - Ro; // for generic (non-diffuse) substrates [Kulla and Conty, 2017]
  value = value * T;
  return value;
}

OSPRAY_END_ISPC_NAMESPACE

// Eval macros

#define MICROFACETDIELECTRICLAYER_EVAL(                                        \
    mdl_name, scatteringType, substrate_bsdf, substrate_fn)                    \
  const varying MicrofacetDielectricLayer *uniform mdl =                       \
      (const varying MicrofacetDielectricLayer *uniform) & mdl_name;           \
                                                                               \
  float cosThetaO = dot(wo, getN(&mdl->super));                                \
  if (cosThetaO <= 0.f)                                                        \
    return make_BSDF_EvalRes_zero();                                           \
  float cosThetaI = dot(wi, getN(&mdl->super));                                \
                                                                               \
  /* Evaluate the substrate */                                                 \
  /* Ignore refraction */                                                      \
  BSDF_EvalRes substrate = substrate_fn(substrate_bsdf, wo, wi);               \
                                                                               \
  /* Apply the coating medium absorption */                                    \
  /* Use refracted angles for computing the absorption path length */          \
  float lengthO1 = rcp(refract(cosThetaO, mdl->eta)); /* rcp(cosThetaO1) */    \
  float lengthI1 =                                                             \
      rcp(refract(abs(cosThetaI), mdl->eta)); /* rcp(cosThetaI1) */            \
  float length = lengthO1 + lengthI1;                                          \
  if (cosThetaI <= 0.f)                                                        \
    length *= 0.5f; /* for transmission, use the average length */             \
  substrate.value = lerp(mdl->weight,                                          \
      substrate.value,                                                         \
      substrate.value * pow(mdl->transmittance, mdl->thickness * length));     \
                                                                               \
  /* Energy conservation */                                                    \
  float Eo = MicrofacetAlbedoTable_eval(                                       \
      mdl->microfacetAlbedoTables, cosThetaO, mdl->roughness);                 \
  float Ei = MicrofacetAlbedoTable_eval(                                       \
      mdl->microfacetAlbedoTables, abs(cosThetaI), mdl->roughness);            \
  float Ro =                                                                   \
      (MicrofacetDielectricReflectionAlbedoTable_eval(                         \
           mdl->microfacetAlbedoTables, cosThetaO, mdl->eta, mdl->roughness)   \
          + mdl->fmsScale * (1.f - Eo))                                        \
      * mdl->weight; /* add missing energy */                                  \
  float Ri = (MicrofacetDielectricReflectionAlbedoTable_eval(                  \
                  mdl->microfacetAlbedoTables,                                 \
                  abs(cosThetaI),                                              \
                  mdl->eta,                                                    \
                  mdl->roughness)                                              \
                 + mdl->fmsScale * (1.f - Ei))                                 \
      * mdl->weight; /* add missing energy */                                  \
  float T;                                                                     \
  if (scatteringType & ~BSDF_DIFFUSE)                                          \
    T = min(1.f - Ro, 1.f - Ri); /* for generic (non-diffuse) substrates       \
                                    [Kulla and Conty, 2017] */                 \
  else                                                                         \
    T = (1.f - Ro) * (1.f - Ri)                                                \
        * rcp(1.f - mdl->Ravg); /* for diffuse substrates [Kelemen and         \
                                    Szirmay-Kalos, 2001] */                    \
  substrate.value = substrate.value * T;                                       \
                                                                               \
  float coatPickProb = Ro;                                                     \
  float substratePickProb = 1.f - coatPickProb;                                \
                                                                               \
  if (cosThetaI > 0.f) {                                                       \
    /* Compute the microfacet normal */                                        \
    vec3f wh = normalize(wo + wi);                                             \
    float cosThetaOH = dot(wo, wh);                                            \
                                                                               \
    /* Fresnel term */                                                         \
    float F = fresnelDielectric(cosThetaOH, mdl->eta) * mdl->weight;           \
                                                                               \
    /* Evaluate the coating reflection */                                      \
    float cosThetaIH = dot(wi, wh);                                            \
                                                                               \
    linear3f toLocal = transposed(getFrame(&mdl->super));                      \
    vec3f wo0 = toLocal * wo;                                                  \
    vec3f wi0 = toLocal * wi;                                                  \
    vec3f wh0 = toLocal * wh;                                                  \
                                                                               \
    float whPdf;                                                               \
    float D = evalVisible(mdl->microfacet, wh0, wo0, cosThetaOH, whPdf);       \
    float G = evalG2(mdl->microfacet, wo0, wi0, cosThetaOH, cosThetaIH);       \
                                                                               \
    /* Energy compensation */                                                  \
    float fms = mdl->fmsScale                                                  \
        * ((1.f - Eo) * (1.f - Ei) * rcp((float)pi * (1.f - mdl->Eavg))        \
            * cosThetaI);                                                      \
                                                                               \
    BSDF_EvalRes coat;                                                         \
    coat.pdf = whPdf * rcp(4.f * cosThetaOH);                                  \
    coat.value = make_vec3f(F * D * G * rcp(4.f * cosThetaO) + fms);           \
                                                                               \
    /* Compute the total reflection */                                         \
    substrate.pdf =                                                            \
        coatPickProb * coat.pdf + substratePickProb * substrate.pdf;           \
    substrate.value = coat.value + substrate.value;                            \
  } else {                                                                     \
    /* Return the substrate transmission */                                    \
    substrate.pdf *= substratePickProb;                                        \
  }

#define MICROFACETDIELECTRICLAYER_EVAL_GET() substrate

// Sample macros

#define MICROFACETDIELECTRICLAYER_SAMPLE(mdl_name,                             \
    scatteringType,                                                            \
    substrate_bsdf,                                                            \
    substrate_eval,                                                            \
    substrate_sample)                                                          \
  const varying MicrofacetDielectricLayer *uniform mdl =                       \
      (const varying MicrofacetDielectricLayer *uniform) & mdl_name;           \
  BSDF_SampleRes res;                                                          \
                                                                               \
  float cosThetaO = dot(wo, getN(&mdl->super));                                \
  if (cosThetaO <= 0.f)                                                        \
    return make_BSDF_SampleRes_zero();                                         \
                                                                               \
  linear3f toGlobal = getFrame(&mdl->super);                                   \
  linear3f toLocal = transposed(getFrame(&mdl->super));                        \
  vec3f wo0 = toLocal * wo;                                                    \
                                                                               \
  /* Energy conservation */                                                    \
  float Eo = MicrofacetAlbedoTable_eval(                                       \
      mdl->microfacetAlbedoTables, cosThetaO, mdl->roughness);                 \
  float Ro =                                                                   \
      MicrofacetDielectricReflectionAlbedoTable_eval(                          \
          mdl->microfacetAlbedoTables, cosThetaO, mdl->eta, mdl->roughness)    \
      + mdl->fmsScale * (1.f - Eo) * mdl->weight; /* add missing energy */     \
                                                                               \
  /* Sample the coating or the substrate */                                    \
  vec3f wh;                                                                    \
  BSDF_EvalRes substrate;                                                      \
                                                                               \
  float coatPickProb = Ro;                                                     \
  float substratePickProb = 1.f - coatPickProb;                                \
                                                                               \
  if (ss < coatPickProb) {                                                     \
    /* Sample the microfacet normal */                                         \
    float whPdf;                                                               \
    wh = toGlobal * sampleVisible(mdl->microfacet, wo0, whPdf, s);             \
    float cosThetaOH = dot(wo, wh);                                            \
                                                                               \
    res.type = BSDF_GLOSSY_REFLECTION;                                         \
    res.wi = reflect(wo, wh, cosThetaOH);                                      \
                                                                               \
    /* Evaluate the substrate */                                               \
    /* Ignore refraction */                                                    \
    substrate = substrate_eval(substrate_bsdf, wo, res.wi);                    \
  } else {                                                                     \
    /* Sample the substrate */                                                 \
    /* Ignore refraction */                                                    \
    float ss1 =                                                                \
        (ss - coatPickProb) * rcp(substratePickProb); /* reallocate sample */  \
    res = substrate_sample(substrate_bsdf, wo, s, ss1);                        \
    if (reduce_max(res.weight) <= 0.f)                                         \
      return res;                                                              \
                                                                               \
    substrate.pdf = res.pdf;                                                   \
    substrate.value = res.weight                                               \
        * ((res.type & BSDF_SPECULAR)                                          \
                ? 1.f                                                          \
                : res.pdf); /* correctly handle delta distributions */         \
                                                                               \
    /* Compute the microfacet normal */                                        \
    wh = normalize(wo + res.wi);                                               \
  }                                                                            \
                                                                               \
  float cosThetaI = dot(res.wi, getN(&mdl->super));                            \
  float cosThetaOH = dot(wo, wh);                                              \
                                                                               \
  /* Apply the coating medium absorption */                                    \
  /* Use refracted angles for computing the absorption path length */          \
  float lengthO1 = rcp(refract(cosThetaO, mdl->eta)); /* rcp(cosThetaO1) */    \
  float lengthI1 =                                                             \
      rcp(refract(abs(cosThetaI), mdl->eta)); /* rcp(cosThetaI1) */            \
  float length = lengthO1 + lengthI1;                                          \
  if (cosThetaI <= 0.f)                                                        \
    length *= 0.5f; /* for transmission, use the average length */             \
  substrate.value = lerp(mdl->weight,                                          \
      substrate.value,                                                         \
      substrate.value * pow(mdl->transmittance, mdl->thickness * length));     \
                                                                               \
  /* Energy conservation */                                                    \
  float Ei = MicrofacetAlbedoTable_eval(                                       \
      mdl->microfacetAlbedoTables, abs(cosThetaI), mdl->roughness);            \
  float Ri = (MicrofacetDielectricReflectionAlbedoTable_eval(                  \
                  mdl->microfacetAlbedoTables,                                 \
                  abs(cosThetaI),                                              \
                  mdl->eta,                                                    \
                  mdl->roughness)                                              \
                 + mdl->fmsScale * (1.f - Ei))                                 \
      * mdl->weight; /* add missing energy */                                  \
  float T;                                                                     \
  if (scatteringType & ~BSDF_DIFFUSE)                                          \
    T = min(1.f - Ro, 1.f - Ri); /* for generic (non-diffuse) substrates       \
                                    [Kulla and Conty, 2017] */                 \
  else                                                                         \
    T = (1.f - Ro) * (1.f - Ri)                                                \
        * rcp(1.f - mdl->Ravg); /* for diffuse substrates [Kelemen and         \
                                    Szirmay-Kalos, 2001] */                    \
  substrate.value = substrate.value * T;                                       \
                                                                               \
  if (res.type & BSDF_SPECULAR) {                                              \
    /* If we sampled a delta distribution, we don't have to evaluate the       \
     * coating reflection (see MultiBSDF_sample) */                            \
    res.weight = substrate.value * rcp(substratePickProb);                     \
  } else if (cosThetaI <= 0.f) {                                               \
    /* If we sampled transmission, we just have to return the substrate        \
     * transmission */                                                         \
    res.pdf = substratePickProb * substrate.pdf;                               \
    res.weight = substrate.value * rcp(res.pdf);                               \
  } else {                                                                     \
    /* Fresnel term */                                                         \
    float F = fresnelDielectric(cosThetaOH, mdl->eta) * mdl->weight;           \
                                                                               \
    /* Evaluate the coating reflection */                                      \
    float cosThetaIH = dot(res.wi, wh);                                        \
    vec3f wi0 = toLocal * res.wi;                                              \
    vec3f wh0 = toLocal * wh;                                                  \
                                                                               \
    float whPdf;                                                               \
    float D = evalVisible(mdl->microfacet, wh0, wo0, cosThetaOH, whPdf);       \
    float G = evalG2(mdl->microfacet, wo0, wi0, cosThetaOH, cosThetaIH);       \
                                                                               \
    /* Energy compensation */                                                  \
    float fms = mdl->fmsScale                                                  \
        * ((1.f - Eo) * (1.f - Ei) * rcp((float)pi * (1.f - mdl->Eavg))        \
            * cosThetaI);                                                      \
                                                                               \
    BSDF_EvalRes coat;                                                         \
    coat.pdf = whPdf * rcp(4.f * cosThetaOH);                                  \
    coat.value = make_vec3f(F * D * G * rcp(4.f * cosThetaO) + fms);           \
                                                                               \
    /* Compute the total reflection */                                         \
    res.pdf = coatPickProb * coat.pdf + substratePickProb * substrate.pdf;     \
    res.weight = (coat.value + substrate.value) * rcp(res.pdf);                \
  }

#define MICROFACETDIELECTRICLAYER_SAMPLE_GET() res
