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

#pragma once

#include "Dielectric.ih"

OSPRAY_BEGIN_ISPC_NAMESPACE

// Simplified, energy conserving Weidlich-Wilkie smooth coating BSDF.
// Refraction is ignored, but absorption is computed from the refracted ray
// lengths. [Weidlich and Wilkie, 2007, "Arbitrarily Layered Micro-Facet
// 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"]
struct DielectricLayer
{
  BSDF super;

  float weight;

  float eta;
  vec3f transmittance;
  float thickness;

  // Energy conservation [Kulla and Conty, 2017]
  float Favg;
};

inline void DielectricLayer_Constructor(varying DielectricLayer *uniform self,
    const varying linear3f *uniform frame,
    varying BSDF *varying substrate,
    float eta,
    vec3f transmittance,
    float thickness,
    float weight)
{
  self->Favg = fresnelDielectricAvg(eta) * weight;
  // XXX check
  const vec3f albedo = substrate->albedo
      * ((1.0f - weight) + pow(transmittance, thickness) * (1.f - self->Favg));
  BSDF_Constructor(&self->super,
      albedo,
      BSDF_SPECULAR_REFLECTION | substrate->scatteringType,
      BSDF_TYPE_DIELECTRIC_LAYER,
      frame);

  self->eta = (eta <= 1.f) ? eta : rcp(eta);
  self->transmittance = transmittance;
  self->thickness = thickness;
  self->weight = weight;
}

// Helper function for transparent shadow rays
inline vec3f DielectricLayer_getTransparency(float cosThetaO,
    float eta,
    vec3f transmittance,
    float thickness,
    float weight)
{
  if (eta < 1.f)
    eta = rcp(eta);

  // Fresnel term
  float cosThetaO1; // positive
  float Fo = fresnelDielectricEx(cosThetaO, cosThetaO1, eta) * weight;

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

  // Energy conservation
  float T =
      1.f - Fo; // for generic (non-diffuse) substrates [Kulla and Conty, 2017]
  value = value * T;
  return value;
}

OSPRAY_END_ISPC_NAMESPACE

// Eval macros

#define DIELECTRICLAYER_EVAL(                                                  \
    dl_name, scatteringType, substrate_bsdf, substrate_fn)                     \
  const varying DielectricLayer *uniform dl =                                  \
      (const varying DielectricLayer *uniform) & dl_name;                      \
                                                                               \
  float cosThetaO = dot(wo, getN(&dl->super));                                 \
  if (cosThetaO <= 0.f)                                                        \
    return make_BSDF_EvalRes_zero();                                           \
                                                                               \
  /* Evaluate the substrate */                                                 \
  /* Ignore refraction */                                                      \
  BSDF_EvalRes substrate = substrate_fn(substrate_bsdf, wo, wi);               \
                                                                               \
  /* Fresnel term */                                                           \
  float cosThetaO1; /* positive */                                             \
  float Fo = fresnelDielectricEx(cosThetaO, cosThetaO1, dl->eta) * dl->weight; \
                                                                               \
  float cosThetaI = dot(wi, getN(&dl->super));                                 \
  float cosThetaI1; /* positive */                                             \
  float Fi =                                                                   \
      fresnelDielectricEx(abs(cosThetaI), cosThetaI1, dl->eta) * dl->weight;   \
                                                                               \
  /* Apply the coating medium absorption */                                    \
  /* Use refracted angles for computing the absorption path length */          \
  float lengthO1 = rcp(cosThetaO1);                                            \
  float lengthI1 = rcp(cosThetaI1);                                            \
  float length = lengthO1 + lengthI1;                                          \
  if (cosThetaI <= 0.f)                                                        \
    length *= 0.5f; /* for transmission, use the average length */             \
  substrate.value = lerp(dl->weight,                                           \
      substrate.value,                                                         \
      substrate.value * pow(dl->transmittance, dl->thickness * length));       \
                                                                               \
  /* Energy conservation */                                                    \
  float T;                                                                     \
  if (scatteringType & ~BSDF_DIFFUSE)                                          \
    T = min(1.f - Fo, 1.f - Fi); /* for generic (non-diffuse) substrates       \
                                    [Kulla and Conty, 2017] */                 \
  else                                                                         \
    T = (1.f - Fo) * (1.f - Fi)                                                \
        * rcp(1.f - dl->Favg); /* for diffuse substrates [Kelemen and          \
                                    Szirmay-Kalos, 2001] */                    \
  substrate.value = substrate.value * T;                                       \
  substrate.pdf *= (1.f - Fo)

