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

#pragma once

#include <openvkl/common.isph>

#include "max_iterator_size.h"

///////////////////////////////////////////////////////////////////////////////
// Sampling ///////////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////

typedef uniform VKLObject VKLSampler;

VKL_API void vklComputeSample4(const int *uniform valid,
                               const VKLSampler *uniform sampler,
                               const varying struct vkl_vec3f *uniform
                                   objectCoordinates,
                               varying float *uniform samples,
                               uniform unsigned int attributeIndex,
                               const varying float *uniform time);

VKL_API void vklComputeSample8(const int *uniform valid,
                               const VKLSampler *uniform sampler,
                               const varying struct vkl_vec3f *uniform
                                   objectCoordinates,
                               varying float *uniform samples,
                               uniform unsigned int attributeIndex,
                               const varying float *uniform time);

VKL_API void vklComputeSample16(const int *uniform valid,
                                const VKLSampler *uniform sampler,
                                const varying struct vkl_vec3f *uniform
                                    objectCoordinates,
                                varying float *uniform samples,
                                uniform unsigned int attributeIndex,
                                const varying float *uniform time);

VKL_FORCEINLINE varying float vklComputeSampleV(
    const VKLSampler *uniform sampler,
    const varying vkl_vec3f *uniform objectCoordinates,
    uniform unsigned int attributeIndex = 0,
    const varying float *uniform time   = NULL)
{
  varying bool mask = __mask;
  unmasked
  {
    varying int imask = mask ? -1 : 0;
  }

  varying float samples;

  if (sizeof(varying float) == 16) {
    vklComputeSample4((uniform int *uniform) & imask,
                      sampler,
                      objectCoordinates,
                      &samples,
                      attributeIndex,
                      time);
  } else if (sizeof(varying float) == 32) {
    vklComputeSample8((uniform int *uniform) & imask,
                      sampler,
                      objectCoordinates,
                      &samples,
                      attributeIndex,
                      time);
  } else if (sizeof(varying float) == 64) {
    vklComputeSample16((uniform int *uniform) & imask,
                       sampler,
                       objectCoordinates,
                       &samples,
                       attributeIndex,
                       time);
  }

  return samples;
}

VKL_API void vklComputeGradient4(const int *uniform valid,
                                 const VKLSampler *uniform sampler,
                                 const varying struct vkl_vec3f *uniform
                                     objectCoordinates,
                                 varying vkl_vec3f *uniform gradients,
                                 uniform unsigned int attributeIndex,
                                 const varying float *uniform time);

VKL_API void vklComputeGradient8(const int *uniform valid,
                                 const VKLSampler *uniform sampler,
                                 const varying struct vkl_vec3f *uniform
                                     objectCoordinates,
                                 varying vkl_vec3f *uniform gradients,
                                 uniform unsigned int attributeIndex,
                                 const varying float *uniform time);

VKL_API void vklComputeGradient16(const int *uniform valid,
                                  const VKLSampler *uniform sampler,
                                  const varying struct vkl_vec3f *uniform
                                      objectCoordinates,
                                  varying vkl_vec3f *uniform gradients,
                                  uniform unsigned int attributeIndex,
                                  const varying float *uniform time);

VKL_FORCEINLINE varying vkl_vec3f
vklComputeGradientV(const VKLSampler *uniform sampler,
                    const varying vkl_vec3f *uniform objectCoordinates,
                    uniform unsigned int attributeIndex = 0,
                    const varying float *uniform time   = NULL)
{
  varying bool mask = __mask;
  unmasked
  {
    varying int imask = mask ? -1 : 0;
  }

  varying vkl_vec3f gradients;

  if (sizeof(varying float) == 16) {
    vklComputeGradient4((uniform int *uniform) & imask,
                        sampler,
                        objectCoordinates,
                        &gradients,
                        attributeIndex,
                        time);
  } else if (sizeof(varying float) == 32) {
    vklComputeGradient8((uniform int *uniform) & imask,
                        sampler,
                        objectCoordinates,
                        &gradients,
                        attributeIndex,
                        time);
  } else if (sizeof(varying float) == 64) {
    vklComputeGradient16((uniform int *uniform) & imask,
                         sampler,
                         objectCoordinates,
                         &gradients,
                         attributeIndex,
                         time);
  }

  return gradients;
}

