// 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 dielectric BSDF with the Smith microsurface model and approximate
// multiple scattering. [Walter et al., 2007, "Microfacet Models for Refraction
// through Rough Surfaces"] [Kulla and Conty, 2017, "Revisiting Physically Based
// Shading at Imageworks"] [Jakob et al., 2014, "A Comprehensive Framework for
// Rendering Layered Materials", Extended Technical Report]
// FIXME: improve robustness
struct MicrofacetDielectric
{
  BSDF super;

  float eta; // etaO / etaI
  GGXDistribution microfacet;
  float roughness;
  const MicrofacetAlbedoTables *uniform microfacetAlbedoTables;

  // Energy compensation [Kulla and Conty, 2017]
  float EavgEta, EavgRcpEta;
  float FavgEta, FavgRcpEta;
  float fmsRatio;
};

// Single-scattering lobe
inline BSDF_EvalRes MicrofacetDielectric_evalSingle(
    const varying BSDF *uniform super, const vec3f &wo, const vec3f &wi)
{
  const varying MicrofacetDielectric *uniform self =
      (const varying MicrofacetDielectric *uniform)super;
  BSDF_EvalRes res;

  float cosThetaO = dot(wo, getN(super));
  if (cosThetaO <= 0.f)
    return make_BSDF_EvalRes_zero();
  float cosThetaI = dot(wi, getN(super));
  bool isReflection = cosThetaI > 0.f;

  // Compute the microfacet normal
  vec3f wh;

  if (isReflection) {
    // Reflection
    wh = wo + wi;
  } else {
    // Transmission
    wh = self->eta * wo + wi;
  }

  wh = normalize(wh);
  float cosThetaH = dot(wh, getN(super));
  // Ensure that the micronormal is in the same hemisphere as the macronormal
  if (cosThetaH < 0.f)
    wh = neg(wh);

  float cosThetaOH = dot(wo, wh);
  float cosThetaIH = dot(wi, wh);

  // Fresnel term
  float F = fresnelDielectric(cosThetaOH, self->eta);

  float value;
  if (isReflection) {
    // Reflection
    res.pdf = F * rcp(4.f * abs(cosThetaOH));
    value = F * rcp(4.f * cosThetaO);
  } else {
    // Transmission
    res.pdf = (1.f - F) * rcp(sqr(self->eta)) * abs(cosThetaIH)
        * rcp(sqr(cosThetaOH + rcp(self->eta) * cosThetaIH));
    value = res.pdf * abs(cosThetaOH)
        * rcp(cosThetaO); // ignore solid angle compression
  }

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

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

  res.pdf *= whPdf;
  res.value = make_vec3f(value * D * G);
  return res;
}

inline BSDF_EvalRes MicrofacetDielectric_eval(
    const varying BSDF *uniform super, const vec3f &wo, const vec3f &wi)
{
  const varying MicrofacetDielectric *uniform self =
      (const varying MicrofacetDielectric *uniform)super;
  BSDF_EvalRes res;

  float cosThetaO = dot(wo, getN(super));
  if (cosThetaO <= 0.f)
    return make_BSDF_EvalRes_zero();
  float cosThetaI = dot(wi, getN(super));
  bool isReflection = cosThetaI > 0.f;

  // Evaluate the energy compensation lobe [Kulla and Conty, 2017]
  float Eo = MicrofacetDielectricAlbedoTable_eval(
      self->microfacetAlbedoTables, cosThetaO, self->eta, self->roughness);

  BSDF_EvalRes fms;
  if (isReflection) {
    // Reflection
    float Ei = MicrofacetDielectricAlbedoTable_eval(
        self->microfacetAlbedoTables, cosThetaI, self->eta, self->roughness);
    fms.pdf = self->fmsRatio * cosineSampleHemispherePDF(cosThetaI);
    fms.value = make_vec3f(self->fmsRatio * (1.f - Eo) * (1.f - Ei)
        * rcp((float)pi * (1.f - self->EavgEta)) * cosThetaI);
  } else {
    // Transmission
    float Ei =
        MicrofacetDielectricAlbedoTable_eval(self->microfacetAlbedoTables,
            abs(cosThetaI),
            rcp(self->eta),
            self->roughness);
    fms.pdf =
        (1.f - self->fmsRatio) * cosineSampleHemispherePDF(abs(cosThetaI));
    fms.value = make_vec3f((1.f - self->fmsRatio) * (1.f - Eo) * (1.f - Ei)
        * rcp((float)pi * (1.f - self->EavgRcpEta)) * abs(cosThetaI));
  }

  // Evaluate the single-scattering lobe
  BSDF_EvalRes single = MicrofacetDielectric_evalSingle(super, wo, wi);

  // Compute the final result
  float singlePickProb = Eo;
  res.pdf = singlePickProb * single.pdf + (1.f - singlePickProb) * fms.pdf;
  res.value = single.value + fms.value;
  return res;
}

