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

#pragma once

#include "../math/sampling.isph"

/*! \addtogroup rivl_render_embree_ivl */
/*! @{ */

/*! Reflects a viewing vector V at a normal N. */
inline uniform Vec3f reflect(const uniform Vec3f& V, const uniform Vec3f& N) { return 2.0f*dot(V,N)*N-V; }
#ifdef ISPC
inline varying Vec3f reflect(const uniform Vec3f& V, const varying Vec3f& N) { return 2.0f*dot(V,N)*N-V; }
inline varying Vec3f reflect(const varying Vec3f& V, const uniform Vec3f& N) { return 2.0f*dot(V,N)*N-V; }
inline varying Vec3f reflect(const varying Vec3f& V, const varying Vec3f& N) { return 2.0f*dot(V,N)*N-V; }
#endif

/*! Reflects a viewing vector V at a normal N. Cosine between V
 *  and N is given as input. */
inline Vec3f reflect(const Vec3f &V, const Vec3f &N, const float cosi) {
  return 2.0f*cosi*N-V;
}

// =======================================================
/*!Refracts a viewing vector V at a normal N using the relative
 *  refraction index eta. Eta is refraction index of outside medium
 *  (where N points into) divided by refraction index of the inside
 *  medium. The vectors V and N have to point towards the same side
 *  of the surface. The cosine between V and N is given as input and
 *  the cosine of -N and transmission ray is computed as output. */
inline Sample3f refract(const Vec3f& V, const Vec3f& N, const float eta,
                        const float cosi, float &cost)
{
  const float k = 1.0f-eta*eta*(1.0f-cosi*cosi);
  if (k < 0.0f) { cost = 0.0f; return make_Sample3f(make_Vec3f(0.f),0.0f); }
  cost = sqrt(k);
  return make_Sample3f(eta*(cosi*N-V)-cost*N, sqr(eta));
}

/*! Computes fresnel coefficient for media interface with relative
 *  refraction index eta. Eta is the outside refraction index
 *  divided by the inside refraction index. Both cosines have to be
 *  positive. */
inline float fresnelDielectric(const float cosi, const float cost, const float eta)
{
  const float Rper = (eta*cosi -     cost) * rcp(eta*cosi +     cost);
  const float Rpar = (    cosi - eta*cost) * rcp(    cosi + eta*cost);
  return 0.5f*(Rpar*Rpar + Rper*Rper);
}

  /*! Computes fresnel coefficient for media interface with relative
   *  refraction index eta. Eta is the outside refraction index
   *  divided by the inside refraction index. The cosine has to be
   *  positive. */
inline float fresnelDielectric(const float cosi, const float eta)
{
  const float k = 1.0f-eta*eta*(1.0f-cosi*cosi);
  if (k < 0.0f) return 1.0f;
  const float cost = sqrt(k);
  return fresnelDielectric(cosi, cost, eta);
}

/*! Computes fresnel coefficient for conductor medium with complex
 *  refraction index (eta,k). The cosine has to be positive. */
inline Vec3f fresnelConductor(const float cosi, const Vec3f& eta, const Vec3f& k)
{
  const Vec3f tmp = eta*eta + k*k;
  const Vec3f Rpar = (tmp * (cosi*cosi) - 2.0f*eta*cosi + make_Vec3f(1.0f)) *
    rcp(tmp * (cosi*cosi) + 2.0f*eta*cosi + make_Vec3f(1.0f));
  const Vec3f Rper = (tmp - 2.0f*eta*cosi + make_Vec3f(cosi*cosi)) *
    rcp(tmp + 2.0f*eta*cosi + make_Vec3f(cosi*cosi));
  return 0.5f * (Rpar + Rper);
}

// =======================================================
struct FresnelConductor {
  Vec3f eta;  //!< Real part of refraction index
  Vec3f k;    //!< Imaginary part of refraction index
};

inline Vec3f eval(varying const FresnelConductor& This, const float cosTheta) {
  return fresnelConductor(cosTheta,This.eta,This.k);
}

inline uniform FresnelConductor make_FresnelConductor(const uniform Vec3f& eta, const uniform Vec3f& k) {
  uniform FresnelConductor m; m.eta = eta; m.k = k; return m;
}

#if defined(ISPC)
inline varying FresnelConductor make_FresnelConductor(const varying Vec3f& eta, const varying Vec3f& k) {
  varying FresnelConductor m; m.eta = eta; m.k = k; return m;
}
#endif

// =======================================================
struct FresnelDielectric
{
  /*! refraction index of the medium the incident ray travels in */
  float etai;

  /*! refraction index of the medium the outgoing transmission rays
   *  travels in */
  float etat;
};

inline Vec3f eval(const uniform FresnelDielectric& This, const float cosTheta) {
  return make_Vec3f(fresnelDielectric(cosTheta,This.etai/This.etat));
}

inline uniform FresnelDielectric make_FresnelDielectric(const uniform float etai, const uniform float etat) {
  uniform FresnelDielectric m; m.etai = etai; m.etat = etat; return m;
}

#if defined(ISPC)
inline varying FresnelDielectric make_FresnelDielectric(const varying float etai, const varying float etat) {
  varying FresnelDielectric m; m.etai = etai; m.etat = etat; return m;
}
#endif

// =======================================================
struct PowerCosineDistribution {
  float exp;
};

inline float eval(const uniform PowerCosineDistribution &This, const float cosThetaH) {
  return (This.exp+2) * (1.0f/(2.0f*(M_PI))) * pow(abs(cosThetaH), This.exp);
}

#if defined(ISPC)

inline float eval(const varying PowerCosineDistribution &This, const float cosThetaH) {
  return (This.exp+2) * (1.0f/(2.0f*(M_PI))) * pow(abs(cosThetaH), This.exp);
}
#endif

/*! Samples the power cosine distribution. */
inline void sample(const uniform PowerCosineDistribution& This, const Vec3f& wo, const Vec3f& N, Sample3f &wi, const Vec2f s)
{
  Vec3f dir = powerCosineSampleHemisphere(This.exp,s);
  Sample3f wh;
  wh.v = frame(N) * dir;
  wh.pdf = powerCosineSampleHemispherePDF(dir,This.exp);
  Sample3f r = make_Sample3f(reflect(wo,wh.v),1.0f);
  wi = make_Sample3f(r.v,wh.pdf/(4.0f*abs(dot(wo,wh.v))));
}

/*! Samples the power cosine distribution. */
#if defined(ISPC)
inline void sample(const varying PowerCosineDistribution& This, const Vec3f& wo, const Vec3f& N, Sample3f &wi, const Vec2f s)
{
  Vec3f dir = powerCosineSampleHemisphere(This.exp,s);
  Sample3f wh;
  wh.v = frame(N) * dir;
  wh.pdf = powerCosineSampleHemispherePDF(dir,This.exp);
  Sample3f r = make_Sample3f(reflect(wo,wh.v),1.0f);
  wi = make_Sample3f(r.v,wh.pdf/(4.0f*abs(dot(wo,wh.v))));
}
#endif

inline uniform PowerCosineDistribution make_PowerCosineDistribution(const uniform float _exp) {
  uniform PowerCosineDistribution m; m.exp = _exp; return m;
}

#if defined(ISPC)
inline varying PowerCosineDistribution make_PowerCosineDistribution(const varying float _exp) {
  varying PowerCosineDistribution m; m.exp = _exp; return m;
}
#endif
