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

#pragma once

#include "VdbGrid.h"
#include "VdbLeafAccessObserver.ih"
#include "VdbSamplerShared.h"
#include "VdbSampler_denseHandler.ih"
#include "VdbSampler_leafHandler.ih"
#include "VdbSampler_packedHandler.ih"

#include "openvkl_vdb/VdbSamplerDispatchInner.ih"

// ---------------------------------------------------------------------------
// Find the origin of the leaf node containing domainOffset.
// Note: Performance is significantly better when these functions return
//       the leaf origin as opposed to having a reference parameter.
// ---------------------------------------------------------------------------

inline uniform vec3ui
VdbSampler_toLeafOrigin(const uniform vec3ui &domainOffset)
{
  static const uniform uint32 originMask =
      ~(((uniform uint32)(VKL_VDB_RES_LEAF - 1)));
  return make_vec3ui(domainOffset.x & originMask,
                     domainOffset.y & originMask,
                     domainOffset.z & originMask);
}

inline vec3ui VdbSampler_toLeafOrigin(const vec3ui &domainOffset)
{
  static const uniform uint32 originMask =
      ~(((uniform uint32)(VKL_VDB_RES_LEAF - 1)));
  return make_vec3ui(domainOffset.x & originMask,
                     domainOffset.y & originMask,
                     domainOffset.z & originMask);
}

// ---------------------------------------------------------------------------
// Map the given index space coordinate to a (root-relative) domain offset.
// Note that we do not check if the ic-rootOrigin is positive here; instead
// we rely on the fact that negative numbers in the subtraction below will be
// mapped very large values in the conversion.
// ---------------------------------------------------------------------------

inline uniform vec3ui VdbSampler_toDomainOffset(const uniform vec3i &indexCoord,
                                                const uniform vec3i &rootOrigin)
{
  return make_vec3ui(indexCoord.x - rootOrigin.x,
                     indexCoord.y - rootOrigin.y,
                     indexCoord.z - rootOrigin.z);
}

inline vec3ui VdbSampler_toDomainOffset(const vec3i &indexCoord,
                                        const uniform vec3i &rootOrigin)
{
  return make_vec3ui(indexCoord.x - rootOrigin.x,
                     indexCoord.y - rootOrigin.y,
                     indexCoord.z - rootOrigin.z);
}

// ---------------------------------------------------------------------------

inline uniform bool VdbSampler_isInDomain(const uniform vec3ui &activeSize,
                                          const uniform vec3ui &domainOffset)
{
  return domainOffset.x < activeSize.x && domainOffset.y < activeSize.y &&
         domainOffset.z < activeSize.z;
}

inline bool VdbSampler_isInDomain(const uniform vec3ui &activeSize,
                                  const vec3ui &domainOffset)
{
  return domainOffset.x < activeSize.x && domainOffset.y < activeSize.y &&
         domainOffset.z < activeSize.z;
}

// ---------------------------------------------------------------------------
// Clip the index space coordinate to the grid bounding box, and traverse
// the hierarchy to find the correct voxel.
// ---------------------------------------------------------------------------

// Uniform traversal.
noinline void VdbSampler_traverse(const VdbSamplerShared *uniform sampler,
                                  const uniform vec3i &ic,
                                  uniform uint64 &voxel,
                                  uniform vec3ui &domainOffset)
{
  assert(sampler);
  assert(sampler->grid);
  assert(sampler->grid->levels[0].numNodes == 1);
  assert(!sampler->grid->dense);

  voxel        = vklVdbVoxelMakeEmpty();
  domainOffset = VdbSampler_toDomainOffset(ic, sampler->grid->rootOrigin);

  if (VdbSampler_isInDomain(sampler->grid->activeSize, domainOffset)) {
    const uniform vec3ui leafOrigin = VdbSampler_toLeafOrigin(domainOffset);
    VdbSampler_dispatchInner_uniform_uniform_0(sampler, 0ul, leafOrigin, voxel);
  }
}

// Varying traversal.
noinline void VdbSampler_traverse(const VdbSamplerShared *uniform sampler,
                                  const vec3i &ic,
                                  uint64 &voxel,
                                  vec3ui &domainOffset)
{
  assert(sampler);
  assert(sampler->grid);
  assert(sampler->grid->levels[0].numNodes == 1);
  assert(!sampler->grid->dense);

  voxel        = vklVdbVoxelMakeEmpty();
  domainOffset = VdbSampler_toDomainOffset(ic, sampler->grid->rootOrigin);

  if (VdbSampler_isInDomain(sampler->grid->activeSize, domainOffset)) {
    const vec3ui leafOrigin = VdbSampler_toLeafOrigin(domainOffset);
    uniform vec3ui uniformLeafOrigin;
    if (reduce_equal(leafOrigin.x, &uniformLeafOrigin.x) &&
        reduce_equal(leafOrigin.y, &uniformLeafOrigin.y) &&
        reduce_equal(leafOrigin.z, &uniformLeafOrigin.z)) {
      // All queries are in the same leaf node. Do uniform traversal!
      uniform uint64 voxelU = vklVdbVoxelMakeEmpty();
      VdbSampler_dispatchInner_uniform_uniform_0(
          sampler, 0ul, uniformLeafOrigin, voxelU);
      voxel = voxelU;
    } else {
      // Varying traversal.
      VdbSampler_dispatchInner_uniform_varying_0(
          sampler, 0ul, leafOrigin, voxel);
    }
  }
}

