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

#pragma once

#include "openvkl/vdb.h"

// Note: Changing this value will probably break our API contract.
// Tread carefully.
#define DDA_STATE_MAX_NUM_LEVELS VKL_VDB_NUM_LEVELS

// Helpers for the offsets in three-dimensional arrays below.
#define DDA_STATE_X_OFFSET(LEVEL) (LEVEL)
#define DDA_STATE_Y_OFFSET(LEVEL) (DDA_STATE_MAX_NUM_LEVELS + (LEVEL))
#define DDA_STATE_Z_OFFSET(LEVEL) (2 * DDA_STATE_MAX_NUM_LEVELS + (LEVEL))

struct DdaState
{
  // The level at which we iterate. This is the same for all lanes
  // as it only depends on the sampler configuration passed into Open VKL.
  uniform uint32 level;

  // The current ray parameter.
  float t;

  // The ray. We need this to reinitialize lower levels when we skip them.
  vec3f rayOrg;
  vec3f rayDir;

  // The current index. This is in units of leaf level voxels for all levels.
  // NOTE: See DDA_STATE_*_OFFSET macros above for the storage layout.
  int32 idx[DDA_STATE_MAX_NUM_LEVELS * 3];
  float tNext[DDA_STATE_MAX_NUM_LEVELS * 3];
  float tDelta[DDA_STATE_MAX_NUM_LEVELS * 3];
  int32 idxDelta[DDA_STATE_MAX_NUM_LEVELS * 3];
};

// -----------------------------------------------------------------------------
// Utilities.
// -----------------------------------------------------------------------------

/*
 * Returns -1 for negative v, 0 for 0, and 1 for positive v.
 */
inline int dir_safe_sign(float v)
{
  return ((int)(0 < v)) - ((int)(v < 0));
}

inline vec3i dir_safe_sign(const vec3f &v)
{
  return make_vec3i(dir_safe_sign(v.x), dir_safe_sign(v.y), dir_safe_sign(v.z));
}

/*
 * Compute 1/v, returning infinity for 0.
 */
inline float dir_safe_rcp(float v)
{
  return (v == 0) ? inf : divide_safe(v);
}

inline vec3f dir_safe_rcp(const vec3f &v)
{
  return make_vec3f(dir_safe_rcp(v.x), dir_safe_rcp(v.y), dir_safe_rcp(v.z));
}

/*
 * Intersect a set of three axis-aligned hyperplanes.
 */
inline vec3f intersect_planes(const vec3f &rayOrg,
                              const vec3f &rayInvDir,
                              const vec3f &planes)
{
  const float tx = (planes.x - rayOrg.x) * rayInvDir.x;
  const float ty = (planes.y - rayOrg.y) * rayInvDir.y;
  const float tz = (planes.z - rayOrg.z) * rayInvDir.z;

  // Do not intersect planes the origin is on!
#if defined(ISPC)
  return make_vec3f(select((planes.x == rayOrg.x || isnan(tx)), inf, tx),
                    select((planes.y == rayOrg.y || isnan(ty)), inf, ty),
                    select((planes.z == rayOrg.z || isnan(tz)), inf, tz));
#else
  return make_vec3f((planes.x == rayOrg.x || isnan(tx)) ? inf : tx,
                    (planes.y == rayOrg.y || isnan(ty)) ? inf : ty,
                    (planes.z == rayOrg.z || isnan(tz)) ? inf : tz);

#endif
}

/*
 * Compare __vkl_vdb_map_domain_offset_to_voxel_uniform_*
 * in openvkl/vdb_topology.h.in.
 */
inline void clampOffsetToCell(
    int32 &x,
    int32 &y,
    int32 &z,
    uniform int32 resolution)  // ... of a single cell on the current level.
{
  // We may map a voxel coordinate to the origin of a voxel with resolution
  // logVoxelRes using this simple mask because resolutions are powers of two.
#if defined(ISPC)
  assert(popcnt(resolution) == 1);
#endif
  const uniform int mask = ~((resolution)-1);
  x                      = x & mask;
  y                      = y & mask;
  z                      = z & mask;
}