VKL_API void vklComputeSampleM4(const int *uniform valid,
                                const VKLSampler *uniform sampler,
                                const varying struct vkl_vec3f *uniform
                                    objectCoordinates,
                                varying float *uniform samples,
                                uniform unsigned int M,
                                const unsigned int *uniform attributeIndices,
                                const varying float *uniform time);

VKL_API void vklComputeSampleM8(const int *uniform valid,
                                const VKLSampler *uniform sampler,
                                const varying struct vkl_vec3f *uniform
                                    objectCoordinates,
                                varying float *uniform samples,
                                uniform unsigned int M,
                                const unsigned int *uniform attributeIndices,
                                const varying float *uniform time);

VKL_API void vklComputeSampleM16(const int *uniform valid,
                                 const VKLSampler *uniform sampler,
                                 const varying struct vkl_vec3f *uniform
                                     objectCoordinates,
                                 varying float *uniform samples,
                                 uniform unsigned int M,
                                 const unsigned int *uniform attributeIndices,
                                 const varying float *uniform time);

VKL_FORCEINLINE void vklComputeSampleMV(
    const VKLSampler *uniform sampler,
    const varying vkl_vec3f *uniform objectCoordinates,
    varying float *uniform samples,
    uniform unsigned int M,
    const unsigned int *uniform attributeIndices,
    const varying float *uniform time = NULL)
{
  varying bool mask = __mask;
  unmasked
  {
    varying int imask = mask ? -1 : 0;
  }

  if (sizeof(varying float) == 16) {
    vklComputeSampleM4((uniform int *uniform) & imask,
                       sampler,
                       objectCoordinates,
                       samples,
                       M,
                       attributeIndices,
                       time);
  } else if (sizeof(varying float) == 32) {
    vklComputeSampleM8((uniform int *uniform) & imask,
                       sampler,
                       objectCoordinates,
                       samples,
                       M,
                       attributeIndices,
                       time);
  } else if (sizeof(varying float) == 64) {
    vklComputeSampleM16((uniform int *uniform) & imask,
                        sampler,
                        objectCoordinates,
                        samples,
                        M,
                        attributeIndices,
                        time);
  }
}

///////////////////////////////////////////////////////////////////////////////
// Interval iterators /////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////

struct IntervalIterator;
typedef IntervalIterator *uniform VKLIntervalIterator;

struct VKLInterval
{
  vkl_range1f tRange;
  vkl_range1f valueRange;
  float nominalDeltaT;
};

VKL_API uniform size_t
vklGetIntervalIteratorSize4(const VKLIntervalIteratorContext *uniform context);

VKL_API uniform size_t
vklGetIntervalIteratorSize8(const VKLIntervalIteratorContext *uniform context);

VKL_API uniform size_t
vklGetIntervalIteratorSize16(const VKLIntervalIteratorContext *uniform context);

VKL_FORCEINLINE uniform size_t
vklGetIntervalIteratorSizeV(const VKLIntervalIteratorContext *uniform context)
{
  varying bool mask = __mask;
  unmasked
  {
    varying int imask = mask ? -1 : 0;
  }

  if (sizeof(varying float) == 16)
    return vklGetIntervalIteratorSize4(context);
  else if (sizeof(varying float) == 32)
    return vklGetIntervalIteratorSize8(context);
  else if (sizeof(varying float) == 64)
    return vklGetIntervalIteratorSize16(context);
}

VKL_API VKLIntervalIterator
vklInitIntervalIterator4(const int *uniform valid,
                         const VKLIntervalIteratorContext *uniform context,
                         const varying vkl_vec3f *uniform origin,
                         const varying vkl_vec3f *uniform direction,
                         const varying vkl_range1f *uniform tRange,
                         const varying float *uniform time,
                         void *uniform buffer);