#define DIELECTRICLAYER_EVAL_GET() substrate

// Sample macros

#define DIELECTRICLAYER_SAMPLE(                                                \
    dl_name, scatteringType, substrate_bsdf, substrate_fn)                     \
  const varying DielectricLayer *uniform dl =                                  \
      (const varying DielectricLayer *uniform) & dl_name;                      \
  BSDF_SampleRes res;                                                          \
                                                                               \
  float cosThetaO = dot(wo, getN(&dl->super));                                 \
  if (cosThetaO <= 0.f)                                                        \
    return make_BSDF_SampleRes_zero();                                         \
                                                                               \
  /* Fresnel term */                                                           \
  float cosThetaO1; /* positive */                                             \
  float Fo = fresnelDielectricEx(cosThetaO, cosThetaO1, dl->eta) * dl->weight; \
                                                                               \
  if (ss < Fo) {                                                               \
    /* Sample the coating */                                                   \
    res.type = BSDF_SPECULAR_REFLECTION;                                       \
    res.wi = reflect(wo, getN(&dl->super), cosThetaO);                         \
    res.pdf = inf;                                                             \
    res.weight = make_vec3f(1.f);                                              \
  } else {                                                                     \
    /* Sample the substrate */                                                 \
    /* Ignore refraction */                                                    \
    float ss1 = (ss - Fo) * rcp(1.f - Fo); /* reallocate sample */             \
    res = substrate_fn(substrate_bsdf, wo, s, ss1);                            \
    if (reduce_max(res.weight) <= 0.f)                                         \
      return res;                                                              \
    float cosThetaI = dot(res.wi, getN(&dl->super));                           \
                                                                               \
    float cosThetaI1; /* positive */                                           \
    float Fi =                                                                 \
        fresnelDielectricEx(abs(cosThetaI), cosThetaI1, dl->eta) * dl->weight; \
                                                                               \
    /* Apply the coating medium absorption */                                  \
    /* Use refracted angles for computing the absorption path length */        \
    float lengthO1 = rcp(cosThetaO1);                                          \
    float lengthI1 = rcp(cosThetaI1);                                          \
    float length = lengthO1 + lengthI1;                                        \
    if (cosThetaI <= 0.f)                                                      \
      length *= 0.5f; /* for transmission, use the average length */           \
    res.weight = lerp(dl->weight,                                              \
        res.weight,                                                            \
        res.weight * pow(dl->transmittance, dl->thickness * length));          \
                                                                               \
    /* Energy conservation */                                                  \
    float T;                                                                   \
    if (scatteringType & ~BSDF_DIFFUSE)                                        \
      T = min(1.f - Fo, 1.f - Fi); /* for generic (non-diffuse) substrates     \
                                      [Kulla and Conty, 2017] */               \
    else                                                                       \
      T = (1.f - Fo) * (1.f - Fi)                                              \
          * rcp(1.f - dl->Favg); /* for diffuse substrates [Kelemen and        \
                                      Szirmay-Kalos, 2001] */                  \
    res.weight = res.weight * (T * rcp(1.f - Fo));                             \
    res.pdf *= (1.f - Fo);                                                     \
  }

#define DIELECTRICLAYER_SAMPLE_GET() res
