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

#pragma once

#include "VdbSampler_stencilDispatch.ih"
#include "VdbSampler_traverseAndSample.ih"

// ---------------------------------------------------------------------------
// TRILINEAR interpolation.
// ---------------------------------------------------------------------------

// -----------------------------------------------------------------------------
// Filter stencils.
// -----------------------------------------------------------------------------

// Defines required for __vkl_stencil_dispatch
#define VKL_STENCIL_TRILINEAR_SIZE 8
#define VKL_STENCIL_TRILINEAR_OFFSETS                                 \
  {                                                                   \
    {0, 0, 0}, {0, 0, 1}, {0, 1, 0}, {0, 1, 1}, {1, 0, 0}, {1, 0, 1}, \
        {1, 1, 0},                                                    \
    {                                                                 \
      1, 1, 1                                                         \
    }                                                                 \
  }

/*
 * Compute voxel values for the eight corners required in trilinear
 * interpolation.
 * This is used for both sampling and gradient computation!
 */
inline void VdbSampler_computeVoxelValuesTrilinear(
    const VdbSamplerShared *uniform sampler,
    const vec3i &ic,
    const float &time,
    const uniform uint32 attributeIndex,
    float *uniform sample)  // Array of VKL_TARGET_WIDTH * 8 elements!
{
  assert(!sampler->grid->dense);

  __vkl_stencil_dispatch_varying(TRILINEAR, ic, time, {
    sample[tgtIdx] =
        VdbSampler_traverseAndSample(sampler, icDisp, timeDisp, attributeIndex);
  });
}