inline void clampOffsetToCellVarying(
    int32 &x,
    int32 &y,
    int32 &z,
    int32 resolution)  // ... of a single cell on the current level.
{
  // We may map a voxel coordinate to the origin of a voxel with resolution
  // logVoxelRes using this simple mask because resolutions are powers of two.
#if defined(ISPC)
  assert(popcnt(resolution) == 1);
#endif
  const int mask = ~((resolution)-1);
  x              = x & mask;
  y              = y & mask;
  z              = z & mask;
}

inline vec3i clampToCell(
    const vec3f &foffset,
    uniform int32 resolution)  // ... of a single cell on the current level.
{
  // Offsets are non-negative, but numerical errors might cause problems.
  vec3i offset = make_vec3i(((int)floor(max(foffset.x, 0.f))),
                            ((int)floor(max(foffset.y, 0.f))),
                            ((int)floor(max(foffset.z, 0.f))));
  clampOffsetToCell(offset.x, offset.y, offset.z, resolution);
  return offset;
}

/*
 * Initialize the dda state.
 */
inline void ddaInit(const vec3f &rayOrg,  // ray in index space.
                    const vec3f &rayDir,
                    const float &tMin,
                    uniform uint32 level,
                    DdaState &dda)
{
  const vec3f oo   = rayOrg + tMin * rayDir;
  const vec3i sgn  = dir_safe_sign(rayDir);
  const vec3f iDir = dir_safe_rcp(rayDir);
  const vec3i sgnNonNegative =
      make_vec3i(((int)(sgn.x >= 0)), ((int)(sgn.y >= 0)), ((int)(sgn.z >= 0)));

#if defined(ISPC)
  dda.level = clamp(level, 0u, ((uniform uint32)DDA_STATE_MAX_NUM_LEVELS - 1));
#else
  dda.level = clamp((int)level, 0, ((int)DDA_STATE_MAX_NUM_LEVELS - 1));
#endif
  dda.t      = tMin;
  dda.rayOrg = rayOrg;
  dda.rayDir = rayDir;

  for (uniform uint32 l = 0; l <= dda.level; ++l) {
    const uniform int32 cellRes = vklVdbLevelRes(l + 1);
    // Clamp the origin to a cell on this level, i.e., a multiple of the
    // cell resolution. This results in the index of the cell we are in.
    const vec3i idx = clampToCell(oo, cellRes);

    // Intersect the exiting plane to get the initial tNext for this level.
    const vec3f planes =
        make_vec3f(((float)(idx.x + cellRes * sgnNonNegative.x)),
                   ((float)(idx.y + cellRes * sgnNonNegative.y)),
                   ((float)(idx.z + cellRes * sgnNonNegative.z)));

    const vec3f tn = intersect_planes(rayOrg, iDir, planes);

    const uniform int32 ix = DDA_STATE_X_OFFSET(l);
    const uniform int32 iy = DDA_STATE_Y_OFFSET(l);
    const uniform int32 iz = DDA_STATE_Z_OFFSET(l);
    dda.idx[ix]            = idx.x;
    dda.idx[iy]            = idx.y;
    dda.idx[iz]            = idx.z;
    dda.idxDelta[ix]       = sgn.x * cellRes;
    dda.idxDelta[iy]       = sgn.y * cellRes;
    dda.idxDelta[iz]       = sgn.z * cellRes;
    dda.tDelta[ix]         = ((float)cellRes) * abs(iDir.x);
    dda.tDelta[iy]         = ((float)cellRes) * abs(iDir.y);
    dda.tDelta[iz]         = ((float)cellRes) * abs(iDir.z);
    dda.tNext[ix]          = tn.x;
    dda.tNext[iy]          = tn.y;
    dda.tNext[iz]          = tn.z;
  }
}

/*
 * Reinitialize dda.idx and dda.tNext for the given level and index space
 * position.
 *
 * All other members must be set (including t, idx, ray members).
 */