// ---------------------------------------------------------------------------
// Given a domain relative index space coordinate and a voxel code, sample the
// leaf node.
// ---------------------------------------------------------------------------

// These helper functions handle all legal combinations of sampling uniform /
// varying voxels at uniform / varying offsets, and are used just below.
#define template_VdbSampler_sample_inner(univaryVoxel, univaryOffset)          \
  inline univaryOffset float VdbSampler_sample_inner(                          \
      const VdbSamplerShared *uniform sampler,                                 \
      const univaryVoxel uint64 &voxel,                                        \
      const univaryOffset vec3ui &domainOffset,                                \
      const univaryOffset float &time,                                         \
      const uniform uint32 attributeIndex)                                     \
  {                                                                            \
    assert(!sampler->grid->dense);                                             \
    univaryOffset float sample  = 0.f;                                         \
    const VdbGrid *uniform grid = sampler->grid;                               \
                                                                               \
    if (vklVdbVoxelIsLeafPtr(voxel)) {                                         \
      const uniform vkl_uint32 voxelType =                                     \
          sampler->grid->attributeTypes[attributeIndex];                       \
      const univaryVoxel uint64 leafIndex = vklVdbVoxelLeafGetIndex(voxel);    \
      if (sampler->grid->nodesPackedDense) {                                   \
        __vkl_vdb_packed_handler(                                              \
            sample =                                                           \
                VdbSampler_sample_##univaryVoxel##_##univaryOffset##_packed,   \
            grid->packedAddressing32,                                          \
            voxelType,                                                         \
            vklVdbVoxelLeafGetFormat(voxel),                                   \
            grid,                                                              \
            leafIndex,                                                         \
            attributeIndex,                                                    \
            domainOffset,                                                      \
            time)                                                              \
      } else {                                                                 \
        __vkl_vdb_leaf_handler(                                                \
            sample = VdbSampler_sample_##univaryVoxel##_##univaryOffset,       \
            voxelType,                                                         \
            vklVdbVoxelLeafGetFormat(voxel),                                   \
            grid->allLeavesConstant,                                           \
            vklVdbVoxelLeafGetTemporalFormat(voxel),                           \
            grid,                                                              \
            leafIndex,                                                         \
            vklVdbGetLeafDataIndex(grid, leafIndex, attributeIndex),           \
            domainOffset,                                                      \
            time)                                                              \
      }                                                                        \
    } else if (vklVdbVoxelIsError(voxel)) {                                    \
      univaryVoxel uint8 level;                                                \
      univaryVoxel uint32 voxelOffset;                                         \
      vklVdbVoxelErrorGet(voxel, level, voxelOffset);                          \
      const univaryVoxel range1f valueRange =                                  \
          grid->levels[level]                                                  \
              .valueRange[voxelOffset * grid->numAttributes + attributeIndex]; \
      sample = 0.5f * (valueRange.lower + valueRange.upper);                   \
    } else if (!VdbSampler_isInDomain(grid->activeSize, domainOffset)) {       \
      sample = sampler->super.super.volume->background[attributeIndex];        \
    }                                                                          \
    return sample;                                                             \
  }

template_VdbSampler_sample_inner(varying, varying);
template_VdbSampler_sample_inner(uniform, varying);
template_VdbSampler_sample_inner(uniform, uniform);
#undef template_VdbSampler_sample_inner

// Uniform sampling.
// Note: This function may seem unnecessary in the uniform path. However,
// keeping this thin wrapper ensures that downstream code is as uniform as
// possible.
noinline uniform float VdbSampler_sample(const VdbSamplerShared *uniform sampler,
                                         const uniform uint64 voxel,
                                         const uniform vec3ui &domainOffset,
                                         const uniform float time,
                                         const uniform uint32 attributeIndex)
{
  assert(!sampler->grid->dense);

  return VdbSampler_sample_inner(
      sampler, voxel, domainOffset, time, attributeIndex);
}

// Varying sampling.
noinline float VdbSampler_sample(const VdbSamplerShared *uniform sampler,
                                 const uint64 &voxel,
                                 const vec3ui &domainOffset,
                                 const float &time,
                                 const uniform uint32 attributeIndex)
{
  assert(!sampler->grid->dense);

  uniform uint64 uniformVoxel;
  if (reduce_equal(voxel, &uniformVoxel)) {
    return VdbSampler_sample_inner(
        sampler, uniformVoxel, domainOffset, time, attributeIndex);
  } else {
    return VdbSampler_sample_inner(
        sampler, voxel, domainOffset, time, attributeIndex);
  }
}