#define template_computeVoxelValuesTrilinear_dense_fast_inner_32(type)     \
  inline void computeVoxelValuesTrilinear_dense_fast_inner_##type##_32(    \
      const VdbGrid *uniform grid,                                         \
      const varying vec3ui &domainOffset,                                  \
      const uniform uint32 attributeIndex,                                 \
      float *uniform sample /* Array of VKL_TARGET_WIDTH * 8 elements! */) \
  {                                                                        \
    assert(grid->dense);                                                   \
    assert(grid->denseTemporalFormat == VKL_TEMPORAL_FORMAT_CONSTANT);     \
    assert(safe_32bit_indexing(grid->denseData[attributeIndex]));          \
    assert(all(VdbSampler_isInDomain(grid->activeSize, domainOffset)));    \
    assert((VdbSampler_isInDomain(grid->activeSize, domainOffset + 1)));   \
                                                                           \
    const uniform uint64 dx = 1;                                           \
    const uniform uint64 dy = grid->activeSize.x;                          \
    const uniform uint64 dz = grid->activeSize.x * grid->activeSize.y;     \
                                                                           \
    const varying uint32 voxelOfs =                                        \
        domainOffset.x * dx + domainOffset.y * dy + domainOffset.z * dz;   \
                                                                           \
    const uniform Data1D voxelData = grid->denseData[attributeIndex];      \
                                                                           \
    const uniform uint64 ofs000 = 0;                                       \
    const uniform uint64 ofs001 = dx;                                      \
    const varying float val000 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs000, voxelOfs);     \
    const varying float val001 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs001, voxelOfs);     \
                                                                           \
    const uniform uint64 ofs010 = dy;                                      \
    const uniform uint64 ofs011 = dy + dx;                                 \
    const varying float val010 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs010, voxelOfs);     \
    const varying float val011 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs011, voxelOfs);     \
                                                                           \
    const uniform uint64 ofs100 = dz;                                      \
    const uniform uint64 ofs101 = ofs100 + ofs001;                         \
    const varying float val100 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs100, voxelOfs);     \
    const varying float val101 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs101, voxelOfs);     \
                                                                           \
    const uniform uint64 ofs110 = ofs100 + ofs010;                         \
    const uniform uint64 ofs111 = ofs100 + ofs011;                         \
    const varying float val110 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs110, voxelOfs);     \
    const varying float val111 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs111, voxelOfs);     \
                                                                           \
    /* note the slightly different assignment order from the computations  \
     * above*/                                                             \
    sample[0 * VKL_TARGET_WIDTH + programIndex] = val000;                  \
    sample[1 * VKL_TARGET_WIDTH + programIndex] = val100;                  \
    sample[2 * VKL_TARGET_WIDTH + programIndex] = val010;                  \
    sample[3 * VKL_TARGET_WIDTH + programIndex] = val110;                  \
    sample[4 * VKL_TARGET_WIDTH + programIndex] = val001;                  \
    sample[5 * VKL_TARGET_WIDTH + programIndex] = val101;                  \
    sample[6 * VKL_TARGET_WIDTH + programIndex] = val011;                  \
    sample[7 * VKL_TARGET_WIDTH + programIndex] = val111;                  \
  }                                                                        \
                                                                           \
  inline void computeVoxelValuesTrilinear_dense_fast_inner_##type##_32(    \
      const VdbGrid *uniform grid,                                         \
      const uniform vec3ui &domainOffset,                                  \
      const uniform uint32 attributeIndex,                                 \
      float *uniform sample /* Array of 8 elements! */)                    \
  {                                                                        \
    assert(grid->dense);                                                   \
    assert(grid->denseTemporalFormat == VKL_TEMPORAL_FORMAT_CONSTANT);     \
    assert(safe_32bit_indexing(grid->denseData[attributeIndex]));          \
    assert(all(VdbSampler_isInDomain(grid->activeSize, domainOffset)));    \
    assert((VdbSampler_isInDomain(grid->activeSize, domainOffset + 1)));   \
                                                                           \
    const uniform uint64 dx = 1;                                           \
    const uniform uint64 dy = grid->activeSize.x;                          \
    const uniform uint64 dz = grid->activeSize.x * grid->activeSize.y;     \
                                                                           \
    const uniform uint32 voxelOfs =                                        \
        domainOffset.x * dx + domainOffset.y * dy + domainOffset.z * dz;   \
                                                                           \
    const uniform Data1D voxelData = grid->denseData[attributeIndex];      \
                                                                           \
    const uniform uint64 ofs000 = 0;                                       \
    const uniform uint64 ofs001 = dx;                                      \
    const uniform float val000 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs000, voxelOfs);     \
    const uniform float val001 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs001, voxelOfs);     \
                                                                           \
    const uniform uint64 ofs010 = dy;                                      \
    const uniform uint64 ofs011 = dy + dx;                                 \
    const uniform float val010 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs010, voxelOfs);     \
    const uniform float val011 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs011, voxelOfs);     \
                                                                           \
    const uniform uint64 ofs100 = dz;                                      \
    const uniform uint64 ofs101 = ofs100 + ofs001;                         \
    const uniform float val100 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs100, voxelOfs);     \
    const uniform float val101 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs101, voxelOfs);     \
                                                                           \
    const uniform uint64 ofs110 = ofs100 + ofs010;                         \
    const uniform uint64 ofs111 = ofs100 + ofs011;                         \
    const uniform float val110 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs110, voxelOfs);     \
    const uniform float val111 =                                           \
        get_##type(grid->denseData[attributeIndex], ofs111, voxelOfs);     \
                                                                           \
    /* note the slightly different assignment order from the computations  \
     * above*/                                                             \
    sample[0] = val000;                                                    \
    sample[1] = val100;                                                    \
    sample[2] = val010;                                                    \
    sample[3] = val110;                                                    \
    sample[4] = val001;                                                    \
    sample[5] = val101;                                                    \
    sample[6] = val011;                                                    \
    sample[7] = val111;                                                    \
  }

template_computeVoxelValuesTrilinear_dense_fast_inner_32(uint8);
template_computeVoxelValuesTrilinear_dense_fast_inner_32(int16);
template_computeVoxelValuesTrilinear_dense_fast_inner_32(uint16);
template_computeVoxelValuesTrilinear_dense_fast_inner_32(half);
template_computeVoxelValuesTrilinear_dense_fast_inner_32(float);
template_computeVoxelValuesTrilinear_dense_fast_inner_32(double);
#undef template_computeVoxelValuesTrilinear_dense_fast_inner_32

