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

#pragma once

#include "VdbGrid.h"
#include "../../common/temporal_data_interpolation.ih"

// ---------------------------------------------------------------------------
// Value range.
// ---------------------------------------------------------------------------

#define __vkl_template_VdbSampler_computeValueRange_denseZYX_constant( \
    voxelType)                                                         \
  inline uniform box1f                                                 \
      VdbSampler_computeValueRange_denseZYX_constant_##voxelType(      \
          const VdbGrid *uniform grid,                                 \
          uniform uint64 /*leafIndex*/,                                \
          uniform uint64 leafDataIndex,                                \
          const uniform vec2ui &xRange,                                \
          const uniform vec2ui &yRange,                                \
          const uniform vec2ui &zRange)                                \
  {                                                                    \
    uniform box1f valueRange = make_box1f(pos_inf, neg_inf);           \
    for (uniform unsigned int x = xRange.x; x < xRange.y; ++x) {       \
      for (uniform unsigned int y = yRange.x; y < yRange.y; ++y) {     \
        for (uniform unsigned int z = zRange.x; z < zRange.y; ++z) {   \
          const uniform uint64 voxelIdx =                              \
              __vkl_vdb_domain_offset_to_linear_uniform_leaf(x, y, z); \
          assert(voxelIdx < ((uniform uint64)1) << 32);                \
          const uniform uint32 v32 = ((uniform uint32)voxelIdx);       \
          extend(valueRange,                                           \
                 get_##voxelType(grid->leafData[leafDataIndex], v32)); \
        }                                                              \
      }                                                                \
    }                                                                  \
    return valueRange;                                                 \
  }

__vkl_template_VdbSampler_computeValueRange_denseZYX_constant(half)
__vkl_template_VdbSampler_computeValueRange_denseZYX_constant(float)

#undef __vkl_template_VdbSampler_computeValueRange_denseZYX_constant

#define __vkl_template_VdbSampler_computeValueRange_denseZYX_structured( \
    voxelType)                                                           \
  inline uniform box1f                                                   \
      VdbSampler_computeValueRange_denseZYX_structured_##voxelType(      \
          const VdbGrid *uniform grid,                                   \
          uniform uint64 leafIndex,                                      \
          uniform uint64 leafDataIndex,                                  \
          const uniform vec2ui &xRange,                                  \
          const uniform vec2ui &yRange,                                  \
          const uniform vec2ui &zRange)                                  \
  {                                                                      \
    const uniform uint32 numTimesteps =                                  \
        grid->leafStructuredTimesteps[leafIndex];                        \
    uniform box1f valueRange = make_box1f(pos_inf, neg_inf);             \
                                                                         \
    for (uniform unsigned int x = xRange.x; x < xRange.y; ++x) {         \
      for (uniform unsigned int y = yRange.x; y < yRange.y; ++y) {       \
        for (uniform unsigned int z = zRange.x; z < zRange.y; ++z) {     \
          const uniform uint64 voxelIdx =                                \
              numTimesteps *                                             \
              __vkl_vdb_domain_offset_to_linear_uniform_leaf(x, y, z);   \
          for (uniform unsigned int t = 0; t < numTimesteps; ++t) {      \
            assert((voxelIdx + t) < ((uniform uint64)1) << 32);          \
            const uniform uint32 v32 = ((uniform uint32)(voxelIdx + t)); \
            extend(valueRange,                                           \
                   get_##voxelType(grid->leafData[leafDataIndex], v32)); \
          }                                                              \
        }                                                                \
      }                                                                  \
    }                                                                    \
    return valueRange;                                                   \
  }

__vkl_template_VdbSampler_computeValueRange_denseZYX_structured(half)
__vkl_template_VdbSampler_computeValueRange_denseZYX_structured(float)

#undef __vkl_template_VdbSampler_computeValueRange_denseZYX_structured

#define __vkl_template_VdbSampler_computeValueRange_denseZYX_unstructured( \
    voxelType)                                                             \
  inline uniform box1f                                                     \
      VdbSampler_computeValueRange_denseZYX_unstructured_##voxelType(      \
          const VdbGrid *uniform grid,                                     \
          uniform uint64 leafIndex,                                        \
          uniform uint64 leafDataIndex,                                    \
          const uniform vec2ui &xRange,                                    \
          const uniform vec2ui &yRange,                                    \
          const uniform vec2ui &zRange)                                    \
  {                                                                        \
    uniform box1f valueRange = make_box1f(pos_inf, neg_inf);               \
                                                                           \
    for (uniform unsigned int x = xRange.x; x < xRange.y; ++x) {           \
      for (uniform unsigned int y = yRange.x; y < yRange.y; ++y) {         \
        for (uniform unsigned int z = zRange.x; z < zRange.y; ++z) {       \
          const uniform uint64 voxelIdx =                                  \
              __vkl_vdb_domain_offset_to_linear_uniform_leaf(x, y, z);     \
          assert(voxelIdx < ((uniform uint64)1) << 32);                    \
          const uniform uint32 v32 = ((uniform uint32)voxelIdx);           \
          valueRange = box_extend(valueRange,                              \
                     computeValueRangeTemporallyUnstructured_##voxelType(  \
                         grid->leafData + leafDataIndex,                   \
                         grid->leafUnstructuredIndices + leafIndex,        \
                         v32));                                            \
        }                                                                  \
      }                                                                    \
    }                                                                      \
                                                                           \
    return valueRange;                                                     \
  }

