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

#pragma once

/*! \brief utility library containing sampling functions */

// convention is to return the sample (vec3f) generated from given vec2f
// 's'ample as last parameter sampling functions often come in pairs: sample and
// pdf (needed later for MIS) good reference is "Total Compendium" by Philip
// Dutre http://people.cs.kuleuven.be/~philip.dutre/GI/

#include "rkcommon/math/vec.ih"

OSPRAY_BEGIN_ISPC_NAMESPACE

inline vec3f cartesian(
    const float phi, const float sinTheta, const float cosTheta)
{
  float sinPhi, cosPhi;
  sincos(phi, &sinPhi, &cosPhi);
  return make_vec3f(cosPhi * sinTheta, sinPhi * sinTheta, cosTheta);
}

inline vec3f cartesian(const float phi, const float cosTheta)
{
  return cartesian(phi, cos2sin(cosTheta), cosTheta);
}

/// uniform sampling of hemisphere oriented along the +z-axis
////////////////////////////////////////////////////////////////////////////////

inline vec3f uniformSampleHemisphere(const vec2f s)
{
  const float phi = (float)two_pi * s.x;
  const float cosTheta = s.y;
  const float sinTheta = cos2sin(s.y);
  return cartesian(phi, sinTheta, cosTheta);
}

inline float uniformSampleHemispherePDF()
{
  return one_over_two_pi;
}

/// cosine-weighted sampling of hemisphere oriented along the +z-axis
////////////////////////////////////////////////////////////////////////////////

inline vec3f cosineSampleHemisphere(const vec2f s)
{
  const float phi = (float)two_pi * s.x;
  const float cosTheta = sqrt(s.y);
  const float sinTheta = sqrt(1.0f - s.y);
  return cartesian(phi, sinTheta, cosTheta);
}

inline float cosineSampleHemispherePDF(const vec3f &dir)
{
  return dir.z * (float)one_over_pi;
}

inline float cosineSampleHemispherePDF(float cosTheta)
{
  return cosTheta * (float)one_over_pi;
}

/// power cosine-weighted sampling of hemisphere oriented along the +z-axis
////////////////////////////////////////////////////////////////////////////////

inline vec3f powerCosineSampleHemisphere(const float n, const vec2f &s)
{
  const float phi = (float)two_pi * s.x;
  const float cosTheta = pow(s.y, 1.0f / (n + 1.0f));
  return cartesian(phi, cosTheta);
}

inline float powerCosineSampleHemispherePDF(
    const float cosTheta, const float n) // TODO: order of arguments
{
  return (n + 1.0f) * (float)one_over_two_pi * pow(cosTheta, n);
}

inline float powerCosineSampleHemispherePDF(
    const vec3f &dir, const float n) // TODO: order of arguments
{
  return (n + 1.0f) * (float)one_over_two_pi * pow(dir.z, n);
}

/// uniform sampling of cone of directions oriented along the +z-axis
////////////////////////////////////////////////////////////////////////////////

inline vec3f uniformSampleCone(const float cosAngle, const vec2f &s)
{
  const float phi = (float)two_pi * s.x;
  const float cosTheta = 1.0f - s.y * (1.0f - cosAngle);
  return cartesian(phi, cosTheta);
}

inline float uniformSampleConePDF(const float cosAngle)
{
  return rcp((float)two_pi * (1.0f - cosAngle));
}

#ifndef OSPRAY_TARGET_SYCL
inline uniform float uniformSampleConePDF(const uniform float cosAngle)
{
  return rcp((float)two_pi * (1.0f - cosAngle));
}
#endif

/// uniform sampling of ring
////////////////////////////////////////////////////////////////////////////////

inline vec3f uniformSampleRing(
    const float radius, const float innerRadius, const vec2f &s)
{
  const float r = sqrtf(lerp(s.x, sqr(innerRadius), sqr(radius)));
  const float phi = (float)two_pi * s.y;
  float sinPhi, cosPhi;
  sincos(phi, &sinPhi, &cosPhi);
  return make_vec3f(r * cosPhi, r * sinPhi, 0.f);
}

inline float uniformSampleRingPDF(const float radius, const float innerRadius)
{
  return rcp((float)pi * (sqr(radius) - sqr(innerRadius)));
}

#ifndef OSPRAY_TARGET_SYCL
inline uniform float uniformSampleRingPDF(
    const uniform float radius, const uniform float innerRadius)
{
  return rcp((float)pi * (sqr(radius) - sqr(innerRadius)));
}
#endif

/// uniform sampling of disk
////////////////////////////////////////////////////////////////////////////////

// TODO: just use uniformSampleRing with innerRadius = 0?
inline vec3f uniformSampleDisk(const float radius, const vec2f &s)
{
  const float r = sqrtf(s.x) * radius;
  const float phi = (float)two_pi * s.y;
  float sinPhi, cosPhi;
  sincos(phi, &sinPhi, &cosPhi);
  return make_vec3f(r * cosPhi, r * sinPhi, 0.f);
}

inline float uniformSampleDiskPDF(const float radius)
{
  return rcp((float)pi * sqr(radius));
}

#ifndef OSPRAY_TARGET_SYCL
inline uniform float uniformSampleDiskPDF(const uniform float radius)
{
  return rcp((float)pi * sqr(radius));
}
#endif

/// uniform sampling of triangle abc
////////////////////////////////////////////////////////////////////////////////

inline vec2f uniformSampleTriangleUV(const vec2f &s)
{
  const float su = sqrtf(s.x);
  vec2f uv;
  uv.x = s.y * su;
  uv.y = (1.0f - s.y) * su;

  return uv;
}

inline vec3f uniformSampleTriangle(
    const vec3f &a, const vec3f &b, const vec3f &c, const vec2f &s)
{
  const vec2f uv = uniformSampleTriangleUV(s);
  return a + uv.x * (b - a) + uv.y * (c - a);
}

inline float uniformSampleTrianglePDF(
    const vec3f &a, const vec3f &b, const vec3f &c)
{
  return 2.0f * rcp(length(cross(a - c, b - c)));
}

/// uniform sampling of sphere
////////////////////////////////////////////////////////////////////////////////

inline vec3f uniformSampleSphere(const float radius, const vec2f s)
{
  const float phi = (float)two_pi * s.x;
  const float cosTheta = radius * (1.f - 2.f * s.y);
  const float sinTheta = 2.f * radius * sqrt(s.y * (1.f - s.y));
  return cartesian(phi, sinTheta, cosTheta);
}

inline float uniformSampleSpherePDF(const float radius)
{
  return rcp((float)four_pi * sqr(radius));
}

#ifndef OSPRAY_TARGET_SYCL
inline uniform float uniformSampleSpherePDF(const uniform float radius)
{
  return rcp((float)four_pi * sqr(radius));
}
#endif
OSPRAY_END_ISPC_NAMESPACE