inline void VdbSampler_computeVoxelValuesTrilinear_dense(
    const VdbSamplerShared *uniform sampler,
    const vec3i &ic,
    const float &time,
    const uniform uint32 attributeIndex,
    float *uniform sample)  // Array of VKL_TARGET_WIDTH * 8 elements!
{
  assert(sampler->grid->dense);

  // if the entire domain of the trilinear interpolation is in bounds, and other
  // constraints are satisfied, we can take a fast path.
  const vec3ui domainOffset =
      VdbSampler_toDomainOffset(ic, sampler->grid->rootOrigin);

  if (sampler->grid->denseTemporalFormat == VKL_TEMPORAL_FORMAT_CONSTANT &&
      safe_32bit_indexing(sampler->grid->denseData[attributeIndex]) &&
      all(VdbSampler_isInDomain(sampler->grid->activeSize, domainOffset)) &&
      all(VdbSampler_isInDomain(sampler->grid->activeSize, domainOffset + 1))) {
    switch (sampler->grid->attributeTypes[attributeIndex]) {
    case (VKL_UCHAR): {
      computeVoxelValuesTrilinear_dense_fast_inner_uint8_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    case (VKL_SHORT): {
      computeVoxelValuesTrilinear_dense_fast_inner_int16_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    case (VKL_USHORT): {
      computeVoxelValuesTrilinear_dense_fast_inner_uint16_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    case (VKL_HALF): {
      computeVoxelValuesTrilinear_dense_fast_inner_half_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    case (VKL_FLOAT): {
      computeVoxelValuesTrilinear_dense_fast_inner_float_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    case (VKL_DOUBLE): {
      computeVoxelValuesTrilinear_dense_fast_inner_double_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    default: {
      assert(false);
    }
    }

  } else {
    __vkl_stencil_dispatch_simple_varying(TRILINEAR, ic, time, {
      sample[tgtIdx] = VdbSampler_traverseAndSample_dense(
          sampler, icDisp, timeDisp, attributeIndex);
    });
  }
}

inline void VdbSampler_computeVoxelValuesTrilinear_dense(
    const VdbSamplerShared *uniform sampler,
    const uniform vec3i &ic,
    const uniform float time,
    const uniform uint32 attributeIndex,
    float *uniform sample)  // Array of 8 elements!
{
  assert(sampler->grid->dense);

  // if the entire domain of the trilinear interpolation is in bounds, and other
  // constraints are satisfied, we can take a fast path.
  const uniform vec3ui domainOffset =
      VdbSampler_toDomainOffset(ic, sampler->grid->rootOrigin);

  if (sampler->grid->denseTemporalFormat == VKL_TEMPORAL_FORMAT_CONSTANT &&
      safe_32bit_indexing(sampler->grid->denseData[attributeIndex]) &&
      all(VdbSampler_isInDomain(sampler->grid->activeSize, domainOffset)) &&
      all(VdbSampler_isInDomain(sampler->grid->activeSize, domainOffset + 1))) {
    switch (sampler->grid->attributeTypes[attributeIndex]) {
    case (VKL_UCHAR): {
      computeVoxelValuesTrilinear_dense_fast_inner_uint8_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    case (VKL_SHORT): {
      computeVoxelValuesTrilinear_dense_fast_inner_int16_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    case (VKL_USHORT): {
      computeVoxelValuesTrilinear_dense_fast_inner_uint16_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    case (VKL_HALF): {
      computeVoxelValuesTrilinear_dense_fast_inner_half_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    case (VKL_FLOAT): {
      computeVoxelValuesTrilinear_dense_fast_inner_float_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    case (VKL_DOUBLE): {
      computeVoxelValuesTrilinear_dense_fast_inner_double_32(
          sampler->grid, domainOffset, attributeIndex, sample);
      break;
    }
    default: {
      assert(false);
    }
    }

  } else {
    __vkl_stencil_dispatch_uniform(TRILINEAR, ic, time, {
      sample[tgtIdx] = VdbSampler_traverseAndSample_dense(
          sampler, icDisp, timeDisp, attributeIndex);
    });
  }
}

/*
 * Traverse for the eight corners required in trilinear interpolation.
 */
inline void VdbSampler_traverseVoxelValuesTrilinear(
    const VdbSamplerShared *uniform sampler,
    const vec3i &ic,
    uint64 *uniform voxel,         // Array of VKL_TARGET_WIDTH * 8 elements!
    vec3ui *uniform domainOffset)  // Array of VKL_TARGET_WIDTH * 8 elements!
{
  assert(!sampler->grid->dense);

  const float time = 0.f;
  __vkl_stencil_dispatch_varying(TRILINEAR, ic, time, {
    uint64 voxelV;
    vec3ui domainOffsetV;
    VdbSampler_traverse(sampler, icDisp, voxelV, domainOffsetV);

    voxel[tgtIdx]        = voxelV;
    domainOffset[tgtIdx] = domainOffsetV;
  });
}

/*
 * Compute voxel values for the eight corners required in trilinear
 * interpolation for given pre-traversed voxel and domain offset.
 */
inline void VdbSampler_computeVoxelValuesTrilinear(
    const VdbSamplerShared *uniform sampler,
    const uint64 *uniform voxel,  // Array of VKL_TARGET_WIDTH * 8 elements!
    const vec3ui *uniform
        domainOffset,  // Array of VKL_TARGET_WIDTH * 8 elements!
    const float &time,
    const uniform uint32 attributeIndex,
    float *uniform sample)  // Array of VKL_TARGET_WIDTH * 8 elements!
{
  assert(!sampler->grid->dense);

  // Note: We might be able to avoid this dummy variable, but that would
  // complicate the dispatch macro.
  const vec3i ic = {0, 0, 0};
  __vkl_stencil_dispatch_varying(TRILINEAR, ic, time, {
    const uint64 voxelV        = voxel[tgtIdx];
    const vec3ui domainOffsetV = domainOffset[tgtIdx];
    sample[tgtIdx]             = VdbSampler_sample(
        sampler, voxelV, domainOffsetV, timeDisp, attributeIndex);
  });
}

/*
 * Trilinear sampling is a good default for directly visible volumes.
 * The implementation is optimized to exploit SIMD.
 */
inline varying float VdbSampler_interpolateTrilinear(
    const VdbSamplerShared *uniform sampler,
    const vec3f &indexCoordinates,
    const float &time,
    const uniform uint32 attributeIndex)
{
  assert(!sampler->grid->dense);

  const vec3i ic    = make_vec3i(floor(indexCoordinates.x),
                              floor(indexCoordinates.y),
                              floor(indexCoordinates.z));
  const vec3f delta = indexCoordinates - make_vec3f(ic);
  uniform float sample[VKL_TARGET_WIDTH * 8];
  VdbSampler_computeVoxelValuesTrilinear(
      sampler, ic, time, attributeIndex, sample);

  const varying float *uniform s = (const varying float *uniform) & sample;
  return lerp(
      delta.x,
      lerp(delta.y, lerp(delta.z, s[0], s[1]), lerp(delta.z, s[2], s[3])),
      lerp(delta.y, lerp(delta.z, s[4], s[5]), lerp(delta.z, s[6], s[7])));
}

inline varying float VdbSampler_interpolate_denseTrilinear(
    const VdbSamplerShared *uniform sampler,
    const vec3f &indexCoordinates,
    const float &time,
    const uniform uint32 attributeIndex)
{
  assert(sampler->grid->dense);

  const vec3i ic    = make_vec3i(floor(indexCoordinates.x),
                              floor(indexCoordinates.y),
                              floor(indexCoordinates.z));
  const vec3f delta = indexCoordinates - make_vec3f(ic);
  uniform float sample[VKL_TARGET_WIDTH * 8];
  VdbSampler_computeVoxelValuesTrilinear_dense(
      sampler, ic, time, attributeIndex, sample);

  const varying float *uniform s = (const varying float *uniform) & sample;
  return lerp(
      delta.x,
      lerp(delta.y, lerp(delta.z, s[0], s[1]), lerp(delta.z, s[2], s[3])),
      lerp(delta.y, lerp(delta.z, s[4], s[5]), lerp(delta.z, s[6], s[7])));
}

inline void VdbSampler_interpolateTrilinear(
    const VdbSamplerShared *uniform sampler,
    const vec3f &indexCoordinates,
    const float &time,
    const uniform uint32 M,
    const uint32 *uniform attributeIndices,
    float *uniform samples)
{
  assert(!sampler->grid->dense);

  const vec3i ic    = make_vec3i(floor(indexCoordinates.x),
                              floor(indexCoordinates.y),
                              floor(indexCoordinates.z));
  const vec3f delta = indexCoordinates - make_vec3f(ic);

  uniform uint64 voxel[VKL_TARGET_WIDTH * 8];
  uniform vec3ui domainOffset[VKL_TARGET_WIDTH * 8];
  VdbSampler_traverseVoxelValuesTrilinear(sampler, ic, voxel, domainOffset);

  for (uniform unsigned int a = 0; a < M; a++) {
    uniform float sample[VKL_TARGET_WIDTH * 8];
    VdbSampler_computeVoxelValuesTrilinear(
        sampler, voxel, domainOffset, time, attributeIndices[a], sample);

    const varying float *uniform s = (const varying float *uniform) & sample;
    varying float samplesA         = lerp(
        delta.x,
        lerp(delta.y, lerp(delta.z, s[0], s[1]), lerp(delta.z, s[2], s[3])),
        lerp(delta.y, lerp(delta.z, s[4], s[5]), lerp(delta.z, s[6], s[7])));

    samples[a * VKL_TARGET_WIDTH + programIndex] = samplesA;
  }
}

inline void VdbSampler_interpolate_denseTrilinear(
    const VdbSamplerShared *uniform sampler,
    const vec3f &indexCoordinates,
    const float &time,
    const uniform uint32 M,
    const uint32 *uniform attributeIndices,
    float *uniform samples)
{
  assert(sampler->grid->dense);

  const vec3i ic    = make_vec3i(floor(indexCoordinates.x),
                              floor(indexCoordinates.y),
                              floor(indexCoordinates.z));
  const vec3f delta = indexCoordinates - make_vec3f(ic);

  for (uniform unsigned int a = 0; a < M; a++) {
    uniform float sample[VKL_TARGET_WIDTH * 8];
    VdbSampler_computeVoxelValuesTrilinear_dense(
        sampler, ic, time, attributeIndices[a], sample);

    const varying float *uniform s = (const varying float *uniform) & sample;
    varying float samplesA         = lerp(
        delta.x,
        lerp(delta.y, lerp(delta.z, s[0], s[1]), lerp(delta.z, s[2], s[3])),
        lerp(delta.y, lerp(delta.z, s[4], s[5]), lerp(delta.z, s[6], s[7])));

    samples[a * VKL_TARGET_WIDTH + programIndex] = samplesA;
  }
}

/*
 * Uniform path. This allows us to skip the selection magic in the function
 * above if we know that there is only one query.
 */
inline uniform float VdbSampler_interpolateTrilinear(
    const VdbSamplerShared *uniform sampler,
    const uniform vec3f &indexCoordinates,
    const uniform float time,
    const uniform uint32 attributeIndex)
{
  assert(!sampler->grid->dense);

  const uniform vec3i ic = make_vec3i(floor(indexCoordinates.x),
                                      floor(indexCoordinates.y),
                                      floor(indexCoordinates.z));

  uniform float sample[8];

  __vkl_stencil_dispatch_uniform(TRILINEAR, ic, time, {
    sample[tgtIdx] =
        VdbSampler_traverseAndSample(sampler, icDisp, timeDisp, attributeIndex);
  });

  const uniform vec3f delta = indexCoordinates - make_vec3f(ic);
  return lerp(delta.x,
              lerp(delta.y,
                   lerp(delta.z, sample[0], sample[1]),
                   lerp(delta.z, sample[2], sample[3])),
              lerp(delta.y,
                   lerp(delta.z, sample[4], sample[5]),
                   lerp(delta.z, sample[6], sample[7])));
}

inline uniform float VdbSampler_interpolate_denseTrilinear(
    const VdbSamplerShared *uniform sampler,
    const uniform vec3f &indexCoordinates,
    const uniform float time,
    const uniform uint32 attributeIndex)
{
  assert(sampler->grid->dense);

  const uniform vec3i ic = make_vec3i(floor(indexCoordinates.x),
                                      floor(indexCoordinates.y),
                                      floor(indexCoordinates.z));

  uniform float sample[8];
  VdbSampler_computeVoxelValuesTrilinear_dense(
      sampler, ic, time, attributeIndex, sample);

  const uniform vec3f delta = indexCoordinates - make_vec3f(ic);
  return lerp(delta.x,
              lerp(delta.y,
                   lerp(delta.z, sample[0], sample[1]),
                   lerp(delta.z, sample[2], sample[3])),
              lerp(delta.y,
                   lerp(delta.z, sample[4], sample[5]),
                   lerp(delta.z, sample[6], sample[7])));
}

inline void VdbSampler_interpolateTrilinear(
    const VdbSamplerShared *uniform sampler,
    const uniform vec3f &indexCoordinates,
    const uniform float time,
    const uniform uint32 M,
    const uint32 *uniform attributeIndices,
    float *uniform samples)
{
  assert(!sampler->grid->dense);

  const uniform vec3i ic    = make_vec3i(floor(indexCoordinates.x),
                                      floor(indexCoordinates.y),
                                      floor(indexCoordinates.z));
  const uniform vec3f delta = indexCoordinates - make_vec3f(ic);

  uniform uint64 voxel[8];
  uniform vec3ui domainOffset[8];

  __vkl_stencil_dispatch_uniform(TRILINEAR, ic, time, {
    uint64 voxelV;
    vec3ui domainOffsetV;
    VdbSampler_traverse(sampler, icDisp, voxelV, domainOffsetV);

    voxel[tgtIdx]        = voxelV;
    domainOffset[tgtIdx] = domainOffsetV;
  });

  for (uniform unsigned int a = 0; a < M; a++) {
    uniform float sample[8];

    __vkl_stencil_dispatch_uniform(TRILINEAR, ic, time, {
      const uint64 voxelV        = voxel[tgtIdx];
      const vec3ui domainOffsetV = domainOffset[tgtIdx];
      sample[tgtIdx]             = VdbSampler_sample(
          sampler, voxelV, domainOffsetV, timeDisp, attributeIndices[a]);
    });

    samples[a] = lerp(delta.x,
                      lerp(delta.y,
                           lerp(delta.z, sample[0], sample[1]),
                           lerp(delta.z, sample[2], sample[3])),
                      lerp(delta.y,
                           lerp(delta.z, sample[4], sample[5]),
                           lerp(delta.z, sample[6], sample[7])));
  }
}

inline void VdbSampler_interpolate_denseTrilinear(
    const VdbSamplerShared *uniform sampler,
    const uniform vec3f &indexCoordinates,
    const uniform float time,
    const uniform uint32 M,
    const uint32 *uniform attributeIndices,
    float *uniform samples)
{
  assert(sampler->grid->dense);

  const uniform vec3i ic    = make_vec3i(floor(indexCoordinates.x),
                                      floor(indexCoordinates.y),
                                      floor(indexCoordinates.z));
  const uniform vec3f delta = indexCoordinates - make_vec3f(ic);

  for (uniform unsigned int a = 0; a < M; a++) {
    uniform float sample[8];
    VdbSampler_computeVoxelValuesTrilinear_dense(
        sampler, ic, time, attributeIndices[a], sample);

    samples[a] = lerp(delta.x,
                      lerp(delta.y,
                           lerp(delta.z, sample[0], sample[1]),
                           lerp(delta.z, sample[2], sample[3])),
                      lerp(delta.y,
                           lerp(delta.z, sample[4], sample[5]),
                           lerp(delta.z, sample[6], sample[7])));
  }
}

inline void VdbSampler_interpolateTrilinear(
    const VdbSamplerShared *uniform sampler,
    const uniform unsigned int N,
    const vec3f *uniform objectCoordinates,
    const float *uniform times,
    const uint32 uniform attributeIndex,
    float *uniform samples)
{
  assert(!sampler->grid->dense);

  foreach (i = 0 ... N) {
    const vec3f oc               = objectCoordinates[i];
    const vec3f indexCoordinates = xfmPoint(sampler->grid->objectToIndex, oc);
    const float time             = times ? times[i] : 0.f;
    samples[i]                   = VdbSampler_interpolateTrilinear(
        sampler, indexCoordinates, time, attributeIndex);
  }
}

inline void VdbSampler_interpolate_denseTrilinear(
    const VdbSamplerShared *uniform sampler,
    const uniform unsigned int N,
    const vec3f *uniform objectCoordinates,
    const float *uniform times,
    const uint32 uniform attributeIndex,
    float *uniform samples)
{
  assert(sampler->grid->dense);

  foreach (i = 0 ... N) {
    const vec3f oc               = objectCoordinates[i];
    const vec3f indexCoordinates = xfmPoint(sampler->grid->objectToIndex, oc);
    const float time             = times ? times[i] : 0.f;
    samples[i]                   = VdbSampler_interpolate_denseTrilinear(
        sampler, indexCoordinates, time, attributeIndex);
  }
}

inline void VdbSampler_interpolateTrilinear(
    const VdbSamplerShared *uniform sampler,
    const uniform unsigned int N,
    const vec3f *uniform objectCoordinates,
    const float *uniform times,
    const uniform uint32 M,
    const uint32 *uniform attributeIndices,
    float *uniform samples)
{
  assert(!sampler->grid->dense);

  foreach (i = 0 ... N) {
    const vec3f oc               = objectCoordinates[i];
    const vec3f indexCoordinates = xfmPoint(sampler->grid->objectToIndex, oc);

    const vec3i ic    = make_vec3i(floor(indexCoordinates.x),
                                floor(indexCoordinates.y),
                                floor(indexCoordinates.z));
    const vec3f delta = indexCoordinates - make_vec3f(ic);
    const float time  = times ? times[i] : 0.f;

    uniform uint64 voxel[VKL_TARGET_WIDTH * 8];
    uniform vec3ui domainOffset[VKL_TARGET_WIDTH * 8];
    VdbSampler_traverseVoxelValuesTrilinear(sampler, ic, voxel, domainOffset);

    for (uniform unsigned int a = 0; a < M; a++) {
      uniform float sample[VKL_TARGET_WIDTH * 8];
      VdbSampler_computeVoxelValuesTrilinear(
          sampler, voxel, domainOffset, time, attributeIndices[a], sample);

      const varying float *uniform s = (const varying float *uniform) & sample;
      varying float samplesA         = lerp(
          delta.x,
          lerp(delta.y, lerp(delta.z, s[0], s[1]), lerp(delta.z, s[2], s[3])),
          lerp(delta.y, lerp(delta.z, s[4], s[5]), lerp(delta.z, s[6], s[7])));

      samples[i * M + a] = samplesA;
    }
  }
}

inline void VdbSampler_interpolate_denseTrilinear(
    const VdbSamplerShared *uniform sampler,
    const uniform unsigned int N,
    const vec3f *uniform objectCoordinates,
    const float *uniform times,
    const uniform uint32 M,
    const uint32 *uniform attributeIndices,
    float *uniform samples)
{
  assert(sampler->grid->dense);

  foreach (i = 0 ... N) {
    const vec3f oc               = objectCoordinates[i];
    const vec3f indexCoordinates = xfmPoint(sampler->grid->objectToIndex, oc);

    const vec3i ic    = make_vec3i(floor(indexCoordinates.x),
                                floor(indexCoordinates.y),
                                floor(indexCoordinates.z));
    const vec3f delta = indexCoordinates - make_vec3f(ic);
    const float time  = times ? times[i] : 0.f;

    for (uniform unsigned int a = 0; a < M; a++) {
      uniform float sample[VKL_TARGET_WIDTH * 8];
      VdbSampler_computeVoxelValuesTrilinear_dense(
          sampler, ic, time, attributeIndices[a], sample);

      const varying float *uniform s = (const varying float *uniform) & sample;
      varying float samplesA         = lerp(
          delta.x,
          lerp(delta.y, lerp(delta.z, s[0], s[1]), lerp(delta.z, s[2], s[3])),
          lerp(delta.y, lerp(delta.z, s[4], s[5]), lerp(delta.z, s[6], s[7])));

      samples[i * M + a] = samplesA;
    }
  }
}

/*
 * Gradients in trilinear fields.
 */
inline vec3f VdbSampler_computeGradientTrilinear(
    const VdbSamplerShared *uniform sampler,
    const vec3f &indexCoordinates,
    const float &time,
    const uniform uint32 attributeIndex)
{
  assert(!sampler->grid->dense);

  const vec3i ic    = make_vec3i(floor(indexCoordinates.x),
                              floor(indexCoordinates.y),
                              floor(indexCoordinates.z));
  const vec3f delta = indexCoordinates - make_vec3f(ic);
  uniform float sample[VKL_TARGET_WIDTH * 8];
  VdbSampler_computeVoxelValuesTrilinear(
      sampler, ic, time, attributeIndex, sample);

  const varying float *uniform s = (const varying float *uniform) & sample;

  vec3f gradient;
  gradient.x = lerp(delta.y,
                    lerp(delta.z, s[4] - s[0], s[5] - s[1]),
                    lerp(delta.z, s[6] - s[2], s[7] - s[3]));
  gradient.y = lerp(delta.x,
                    lerp(delta.z, s[2] - s[0], s[3] - s[1]),
                    lerp(delta.z, s[6] - s[4], s[7] - s[5]));
  gradient.z = lerp(delta.x,
                    lerp(delta.y, s[1] - s[0], s[3] - s[2]),
                    lerp(delta.y, s[5] - s[4], s[7] - s[6]));
  return gradient;
}

inline vec3f VdbSampler_computeGradient_denseTrilinear(
    const VdbSamplerShared *uniform sampler,
    const vec3f &indexCoordinates,
    const float &time,
    const uniform uint32 attributeIndex)
{
  assert(sampler->grid->dense);

  const vec3i ic    = make_vec3i(floor(indexCoordinates.x),
                              floor(indexCoordinates.y),
                              floor(indexCoordinates.z));
  const vec3f delta = indexCoordinates - make_vec3f(ic);
  uniform float sample[VKL_TARGET_WIDTH * 8];
  VdbSampler_computeVoxelValuesTrilinear_dense(
      sampler, ic, time, attributeIndex, sample);

  const varying float *uniform s = (const varying float *uniform) & sample;

  vec3f gradient;
  gradient.x = lerp(delta.y,
                    lerp(delta.z, s[4] - s[0], s[5] - s[1]),
                    lerp(delta.z, s[6] - s[2], s[7] - s[3]));
  gradient.y = lerp(delta.x,
                    lerp(delta.z, s[2] - s[0], s[3] - s[1]),
                    lerp(delta.z, s[6] - s[4], s[7] - s[5]));
  gradient.z = lerp(delta.x,
                    lerp(delta.y, s[1] - s[0], s[3] - s[2]),
                    lerp(delta.y, s[5] - s[4], s[7] - s[6]));
  return gradient;
}

inline void VdbSampler_computeGradientTrilinear(
    const VdbSamplerShared *uniform sampler,
    uniform unsigned int N,
    const vec3f *uniform objectCoordinates,
    const float *uniform times,
    const uniform uint32 attributeIndex,
    vec3f *uniform gradients)
{
  assert(!sampler->grid->dense);

  foreach (i = 0 ... N) {
    const vec3f oc               = objectCoordinates[i];
    const vec3f indexCoordinates = xfmPoint(sampler->grid->objectToIndex, oc);
    const float time             = times ? times[i] : 0.f;
    const vec3f gradient         = VdbSampler_computeGradientTrilinear(
        sampler, indexCoordinates, time, attributeIndex);
    // Note: xfmNormal takes inverse!
    gradients[i] = xfmNormal(sampler->grid->objectToIndex, gradient);
  }
}

inline void VdbSampler_computeGradient_denseTrilinear(
    const VdbSamplerShared *uniform sampler,
    uniform unsigned int N,
    const vec3f *uniform objectCoordinates,
    const float *uniform times,
    const uniform uint32 attributeIndex,
    vec3f *uniform gradients)
{
  assert(sampler->grid->dense);

  foreach (i = 0 ... N) {
    const vec3f oc               = objectCoordinates[i];
    const vec3f indexCoordinates = xfmPoint(sampler->grid->objectToIndex, oc);
    const float time             = times ? times[i] : 0.f;
    const vec3f gradient         = VdbSampler_computeGradient_denseTrilinear(
        sampler, indexCoordinates, time, attributeIndex);
    // Note: xfmNormal takes inverse!
    gradients[i] = xfmNormal(sampler->grid->objectToIndex, gradient);
  }
}