__vkl_template_VdbSampler_computeValueRange_denseZYX_unstructured(half)
__vkl_template_VdbSampler_computeValueRange_denseZYX_unstructured(float)

#undef __vkl_template_VdbSampler_computeValueRange_denseZYX_unstructured

// ---------------------------------------------------------------------------
// Constant leaf sampling.
// ---------------------------------------------------------------------------

#define __vkl_template_VdbSampler_sample_denseZYX_constant(voxelType)        \
  inline uniform float                                                       \
      VdbSampler_sample_uniform_uniform_denseZYX_constant_##voxelType(       \
          const VdbGrid *uniform grid,                                       \
          uniform vkl_uint64 /*leafIndex*/,                                  \
          uniform vkl_uint64 leafDataIndex,                                  \
          const uniform vec3ui &offset,                                      \
          uniform float /*time*/)                                            \
  {                                                                          \
    const uniform uint64 voxelIdx =                                          \
        __vkl_vdb_domain_offset_to_linear_uniform_leaf(                      \
            offset.x, offset.y, offset.z);                                   \
    assert(voxelIdx < ((uniform uint64)1) << 32);                            \
    const uniform uint32 v32 = ((uniform uint32)voxelIdx);                   \
                                                                             \
    return get_##voxelType(grid->leafData[leafDataIndex], v32);              \
  }                                                                          \
                                                                             \
  inline float                                                               \
      VdbSampler_sample_uniform_varying_denseZYX_constant_##voxelType(       \
          const VdbGrid *uniform grid,                                       \
          uniform vkl_uint64 /*leafIndex*/,                                  \
          uniform vkl_uint64 leafDataIndex,                                  \
          const vec3ui &offset,                                              \
          const float & /*time*/)                                            \
  {                                                                          \
    const uint64 voxelIdx = __vkl_vdb_domain_offset_to_linear_varying_leaf(  \
        offset.x, offset.y, offset.z);                                       \
    assert(voxelIdx < ((uint64)1) << 32);                                    \
    const uint32 v32 = ((varying uint32)voxelIdx);                           \
                                                                             \
    uniform uint32 uv32;                                                     \
    if (reduce_equal(v32, &uv32)) {                                          \
      return get_##voxelType(grid->leafData[leafDataIndex], uv32);           \
    } else {                                                                 \
      return get_##voxelType(grid->leafData[leafDataIndex], v32);            \
    }                                                                        \
  }                                                                          \
                                                                             \
  inline float                                                               \
      VdbSampler_sample_varying_varying_denseZYX_constant_##voxelType(       \
          const VdbGrid *uniform grid,                                       \
          const vkl_uint64 & /*leafIndex*/,                                  \
          const vkl_uint64 &leafDataIndex,                                   \
          const vec3ui &offset,                                              \
          const float & /*time*/)                                            \
  {                                                                          \
    const uint64 voxelIdx = __vkl_vdb_domain_offset_to_linear_varying_leaf(  \
        offset.x, offset.y, offset.z);                                       \
    assert(voxelIdx < ((uint64)1) << 32);                                    \
    const uint32 v32 = ((uint32)voxelIdx);                                   \
                                                                             \
    /* This uniform condition is more efficient that the varying condition   \
     * in get_##voxelType##(varying). */                                     \
    if (grid->allLeavesCompact) {                                            \
      return get_##voxelType##_compact(grid->leafData + leafDataIndex, v32); \
    } else {                                                                 \
      return get_##voxelType##_strided(grid->leafData + leafDataIndex, v32); \
    }                                                                        \
  }