// ---------------------------------------------------------------------------
// Traverse and sample in a single call.
//
// This is generally more efficient than calling traverse() and sample()
// separately if only a single element needs to be looked up.
// ---------------------------------------------------------------------------

noinline uniform float VdbSampler_traverseAndSample(
    const VdbSamplerShared *uniform sampler,
    const uniform vec3i &ic,
    const uniform float time,
    const uniform uint32 attributeIndex)
{
  assert(sampler);
  assert(sampler->grid);
  assert(sampler->grid->levels[0].numNodes == 1);
  assert(!sampler->grid->dense);

  const uniform vec3ui domainOffset =
      VdbSampler_toDomainOffset(ic, sampler->grid->rootOrigin);
  if (VdbSampler_isInDomain(sampler->grid->activeSize, domainOffset)) {
    const uniform vec3ui leafOrigin = VdbSampler_toLeafOrigin(domainOffset);
    uniform uint64 voxelU           = vklVdbVoxelMakeEmpty();
    VdbSampler_dispatchInner_uniform_uniform_0(
        sampler, 0ul, leafOrigin, voxelU);
    return VdbSampler_sample_inner(
        sampler, voxelU, domainOffset, time, attributeIndex);
  }

  return sampler->super.super.volume->background[attributeIndex];
}

inline uniform bool reduce_equal(const vec3ui &v, vec3ui *uniform uv)
{
  return reduce_equal(v.x, &(uv->x)) && reduce_equal(v.y, &(uv->y)) &&
         reduce_equal(v.z, &(uv->z));
}

noinline float VdbSampler_traverseAndSample(const VdbSamplerShared *uniform sampler,
                                            const vec3i &ic,
                                            const float &time,
                                            const uniform uint32 attributeIndex)
{
  assert(sampler);
  assert(sampler->grid);
  assert(sampler->grid->levels[0].numNodes == 1);
  assert(!sampler->grid->dense);

  const vec3ui domainOffset =
      VdbSampler_toDomainOffset(ic, sampler->grid->rootOrigin);
  if (VdbSampler_isInDomain(sampler->grid->activeSize, domainOffset)) {
    const vec3ui leafOrigin = VdbSampler_toLeafOrigin(domainOffset);
    uniform vec3ui uniformLeafOrigin;
    if (reduce_equal(leafOrigin, &uniformLeafOrigin)) {
      // All queries are in the same leaf node. Do uniform traversal!
      uniform uint64 voxelU = vklVdbVoxelMakeEmpty();
      VdbSampler_dispatchInner_uniform_uniform_0(
          sampler, 0ul, uniformLeafOrigin, voxelU);
      return VdbSampler_sample_inner(
          sampler, voxelU, domainOffset, time, attributeIndex);
    } else {
      // Varying traversal.
      uint64 voxel = vklVdbVoxelMakeEmpty();
      VdbSampler_dispatchInner_uniform_varying_0(
          sampler, 0ul, leafOrigin, voxel);
      return VdbSampler_sample_inner(
          sampler, voxel, domainOffset, time, attributeIndex);
    }
  }
  return sampler->super.super.volume->background[attributeIndex];
}

// Specialized versions of the above, for dense volumes.

inline uniform float VdbSampler_traverseAndSample_dense(
    const VdbSamplerShared *uniform sampler,
    const uniform vec3i &ic,
    const uniform float time,
    const uniform uint32 attributeIndex)
{
  assert(sampler);
  assert(sampler->grid);
  assert(sampler->grid->dense);

  const uniform vec3ui domainOffset =
      VdbSampler_toDomainOffset(ic, sampler->grid->rootOrigin);
  if (VdbSampler_isInDomain(sampler->grid->activeSize, domainOffset)) {
    return sampler->denseLeafSample_uniform[attributeIndex](
        sampler->grid, attributeIndex, domainOffset, time);
  }
  return sampler->super.super.volume->background[attributeIndex];
}

inline float VdbSampler_traverseAndSample_dense(
    const VdbSamplerShared *uniform sampler,
    const vec3i &ic,
    const float &time,
    const uniform uint32 attributeIndex)
{
  assert(sampler);
  assert(sampler->grid);
  assert(sampler->grid->dense);

  const vec3ui domainOffset =
      VdbSampler_toDomainOffset(ic, sampler->grid->rootOrigin);
  if (VdbSampler_isInDomain(sampler->grid->activeSize, domainOffset)) {
    return sampler->denseLeafSample_varying[attributeIndex](
        sampler->grid, attributeIndex, domainOffset, time);
  }
  return sampler->super.super.volume->background[attributeIndex];
}
