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

#pragma once

#include "BSDF.ih"
#include "Fresnel.ih"
#include "common/DifferentialGeometry.ih"

OSPRAY_BEGIN_ISPC_NAMESPACE

// FIXME: this should be cleaned up and integrated into Dielectric and
// MicrofacetDielectric
struct RobustDielectric
{
  BSDF super;

  float eta;

  // FIXME: BSDFs should *not* store the DifferentialGeometry or similar
  // parameters (i.e., Ng, Ns). These should be passed to sample/eval
  // DifferentialGeometry dgeom;
  vec3f Ng;
  vec3f Ns;
};

struct RefractionResult
{
  vec3f outgoingDirection;
  float FresnelReflectionCoefficient;
  bool mustReflect;
};

inline RefractionResult RefractionDirection(
    const vec3f &incomingDir, const vec3f &surfaceNormal, const float eta)
{
  RefractionResult res;
  res.mustReflect = false;

  float cos1 = dot(incomingDir, surfaceNormal);
  float cos1_abs = cos1;
  if (cos1_abs < 0.f) {
    cos1_abs = -cos1_abs;
  }
  // For head-on collisions the cosine of the angle of the ray with the surface
  // normal will be equal to one or minus one.
  bool headOn = (1.f - 1e-6f < cos1_abs && cos1_abs < 1.f - 1e-6f);
  if (headOn) {
    res.FresnelReflectionCoefficient = sqr((eta - 1.f) / (eta + 1.f));
    res.outgoingDirection = incomingDir;
  } else {
    vec3f dirParallelToNormal = cos1 * surfaceNormal;
    vec3f dirPerpendicularToNormal = incomingDir - dirParallelToNormal;
    vec3f b = normalize(dirPerpendicularToNormal);
    float sin1 = dot(dirPerpendicularToNormal, b);
    float sin2 = sin1 / eta;
    if (sin2 > 1.0f) { // Total internal reflection.
      res.outgoingDirection = make_vec3f(0.f); // Can't refract.
      res.FresnelReflectionCoefficient = 1.f;
      res.mustReflect = true;
    } else { // Refracting.
      float cos2 = sqrt(clamp(1.f - sin2 * sin2));
      vec3f N2 = normalize(dirParallelToNormal);
      res.outgoingDirection = N2 * cos2 + b * sin2;
      float c1 = cos1_abs;
      float c2 = cos2;
      float r1 = (eta * c1 - c2) / (eta * c1 + c2);
      float r2 = (c1 - eta * c2) / (c1 + eta * c2);
      res.FresnelReflectionCoefficient = 0.5f * (r1 * r1 + r2 * r2);
    }
  }
  // Clamp in case of roundoff error.
  res.FresnelReflectionCoefficient = clamp(res.FresnelReflectionCoefficient);

  return res;
}

inline vec3f ReflectionDirection(const vec3f &incomingDir, const vec3f &normal)
{ // Imagine a particle traveling toward the surface and then bouncing off in
  // the
  // reflection direction.

  // Parallel or anti-parallel direction.
  vec3f dirParallelToNormal = dot(incomingDir, normal) * normal;
  vec3f dirPerpendicularToNormal = incomingDir - dirParallelToNormal;

  // Outgoing direction.
  return (dirPerpendicularToNormal - dirParallelToNormal);
}