__vkl_template_VdbSampler_sample_denseZYX_constant(half)
__vkl_template_VdbSampler_sample_denseZYX_constant(float)

#undef __vkl_template_VdbSampler_sample_denseZYX_constant

// ---------------------------------------------------------------------------
// Structured leaf sampling.
// ---------------------------------------------------------------------------

#define __vkl_template_VdbSampler_sample_denseZYX_structured(voxelType)  \
  inline uniform float                                                   \
      VdbSampler_sample_uniform_uniform_denseZYX_structured_##voxelType( \
          const VdbGrid *uniform grid,                                   \
          uniform vkl_uint64 leafIndex,                                  \
          uniform vkl_uint64 leafDataIndex,                              \
          const uniform vec3ui &offset,                                  \
          uniform float time)                                            \
  {                                                                      \
    const uniform int32 numTimesteps =                                   \
        grid->leafStructuredTimesteps[leafIndex];                        \
    const uniform uint64 voxelIdx =                                      \
        numTimesteps * __vkl_vdb_domain_offset_to_linear_uniform_leaf(   \
                           offset.x, offset.y, offset.z);                \
    assert(voxelIdx < ((uniform uint64)1) << 32);                        \
    const uniform uint32 v32 = ((uniform uint32)voxelIdx);               \
                                                                         \
    return interpolateTemporallyStructured_##voxelType(                  \
        grid->leafData + leafDataIndex, numTimesteps, v32, time);        \
  }                                                                      \
                                                                         \
  inline float                                                           \
      VdbSampler_sample_uniform_varying_denseZYX_structured_##voxelType( \
          const VdbGrid *uniform grid,                                   \
          uniform vkl_uint64 leafIndex,                                  \
          uniform vkl_uint64 leafDataIndex,                              \
          const vec3ui &offset,                                          \
          const float &time)                                             \
  {                                                                      \
    const uniform int32 numTimesteps =                                   \
        grid->leafStructuredTimesteps[leafIndex];                        \
    const uint64 voxelIdx =                                              \
        numTimesteps * __vkl_vdb_domain_offset_to_linear_varying_leaf(   \
                           offset.x, offset.y, offset.z);                \
    assert(voxelIdx < ((uint64)1) << 32);                                \
    const uint32 v32 = ((uint32)voxelIdx);                               \
                                                                         \
    uniform uint32 uv32;                                                 \
    uniform float utime;                                                 \
    if (reduce_equal(v32, &uv32) && reduce_equal(time, &utime)) {        \
      return interpolateTemporallyStructured_##voxelType(                \
          grid->leafData + leafDataIndex, numTimesteps, uv32, utime);    \
    } else {                                                             \
      return interpolateTemporallyStructured_##voxelType(                \
          grid->leafData + leafDataIndex, numTimesteps, v32, time);      \
    }                                                                    \
  }                                                                      \
                                                                         \
  inline float                                                           \
      VdbSampler_sample_varying_varying_denseZYX_structured_##voxelType( \
          const VdbGrid *uniform grid,                                   \
          const vkl_uint64 &leafIndex,                                   \
          const vkl_uint64 &leafDataIndex,                               \
          const vec3ui &offset,                                          \
          const float &time)                                             \
  {                                                                      \
    const int32 numTimesteps = grid->leafStructuredTimesteps[leafIndex]; \
    const uint64 voxelIdx =                                              \
        numTimesteps * __vkl_vdb_domain_offset_to_linear_varying_leaf(   \
                           offset.x, offset.y, offset.z);                \
    assert(voxelIdx < ((uint64)1) << 32);                                \
    const uint32 v32 = ((uint32)voxelIdx);                               \
                                                                         \
    return interpolateTemporallyStructured_##voxelType(                  \
        grid->leafData + leafDataIndex, numTimesteps, v32, time);        \
  }

