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

#pragma once

#include "BSDF.ih"
#include "MicrofacetAlbedoTables.ih"
#include "SheenDistribution.ih"

OSPRAY_BEGIN_ISPC_NAMESPACE

// Microfacet sheen BSDF for cloth-like materials.
// [Conty and Kulla, 2017, "Production Friendly Microfacet Sheen BRDF"]
// [Kulla and Conty, 2017, "Revisiting Physically Based Shading at Imageworks"]
struct MicrofacetSheenLayer
{
  BSDF super;
  float weight;

  vec3f R;
  SheenDistribution microfacet;
  float roughness;
  const MicrofacetAlbedoTables *uniform microfacetAlbedoTables;
};

inline void MicrofacetSheenLayer_Constructor(
    varying MicrofacetSheenLayer *uniform self,
    const MicrofacetAlbedoTables *uniform microfacetAlbedoTables,
    const varying linear3f *uniform frame,
    varying BSDF *varying substrate,
    vec3f R,
    float roughness,
    float weight)
{
  BSDF_Constructor(&self->super,
      lerp(weight * 0.1f,
          substrate->albedo,
          R), // lean towards substrate (sheen is only at edge)
      BSDF_DIFFUSE_REFLECTION | substrate->scatteringType,
      BSDF_TYPE_MICROFACET_SHEEN_LAYER,
      frame);

  self->R = R;
  self->microfacet = make_SheenDistribution(roughnessToAlpha(roughness));
  self->roughness = roughness;
  self->weight = weight;
  self->microfacetAlbedoTables = microfacetAlbedoTables;
}

inline varying BSDF *uniform MicrofacetSheenLayer_create(
    uniform ShadingContext *uniform ctx,
    const MicrofacetAlbedoTables *uniform microfacetAlbedoTables,
    const varying linear3f *uniform frame,
    varying BSDF *varying substrate,
    vec3f R,
    float roughness,
    float weight)
{
  varying MicrofacetSheenLayer *uniform self =
      (varying MicrofacetSheenLayer * uniform)
          ShadingContext_alloc(ctx, sizeof(MicrofacetSheenLayer));
  MicrofacetSheenLayer_Constructor(
      self, microfacetAlbedoTables, frame, substrate, R, roughness, weight);
  return &self->super;
}

// Helper function for transparent shadow rays
inline vec3f MicrofacetSheenLayer_getTransparency(
    const MicrofacetAlbedoTables *uniform microfacetAlbedoTables,
    float cosThetaO,
    float roughness,
    float weight)
{
  // Energy conservation
  float Ro = MicrofacetSheenAlbedoTable_eval(
                 microfacetAlbedoTables, cosThetaO, roughness)
      * weight;
  float T = 1.f - Ro;
  return make_vec3f(T);
}

OSPRAY_END_ISPC_NAMESPACE

// Eval macros

