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

#pragma once

#include "../common/Hit.ih"
#include "../common/Interval.ih"
#include "../sampler/Sampler.ih"
#include "openvkl/iterator.isph"
#include "rkcommon/math/box.ih"
#include "rkcommon/math/math.ih"

/*
 * Intersect isosurfaces along the given ray using bisection method.
 */

static const uniform float BISECT_T_TOL_ABS           = 1e-6f;
static const uniform float BISECT_T_TOL_FACTOR        = 0.01f;
static const uniform float BISECT_VALUE_TOL           = 1e-6f;
static const uniform float BISECT_MAX_BACKTRACK_ITERS = 10;

#define template_intersectSurfacesBisection(univary)                            \
  inline univary float bisect(const SamplerShared *uniform sampler,             \
                              const univary vec3f &origin,                      \
                              const univary vec3f &direction,                   \
                              const uniform uint32 attributeIndex,              \
                              const univary float &time,                        \
                              univary float t0,                                 \
                              univary float sample0,                            \
                              univary float t,                                  \
                              univary float sample,                             \
                              uniform float isovalue,                           \
                              univary float tol,                                \
                              univary float &error)                             \
  {                                                                             \
    const uniform int maxIter = 10;                                             \
                                                                                \
    uniform int iter = 0;                                                       \
    univary float tMid;                                                         \
                                                                                \
    while (iter < maxIter) {                                                    \
      tMid = 0.5f * (t0 + t);                                                   \
                                                                                \
      error = 0.5f * (t - t0);                                                  \
                                                                                \
      if (error < tol) {                                                        \
        break;                                                                  \
      }                                                                         \
                                                                                \
      univary float sampleMid = sampler->computeSample_##univary(               \
          sampler, origin + tMid * direction, attributeIndex, time);            \
                                                                                \
      /* sampling at boundaries between unstructured cells can rarely lead to   \
      NaN values (indicating outside of cell) due to numerical issues; in this  \
      case we know we have already bracketed an isovalue, so return the         \
      nearest result */                                                         \
      if (isnan(sampleMid)) {                                                   \
        if (abs(isovalue - sample0) < abs(isovalue - sample)) {                 \
          return t0;                                                            \
        } else {                                                                \
          return t;                                                             \
        }                                                                       \
      }                                                                         \
                                                                                \
      if (sampleMid == isovalue ||                                              \
          (isovalue - sample0) * (isovalue - sampleMid) < 0.f) {                \
        /* crossing in (t0, tMid) */                                            \
        t      = tMid;                                                          \
        sample = sampleMid;                                                     \
      } else if (sample == isovalue ||                                          \
                 (isovalue - sample) * (isovalue - sampleMid) < 0.f) {          \
        /* crossing in (tMid, t) */                                             \
        t0      = tMid;                                                         \
        sample0 = sampleMid;                                                    \
      } else {                                                                  \
        /* should never get here */                                             \
        print("bisect error\n");                                                \
        return inf;                                                             \
      }                                                                         \
                                                                                \
      iter++;                                                                   \
    }                                                                           \
                                                                                \
    return tMid;                                                                \
  }                                                                             \
                                                                                \
  inline univary bool intersectSurfacesBisection(                               \
      const SamplerShared *uniform sampler,                                     \
      const univary vec3f &origin,                                              \
      const univary vec3f &direction,                                           \
      const univary box1f &tRange,                                              \
      const uniform uint32 attributeIndex,                                      \
      const univary float &time,                                                \
      const univary float step,                                                 \
      const uniform int numValues,                                              \
      const float *uniform values,                                              \
      univary Hit &hit)                                                         \
  {                                                                             \
    assert(tRange.lower < tRange.upper);                                        \
                                                                                \
    univary float t0      = tRange.lower;                                       \
    univary float sample0 = sampler->computeSample_##univary(                   \
        sampler, origin + t0 * direction, attributeIndex, time);                \
                                                                                \
    {                                                                           \
      univary int iters = 0;                                                    \
      while (iters < BISECT_MAX_BACKTRACK_ITERS && isnan(sample0) &&            \
             t0 < tRange.upper) {                                               \
        t0 += max(1e-5f, 1e-5f * step);                                         \
        sample0 = sampler->computeSample_##univary(                             \
            sampler, origin + t0 * direction, attributeIndex, time);            \
        iters++;                                                                \
      }                                                                         \
    }                                                                           \
                                                                                \
    univary float t;                                                            \
                                                                                \
    while (t0 < tRange.upper) {                                                 \
      const univary float h = min(step, tRange.upper - t0);                     \
      t                     = t0 + h;                                           \
      univary float sample  = sampler->computeSample_##univary(                 \
          sampler, origin + t * direction, attributeIndex, time);               \
      univary float ts = t;                                                     \
                                                                                \
      {                                                                         \
        univary int iters = 0;                                                  \
        while (iters < BISECT_MAX_BACKTRACK_ITERS && isnan(sample) &&           \
               ts > t0) {                                                       \
          ts -= max(1e-5f, 1e-5f * step);                                       \
          sample = sampler->computeSample_##univary(                            \
              sampler, origin + ts * direction, attributeIndex, time);          \
          iters++;                                                              \
        }                                                                       \
      }                                                                         \
                                                                                \
      univary float tHit    = inf;                                              \
      univary float epsilon = inf;                                              \
      univary float value   = inf;                                              \
                                                                                \
      for (uniform int i = 0; i < numValues; i++) {                             \
        /* hit at bracket entrance */                                           \
        if (abs(sample0 - values[i]) < BISECT_VALUE_TOL) {                      \
          if (t0 < tHit && t0 <= tRange.upper) {                                \
            tHit    = t0;                                                       \
            value   = values[i];                                                \
            epsilon = 0.125f * step;                                            \
          }                                                                     \
        } /* hit inside bracket */                                              \
        else if (!isnan(sample0 + sample) &&                                    \
                 (values[i] - sample0) * (values[i] - sample) < 0.f) {          \
          /* we have bracketed a crossing; bisect */                            \
          univary float error;                                                  \
          univary float tIso =                                                  \
              bisect(sampler,                                                   \
                     origin,                                                    \
                     direction,                                                 \
                     attributeIndex,                                            \
                     time,                                                      \
                     t0,                                                        \
                     sample0,                                                   \
                     ts,                                                        \
                     sample,                                                    \
                     values[i],                                                 \
                     min(BISECT_T_TOL_ABS, BISECT_T_TOL_FACTOR * h),            \
                     error);                                                    \
                                                                                \
          if (tIso < tHit && tIso <= tRange.upper) {                            \
            tHit    = tIso;                                                     \
            value   = values[i];                                                \
            epsilon = 0.125f * step;                                            \
          }                                                                     \
        }                                                                       \
                                                                                \
        /* check hit at tRange upper limit, after any bisection has occurred    \
         */                                                                     \
        if (t == tRange.upper && abs(sample - values[i]) < BISECT_VALUE_TOL) {  \
          if (ts < tHit && ts <= tRange.upper) {                                \
            tHit    = ts;                                                       \
            value   = values[i];                                                \
            epsilon = 0.125f * step;                                            \
          }                                                                     \
        }                                                                       \
      }                                                                         \
                                                                                \
      if (tHit < inf) {                                                         \
        hit.t       = tHit;                                                     \
        hit.sample  = value;                                                    \
        hit.epsilon = epsilon * length(direction); /* in object space */        \
        return true;                                                            \
      }                                                                         \
                                                                                \
      t0      = t;                                                              \
      sample0 = sample;                                                         \
    }                                                                           \
                                                                                \
    return false;                                                               \
  }

template_intersectSurfacesBisection(uniform);
template_intersectSurfacesBisection(varying);
#undef template_intersectSurfacesBisection