VKL_API VKLIntervalIterator
vklInitIntervalIterator8(const int *uniform valid,
                         const VKLIntervalIteratorContext *uniform context,
                         const varying vkl_vec3f *uniform origin,
                         const varying vkl_vec3f *uniform direction,
                         const varying vkl_range1f *uniform tRange,
                         const varying float *uniform time,
                         void *uniform buffer);

VKL_API VKLIntervalIterator
vklInitIntervalIterator16(const int *uniform valid,
                          const VKLIntervalIteratorContext *uniform context,
                          const varying vkl_vec3f *uniform origin,
                          const varying vkl_vec3f *uniform direction,
                          const varying vkl_range1f *uniform tRange,
                          const varying float *uniform time,
                          void *uniform buffer);

VKL_FORCEINLINE VKLIntervalIterator
vklInitIntervalIteratorV(const VKLIntervalIteratorContext *uniform context,
                         const varying vkl_vec3f *uniform origin,
                         const varying vkl_vec3f *uniform direction,
                         const varying vkl_range1f *uniform tRange,
                         const varying float *uniform time,
                         void *uniform buffer)
{
  varying bool mask = __mask;
  unmasked
  {
    varying int imask = mask ? -1 : 0;
  }

  if (sizeof(varying float) == 16) {
    return vklInitIntervalIterator4((uniform int *uniform) & imask,
                                    context,
                                    origin,
                                    direction,
                                    tRange,
                                    time,
                                    buffer);
  } else if (sizeof(varying float) == 32) {
    return vklInitIntervalIterator8((uniform int *uniform) & imask,
                                    context,
                                    origin,
                                    direction,
                                    tRange,
                                    time,
                                    buffer);
  } else if (sizeof(varying float) == 64) {
    return vklInitIntervalIterator16((uniform int *uniform) & imask,
                                     context,
                                     origin,
                                     direction,
                                     tRange,
                                     time,
                                     buffer);
  }
}

VKL_API void vklIterateInterval4(const int *uniform valid,
                                 VKLIntervalIterator iterator,
                                 varying VKLInterval *uniform interval,
                                 varying int *uniform result);

VKL_API void vklIterateInterval8(const int *uniform valid,
                                 VKLIntervalIterator iterator,
                                 varying VKLInterval *uniform interval,
                                 varying int *uniform result);

VKL_API void vklIterateInterval16(const int *uniform valid,
                                  VKLIntervalIterator iterator,
                                  varying VKLInterval *uniform interval,
                                  varying int *uniform result);

VKL_FORCEINLINE varying bool vklIterateIntervalV(
    VKLIntervalIterator iterator, varying VKLInterval *uniform interval)
{
  varying bool mask = __mask;
  unmasked
  {
    varying int imask = mask ? -1 : 0;
  }
  if (sizeof(varying float) == 16) {
    varying int iresult;
    vklIterateInterval4(
        (uniform int *uniform) & imask, iterator, interval, &iresult);
    varying bool result = iresult ? true : false;
    return result;
  } else if (sizeof(varying float) == 32) {
    varying int iresult;
    vklIterateInterval8(
        (uniform int *uniform) & imask, iterator, interval, &iresult);
    varying bool result = iresult ? true : false;
    return result;
  } else if (sizeof(varying float) == 64) {
    varying int iresult;
    vklIterateInterval16(
        (uniform int *uniform) & imask, iterator, interval, &iresult);
    varying bool result = iresult ? true : false;
    return result;
  }
}

///////////////////////////////////////////////////////////////////////////////
// Hit iterators //////////////////////////////////////////////////////////////
///////////////////////////////////////////////////////////////////////////////

struct HitIterator;
typedef HitIterator *uniform VKLHitIterator;

struct VKLHit
{
  float t;
  float sample;
  float epsilon;
};

VKL_API uniform size_t
vklGetHitIteratorSize4(const VKLHitIteratorContext *uniform context);

VKL_API uniform size_t
vklGetHitIteratorSize8(const VKLHitIteratorContext *uniform context);

VKL_API uniform size_t
vklGetHitIteratorSize16(const VKLHitIteratorContext *uniform context);