#define MICROFACETSHEENLAYER_EVAL(msl_name, substrate_bsdf, substrate_fn)      \
  const varying MicrofacetSheenLayer *uniform msl =                            \
      (const varying MicrofacetSheenLayer *uniform) & msl_name;                \
                                                                               \
  float cosThetaO = dot(wo, getN(&msl->super));                                \
  if (cosThetaO <= 0.f)                                                        \
    return make_BSDF_EvalRes_zero();                                           \
  float cosThetaI = dot(wi, getN(&msl->super));                                \
                                                                               \
  /* Evaluate the substrate */                                                 \
  BSDF_EvalRes substrate = substrate_fn(substrate_bsdf, wo, wi);               \
                                                                               \
  /* Energy conservation */                                                    \
  float Ro = MicrofacetSheenAlbedoTable_eval(                                  \
                 msl->microfacetAlbedoTables, cosThetaO, msl->roughness)       \
      * msl->weight;                                                           \
  float Ri = MicrofacetSheenAlbedoTable_eval(                                  \
                 msl->microfacetAlbedoTables, abs(cosThetaI), msl->roughness)  \
      * msl->weight;                                                           \
  float T = min(1.f - Ro, 1.f - Ri);                                           \
  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 cosThetaH = dot(wh, getN(&msl->super));                              \
    float cosThetaOH = dot(wo, wh);                                            \
    float cosThetaIH = dot(wi, wh);                                            \
                                                                               \
    /* Evaluate the coating reflection */                                      \
    float D = eval(msl->microfacet, cosThetaH);                                \
    float G =                                                                  \
        evalG2(msl->microfacet, cosThetaO, cosThetaI, cosThetaOH, cosThetaIH); \
                                                                               \
    BSDF_EvalRes coat;                                                         \
    coat.pdf = uniformSampleHemispherePDF();                                   \
    coat.value = msl->R * (D * G * rcp(4.f * cosThetaO) * msl->weight);        \
                                                                               \
    /* 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 MICROFACETSHEENLAYER_EVAL_GET() substrate

// Sample macros

#define MICROFACETSHEENLAYER_SAMPLE(                                           \
    msl_name, substrate_bsdf, substrate_eval, substrate_sample)                \
  const varying MicrofacetSheenLayer *uniform msl =                            \
      (const varying MicrofacetSheenLayer *uniform) & msl_name;                \
  BSDF_SampleRes res;                                                          \
                                                                               \
  float cosThetaO = dot(wo, getN(&msl->super));                                \
  if (cosThetaO <= 0.f)                                                        \
    return make_BSDF_SampleRes_zero();                                         \
                                                                               \
  /* Energy conservation */                                                    \
  float Ro = MicrofacetSheenAlbedoTable_eval(                                  \
                 msl->microfacetAlbedoTables, cosThetaO, msl->roughness)       \
      * msl->weight;                                                           \
                                                                               \
  /* Sample the coating or the substrate */                                    \
  BSDF_EvalRes substrate;                                                      \
                                                                               \
  float coatPickProb = Ro;                                                     \
  float substratePickProb = 1.f - coatPickProb;                                \
                                                                               \
  if (ss < coatPickProb) {                                                     \
    /* Sample the coating reflection */                                        \
    res.type = BSDF_DIFFUSE_REFLECTION;                                        \
    res.wi = getFrame(&msl->super) * uniformSampleHemisphere(s);               \
                                                                               \
    /* Evaluate the substrate */                                               \
    substrate = substrate_eval(substrate_bsdf, wo, res.wi);                    \
  } else {                                                                     \
    /* Sample the substrate */                                                 \
    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 */         \
  }                                                                            \
                                                                               \
  float cosThetaI = dot(res.wi, getN(&msl->super));                            \
                                                                               \
  /* Energy conservation */                                                    \
  float Ri = MicrofacetSheenAlbedoTable_eval(                                  \
                 msl->microfacetAlbedoTables, abs(cosThetaI), msl->roughness)  \
      * msl->weight;                                                           \
  float T = min(1.f - Ro, 1.f - Ri);                                           \
  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 {                                                                     \
    /* Compute the microfacet normal */                                        \
    vec3f wh = normalize(wo + res.wi);                                         \
    float cosThetaH = dot(wh, getN(&msl->super));                              \
    float cosThetaOH = dot(wo, wh);                                            \
    float cosThetaIH = dot(res.wi, wh);                                        \
                                                                               \
    /* Evaluate the coating reflection */                                      \
    float D = eval(msl->microfacet, cosThetaH);                                \
    float G =                                                                  \
        evalG2(msl->microfacet, cosThetaO, cosThetaI, cosThetaOH, cosThetaIH); \
                                                                               \
    BSDF_EvalRes coat;                                                         \
    coat.pdf = uniformSampleHemispherePDF();                                   \
    coat.value = msl->R * (D * G * rcp(4.f * cosThetaO) * msl->weight);        \
                                                                               \
    /* Compute the total reflection */                                         \
    res.pdf = coatPickProb * coat.pdf + substratePickProb * substrate.pdf;     \
    res.weight = (coat.value + substrate.value) * rcp(res.pdf);                \
  }

#define MICROFACETSHEENLAYER_SAMPLE_GET() res