inline void ddaReinitBelow(DdaState &dda,
                           uniform uint32 level,
                           int32 dx, // The parent level moved by this many voxels.
                           int32 dy, // Only one of dx, dy, dz is nonzero!
                           int32 dz)
{
  // Relatively cheap to recompute, so we avoid storing these in DdaState as
  // that impacts memory for all ISPC side iterators.
  const vec3f iDir = dir_safe_rcp(dda.rayDir);
  const vec3i sgn = dir_safe_sign(dda.rayDir);
  const vec3i sgnNonNegative =
      make_vec3i(((int)(sgn.x >= 0)), ((int)(sgn.y >= 0)), ((int)(sgn.z >= 0)));

  // We know that we are on a hyperplane dividing two cells on level.
  // Move the index point off of that plane, or clamping below will not be
  // numerically robust!
  // Note that we know the cell has size 1 in each dimension, so the move
  // by 1e-2f seems perfectly safe.
  // There is no point in moving oo in other dimensions.
  vec3f oo = dda.rayOrg + dda.t * dda.rayDir;
  if (dx != 0)
  {
    assert(dy == 0);
    assert(dz == 0);
    oo.x = oo.x + ((float)sgn.x) * 1e-2f;
  }
  else if (dy != 0)
  {
    assert(dx == 0);
    assert(dz == 0);
    oo.y = oo.y + ((float)sgn.y) * 1e-2f;
  }
  else
  {
    assert(dx == 0);
    assert(dy == 0);
    assert(dz != 0);
    oo.z = oo.z + ((float)sgn.z) * 1e-2f;
  }

  for (uniform uint32 l = level+1; l <= dda.level; ++l)
  {
    const uniform int32 cellRes = vklVdbLevelRes(l + 1);
    const uniform uint32 ix = DDA_STATE_X_OFFSET(l);
    const uniform uint32 iy = DDA_STATE_Y_OFFSET(l);
    const uniform uint32 iz = DDA_STATE_Z_OFFSET(l);
    const vec3i idx = clampToCell(oo, cellRes);

    dda.idx[ix] = idx.x;
    dda.idx[iy] = idx.y;
    dda.idx[iz] = idx.z;

    const vec3f planes =
        make_vec3f(((float)(dda.idx[ix] + cellRes * sgnNonNegative.x)),
                   ((float)(dda.idx[iy] + cellRes * sgnNonNegative.y)),
                   ((float)(dda.idx[iz] + cellRes * sgnNonNegative.z)));
    const vec3f tn = intersect_planes(dda.rayOrg, iDir, planes);

    dda.tNext[ix] = tn.x;
    dda.tNext[iy] = tn.y;
    dda.tNext[iz] = tn.z;
  }
}

/*
 * Step once on the given level, ignoring all other levels.
 */
inline void ddaStep(DdaState &dda, uniform uint32 level)
{
  const uniform uint32 ox = DDA_STATE_X_OFFSET(level);
  const uniform uint32 oy = DDA_STATE_Y_OFFSET(level);
  const uniform uint32 oz = DDA_STATE_Z_OFFSET(level);

  // Enforce gathers on tNext now as we use these values a number of
  // times below.
  const float tNextX = dda.tNext[ox];
  const float tNextY = dda.tNext[oy];
  const float tNextZ = dda.tNext[oz];
  const float minT = min(min(tNextX, tNextY), tNextZ);

  if (minT == tNextX)
  {
    assert(dda.tDelta[ox] != 0);
    assert(dda.idxDelta[ox] != 0);
    dda.t = tNextX;
    dda.tNext[ox] = tNextX + dda.tDelta[ox];
    dda.idx[ox]   = dda.idx[ox] + dda.idxDelta[ox];
  }
  else if (minT == tNextY)
  {
    assert(dda.tDelta[oy] != 0);
    assert(dda.idxDelta[oy] != 0);
    dda.t = tNextY;
    dda.tNext[oy] = tNextY + dda.tDelta[oy];
    dda.idx[oy]   = dda.idx[oy] + dda.idxDelta[oy];
  }
  else
  {
    assert(minT == tNextZ);
    assert(dda.tDelta[oz] != 0);
    assert(dda.idxDelta[oz] != 0);
    dda.t = tNextZ;
    dda.tNext[oz] = tNextZ + dda.tDelta[oz];
    dda.idx[oz]   = dda.idx[oz] + dda.idxDelta[oz];
  }
}