VKL_FORCEINLINE uniform size_t
vklGetHitIteratorSizeV(const VKLHitIteratorContext *uniform context)
{
  varying bool mask = __mask;
  unmasked
  {
    varying int imask = mask ? -1 : 0;
  }

  if (sizeof(varying float) == 16)
    return vklGetHitIteratorSize4(context);
  else if (sizeof(varying float) == 32)
    return vklGetHitIteratorSize8(context);
  else if (sizeof(varying float) == 64)
    return vklGetHitIteratorSize16(context);
}

VKL_API VKLHitIterator
vklInitHitIterator4(const int *uniform valid,
                    const VKLHitIteratorContext *uniform context,
                    const varying vkl_vec3f *uniform origin,
                    const varying vkl_vec3f *uniform direction,
                    const varying vkl_range1f *uniform tRange,
                    const varying float *uniform time,
                    void *uniform buffer);

VKL_API VKLHitIterator
vklInitHitIterator8(const int *uniform valid,
                    const VKLHitIteratorContext *uniform context,
                    const varying vkl_vec3f *uniform origin,
                    const varying vkl_vec3f *uniform direction,
                    const varying vkl_range1f *uniform tRange,
                    const varying float *uniform time,
                    void *uniform buffer);

VKL_API VKLHitIterator
vklInitHitIterator16(const int *uniform valid,
                     const VKLHitIteratorContext *uniform context,
                     const varying vkl_vec3f *uniform origin,
                     const varying vkl_vec3f *uniform direction,
                     const varying vkl_range1f *uniform tRange,
                     const varying float *uniform time,
                     void *uniform buffer);

VKL_FORCEINLINE VKLHitIterator
vklInitHitIteratorV(const VKLHitIteratorContext *uniform context,
                    const varying vkl_vec3f *uniform origin,
                    const varying vkl_vec3f *uniform direction,
                    const varying vkl_range1f *uniform tRange,
                    const varying float *uniform time,
                    void *uniform buffer)
{
  varying bool mask = __mask;
  unmasked
  {
    varying int imask = mask ? -1 : 0;
  }

  if (sizeof(varying float) == 16) {
    return vklInitHitIterator4((uniform int *uniform) & imask,
                               context,
                               origin,
                               direction,
                               tRange,
                               time,
                               buffer);
  } else if (sizeof(varying float) == 32) {
    return vklInitHitIterator8((uniform int *uniform) & imask,
                               context,
                               origin,
                               direction,
                               tRange,
                               time,
                               buffer);
  } else if (sizeof(varying float) == 64) {
    return vklInitHitIterator16((uniform int *uniform) & imask,
                                context,
                                origin,
                                direction,
                                tRange,
                                time,
                                buffer);
  }
}

VKL_API void vklIterateHit4(const int *uniform valid,
                            VKLHitIterator iterator,
                            varying VKLHit *uniform hit,
                            varying int *uniform result);

VKL_API void vklIterateHit8(const int *uniform valid,
                            VKLHitIterator iterator,
                            varying VKLHit *uniform hit,
                            varying int *uniform result);

VKL_API void vklIterateHit16(const int *uniform valid,
                             VKLHitIterator iterator,
                             varying VKLHit *uniform hit,
                             varying int *uniform result);

VKL_FORCEINLINE varying bool vklIterateHitV(VKLHitIterator iterator,
                                            varying VKLHit *uniform hit)
{
  varying bool mask = __mask;
  unmasked
  {
    varying int imask = mask ? -1 : 0;
  }
  if (sizeof(varying float) == 16) {
    varying int iresult;
    vklIterateHit4((uniform int *uniform) & imask, iterator, hit, &iresult);
    varying bool result = iresult ? true : false;
    return result;
  } else if (sizeof(varying float) == 32) {
    varying int iresult;
    vklIterateHit8((uniform int *uniform) & imask, iterator, hit, &iresult);
    varying bool result = iresult ? true : false;
    return result;
  } else if (sizeof(varying float) == 64) {
    varying int iresult;
    vklIterateHit16((uniform int *uniform) & imask, iterator, hit, &iresult);
    varying bool result = iresult ? true : false;
    return result;
  }
}