inline BSDF_SampleRes MicrofacetDielectric_sample(
    const varying BSDF *uniform super,
    const vec3f &wo,
    const vec2f &s,
    float ss)
{
  const varying MicrofacetDielectric *uniform self =
      (const varying MicrofacetDielectric *uniform)super;
  BSDF_SampleRes res;

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

  // Energy compensation
  float Eo = MicrofacetDielectricAlbedoTable_eval(
      self->microfacetAlbedoTables, cosThetaO, self->eta, self->roughness);

  float singlePickProb = Eo;
  bool doReflection;
  bool doSingle = (ss <= singlePickProb);

  if (doSingle) {
    // Sample the single-scattering lobe
    linear3f toGlobal = getFrame(super);
    linear3f toLocal = transposed(getFrame(super));
    vec3f wo0 = toLocal * wo;

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

    // Fresnel term
    float cosThetaTH; // positive
    float F = fresnelDielectricEx(cosThetaOH, cosThetaTH, self->eta);

    // Sample the reflection or the transmission
    float ss1 = ss * rcp(singlePickProb);
    doReflection = (ss1 <= F);

    if (doReflection) {
      // Reflection
      res.type = BSDF_GLOSSY_REFLECTION;
      res.wi = reflect(wo, wh, cosThetaOH);
    } else {
      // Transmission
      res.type = BSDF_GLOSSY_TRANSMISSION;
      res.wi = refract(wo, wh, cosThetaOH, cosThetaTH, self->eta);
    }
  } else {
    // Sample the energy compensation lobe
    float ss1 = (ss - singlePickProb) * rcp(1.f - singlePickProb);
    doReflection = (ss1 <= self->fmsRatio);
    res.wi = getFrame(super) * cosineSampleHemisphere(s);

    if (doReflection) {
      // Reflection
      res.type = BSDF_DIFFUSE_REFLECTION;
    } else {
      // Transmission
      res.type = BSDF_DIFFUSE_TRANSMISSION;
      res.wi = neg(res.wi);
    }
  }

  float cosThetaI = dot(res.wi, getN(super));
  if ((doReflection && cosThetaI <= 0.f) || (!doReflection && cosThetaI >= 0.f))
    return make_BSDF_SampleRes_zero();

  // Evaluate the energy compensation lobe
  BSDF_EvalRes fms;
  if (doReflection) {
    // Reflection
    float Ei = MicrofacetDielectricAlbedoTable_eval(
        self->microfacetAlbedoTables, cosThetaI, self->eta, self->roughness);
    fms.pdf = self->fmsRatio * cosineSampleHemispherePDF(cosThetaI);
    fms.value = make_vec3f(self->fmsRatio * (1.f - Eo) * (1.f - Ei)
        * rcp((float)pi * (1.f - self->EavgEta)) * cosThetaI);
  } else {
    // Transmission
    float Ei =
        MicrofacetDielectricAlbedoTable_eval(self->microfacetAlbedoTables,
            abs(cosThetaI),
            rcp(self->eta),
            self->roughness);
    fms.pdf =
        (1.f - self->fmsRatio) * cosineSampleHemispherePDF(abs(cosThetaI));
    fms.value = make_vec3f((1.f - self->fmsRatio) * (1.f - Eo) * (1.f - Ei)
        * rcp((float)pi * (1.f - self->EavgRcpEta)) * abs(cosThetaI));
  }

  // Evaluate the single-scattering lobe
  BSDF_EvalRes single = MicrofacetDielectric_evalSingle(super, wo, res.wi);

  // Compute the final result
  res.pdf = singlePickProb * single.pdf + (1.f - singlePickProb) * fms.pdf;
  res.weight = (single.value + fms.value) * rcp(res.pdf);
  return res;
}

inline void MicrofacetDielectric_Constructor(
    varying MicrofacetDielectric *uniform self,
    MicrofacetAlbedoTables *uniform microfacetAlbedoTables,
    const varying linear3f *uniform frame,
    float eta,
    float roughness,
    float anisotropy)
{
  BSDF_Constructor(&self->super,
      make_vec3f(0.9f), // TODO better estimate
      BSDF_GLOSSY | BSDF_DIFFUSE,
      BSDF_TYPE_MICROFACET_DIELECTRIC,
      frame);

  self->eta = eta;
  self->microfacet =
      make_GGXDistribution(roughnessToAlpha(roughness, anisotropy));
  self->roughness = roughness;
  self->microfacetAlbedoTables = microfacetAlbedoTables;

  // Energy compensation
  self->EavgEta = MicrofacetDielectricAlbedoTable_evalAvg(
      microfacetAlbedoTables, eta, roughness);
  self->EavgRcpEta = MicrofacetDielectricAlbedoTable_evalAvg(
      microfacetAlbedoTables, rcp(eta), roughness);
  self->FavgEta = fresnelDielectricAvg(eta);
  self->FavgRcpEta = fresnelDielectricAvg(rcp(eta));
  float a = (1.f - self->FavgEta) * rcp(1.f - self->EavgRcpEta);
  float b = (1.f - self->FavgRcpEta)
      * rcp(1.f
          - self->EavgEta) /* * sqr(eta)*/; // ignore solid angle compression
  float x = b * rcp(a + b);
  self->fmsRatio = (1.f - x * (1.f - self->FavgEta));
}

inline varying BSDF *uniform MicrofacetDielectric_create(
    uniform ShadingContext *uniform ctx,
    MicrofacetAlbedoTables *uniform microfacetAlbedoTables,
    const varying linear3f *uniform frame,
    float eta,
    float roughness,
    float anisotropy)
{
  varying MicrofacetDielectric *uniform self =
      (varying MicrofacetDielectric * uniform)
          ShadingContext_alloc(ctx, sizeof(MicrofacetDielectric));
  MicrofacetDielectric_Constructor(
      self, microfacetAlbedoTables, frame, eta, roughness, anisotropy);
  return &self->super;
}

OSPRAY_END_ISPC_NAMESPACE