inline BSDF_SampleRes RobustDielectric_sample(const varying BSDF *uniform super,
    const vec3f &wo,
    const vec2f & /*randomV*/,
    float randomF)
{
  const varying RobustDielectric *uniform self =
      (const varying RobustDielectric *uniform)super;
  BSDF_SampleRes res;
  res.pdf = inf;

  // We want the actual ray direction into the surface.
  vec3f incomingDir = neg(wo);
  vec3f shadingNormal = self->Ns;
  vec3f geometricNormal = self->Ng;
  float eta = 1.f / self->eta; // ior(outside) / ior(inside)
  float glass_epsilon = 1e-6f;
  bool canReflect = (eta < 1.f - glass_epsilon || eta > 1.f + glass_epsilon);

  bool be_careful = true; // Extra test for edge cases when the shading normal
                          // is different from the geometry normal.
  bool nudge =
      true; // When we have grazing angles. Only used if be_careful==true.

  if (be_careful) {
    // Flip the normals to be on the same side of the surface as the incoming
    // ray.
    if (dot(incomingDir, shadingNormal) > 0.f) {
      shadingNormal = neg(shadingNormal);
    }
    if (dot(incomingDir, geometricNormal) > 0.f) {
      geometricNormal = neg(geometricNormal);
    }
  }

  float geometricCosine;
  vec3f dirR = ReflectionDirection(incomingDir, shadingNormal);

  RefractionResult refractionRes =
      RefractionDirection(incomingDir, shadingNormal, eta);
  vec3f dirT = refractionRes.outgoingDirection;
  float FresnelR = refractionRes.FresnelReflectionCoefficient;
  bool mustReflect = refractionRes.mustReflect;

  res.weight = make_vec3f(1.f);

  // Sample the reflection or the transmission using the provided random number.
  if ((randomF < FresnelR && canReflect) || mustReflect) { // Reflection
    res.wi = dirR; // Outgoing ray direction.
    res.type = BSDF_SPECULAR_REFLECTION;

    if (be_careful) {
      geometricCosine = dot(res.wi, geometricNormal);
      if (geometricCosine >= 0.f) {
        // Reflecting as expected.
        if (nudge && (geometricCosine < glass_epsilon)) { // Grazing angle.
          res.wi = res.wi
              + glass_epsilon
                  * geometricNormal; // Nudge the ray away from the surface.
          res.wi = normalize(res.wi);
        }
      } else { // geometricCosine < 0
        // We actually "refracted" since we are now on the other side of the
        // surface. Keep the ray type as reflection, but move the ray above the
        // surface.
        res.wi = res.wi - 2.f * geometricNormal * geometricCosine;
      }
    }
  } else { // Transmission
    res.wi = dirT; // Outgoing ray direction.
    res.type = BSDF_SPECULAR_TRANSMISSION;
    // res.weight = make_vec3f(rsqrt(self->eta));// Solid angle compression.

    if (be_careful) {
      geometricCosine = dot(res.wi, geometricNormal);
      if (geometricCosine <= 0.f) {
        // Transmitting as expected.
        if (nudge && (geometricCosine > -glass_epsilon)) { // Grazing angle.
          res.wi = res.wi
              - glass_epsilon
                  * geometricNormal; // Nudge the ray away from the surface.
          res.wi = normalize(res.wi);
        }
      } else { // geometricCosine > 0
        // We actually "reflected" since we are still on the same side of the
        // surface. Change ray type to reflection.
        res.type = BSDF_SPECULAR_REFLECTION;
      }
    }
  }

  return res;
}

inline void RobustDielectric_Constructor(varying RobustDielectric *uniform self,
    uniform ShadingContext *uniform ctx,
    const DifferentialGeometry &dg,
    float eta)
{
  const varying linear3f *uniform shadingframe =
      LinearSpace3f_create(ctx, frame(dg.Ns));

  BSDF_Constructor(&self->super,
      make_vec3f(1.0f),
      BSDF_SPECULAR,
      BSDF_TYPE_ROBUST_DIELECTRIC,
      shadingframe);
  self->eta = eta;
  self->Ns = dg.Ns;
  self->Ng = dg.Ng;
}

inline varying BSDF *uniform RobustDielectric_create(
    uniform ShadingContext *uniform ctx,
    const DifferentialGeometry &dg,
    float eta)

{
  varying RobustDielectric *uniform self = (varying RobustDielectric * uniform)
      ShadingContext_alloc(ctx, sizeof(RobustDielectric));

  RobustDielectric_Constructor(self, ctx, dg, eta);
  return &self->super;
}

OSPRAY_END_ISPC_NAMESPACE