__vkl_template_VdbSampler_sample_denseZYX_structured(half)
__vkl_template_VdbSampler_sample_denseZYX_structured(float)

#undef __vkl_template_VdbSampler_sample_denseZYX_structured

// ---------------------------------------------------------------------------
// Unstructured leaf sampling.
// ---------------------------------------------------------------------------

#define __vkl_template_VdbSampler_sample_denseZYX_unstructured(voxelType)     \
  inline uniform float                                                        \
      VdbSampler_sample_uniform_uniform_denseZYX_unstructured_##voxelType(    \
          const VdbGrid *uniform grid,                                        \
          uniform vkl_uint64 leafIndex,                                       \
          uniform vkl_uint64 leafDataIndex,                                   \
          const uniform vec3ui &offset,                                       \
          uniform float time)                                                 \
  {                                                                           \
    const uniform uint64 voxelIdx64 =                                         \
        __vkl_vdb_domain_offset_to_linear_uniform_leaf(                       \
            offset.x, offset.y, offset.z);                                    \
    assert(voxelIdx64 < ((uniform uint64)1) << 32);                           \
    const uniform uint32 voxelIdx = ((uniform uint32)voxelIdx64);             \
                                                                              \
    return interpolateTemporallyUnstructured_##voxelType(                     \
        grid->leafData + leafDataIndex,                                       \
        grid->leafUnstructuredIndices + leafIndex,                            \
        grid->leafUnstructuredTimes + leafIndex,                              \
        voxelIdx,                                                             \
        time);                                                                \
  }                                                                           \
                                                                              \
  inline float                                                                \
      VdbSampler_sample_uniform_varying_denseZYX_unstructured_##voxelType(    \
          const VdbGrid *uniform grid,                                        \
          uniform vkl_uint64 leafIndex,                                       \
          uniform vkl_uint64 leafDataIndex,                                   \
          const vec3ui &offset,                                               \
          const float &time)                                                  \
  {                                                                           \
    const uint64 voxelIdx64 = __vkl_vdb_domain_offset_to_linear_varying_leaf( \
        offset.x, offset.y, offset.z);                                        \
    assert(voxelIdx64 < ((uint64)1) << 32);                                   \
    const varying uint32 voxelIdx = ((uint32)voxelIdx64);                     \
                                                                              \
    return interpolateTemporallyUnstructured_##voxelType(                     \
        grid->leafData + leafDataIndex,                                       \
        grid->leafUnstructuredIndices + leafIndex,                            \
        grid->leafUnstructuredTimes + leafIndex,                              \
        voxelIdx,                                                             \
        time);                                                                \
  }                                                                           \
                                                                              \
  inline float                                                                \
      VdbSampler_sample_varying_varying_denseZYX_unstructured_##voxelType(    \
          const VdbGrid *uniform grid,                                        \
          const vkl_uint64 &leafIndex,                                        \
          const vkl_uint64 &leafDataIndex,                                    \
          const vec3ui &offset,                                               \
          const float &time)                                                  \
  {                                                                           \
    const uint64 voxelIdx64 = __vkl_vdb_domain_offset_to_linear_varying_leaf( \
        offset.x, offset.y, offset.z);                                        \
    assert(voxelIdx64 < ((uint64)1) << 32);                                   \
    const varying uint32 voxelIdx = ((uint32)voxelIdx64);                     \
                                                                              \
    return interpolateTemporallyUnstructured_##voxelType(                     \
        grid->leafData + leafDataIndex,                                       \
        grid->leafUnstructuredIndices + leafIndex,                            \
        grid->leafUnstructuredTimes + leafIndex,                              \
        voxelIdx,                                                             \
        time);                                                                \
  }

__vkl_template_VdbSampler_sample_denseZYX_unstructured(half)
__vkl_template_VdbSampler_sample_denseZYX_unstructured(float)

#undef __vkl_template_VdbSampler_sample_denseZYX_unstructured

