// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause
#include "vtkFiniteDifferenceGradientEstimator.h"

#include "vtkArrayDispatch.h"
#include "vtkDataArrayRange.h"
#include "vtkDirectionEncoder.h"
#include "vtkImageData.h"
#include "vtkMultiThreader.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"

#include <cmath>

VTK_ABI_NAMESPACE_BEGIN
vtkStandardNewMacro(vtkFiniteDifferenceGradientEstimator);

// This is the templated function that actually computes the EncodedNormal
// and the GradientMagnitude
struct vtkComputeGradientsFunctor
{
  template <class TArray>
  void operator()(
    TArray* data, vtkFiniteDifferenceGradientEstimator* estimator, int thread_id, int thread_count)
  {
    vtkIdType xstep, ystep, zstep;
    int x, y, z;
    vtkIdType offset;
    int x_start, x_limit;
    int y_start, y_limit;
    int z_start, z_limit;
    int useClip;
    int* clip;
    unsigned char* gptr;
    unsigned short* nptr;
    float n[3], t;
    float gvalue;
    float zeroNormalThreshold;
    int useBounds;
    int bounds[6];
    int size[3];
    float aspect[3];
    int xlow, xhigh;
    float scale, bias;
    int computeGradientMagnitudes;
    vtkDirectionEncoder* direction_encoder;
    int zeroPad;

    estimator->GetInputSize(size);
    estimator->GetInputAspect(aspect);
    computeGradientMagnitudes = estimator->GetComputeGradientMagnitudes();
    scale = estimator->GetGradientMagnitudeScale();
    bias = estimator->GetGradientMagnitudeBias();
    zeroPad = estimator->GetZeroPad();

    // adjust the aspect
    aspect[0] = aspect[0] * 2.0 * estimator->SampleSpacingInVoxels;
    aspect[1] = aspect[1] * 2.0 * estimator->SampleSpacingInVoxels;
    aspect[2] = aspect[2] * 2.0 * estimator->SampleSpacingInVoxels;

    // Compute steps through the volume in x, y, and z
    xstep = 1;
    ystep = size[0];
    zstep = ystep * size[1];

    // Multiply by the spacing used for normal estimation
    xstep *= estimator->SampleSpacingInVoxels;
    ystep *= estimator->SampleSpacingInVoxels;
    zstep *= estimator->SampleSpacingInVoxels;

    // Get the length at or below which normals are considered to
    // be "zero"
    zeroNormalThreshold = estimator->GetZeroNormalThreshold();

    useBounds = estimator->GetBoundsClip();

    // Compute an offset based on the thread_id. The volume will
    // be broken into large slabs (thread_count slabs). For this thread
    // we need to access the correct slab. Also compute the z plane that
    // this slab starts on, and the z limit of this slab (one past the
    // end of the slab)
    if (useBounds)
    {
      estimator->GetBounds(bounds);
      x_start = bounds[0];
      x_limit = bounds[1] + 1;
      y_start = bounds[2];
      y_limit = bounds[3] + 1;
      z_start = static_cast<int>(
                  (thread_id / static_cast<float>(thread_count)) * (bounds[5] - bounds[4] + 1)) +
        bounds[4];
      z_limit = static_cast<int>(((thread_id + 1) / static_cast<float>(thread_count)) *
                  (bounds[5] - bounds[4] + 1)) +
        bounds[4];
    }
    else
    {
      x_start = 0;
      x_limit = size[0];
      y_start = 0;
      y_limit = size[1];
      z_start = static_cast<int>((thread_id / static_cast<float>(thread_count)) * size[2]);
      z_limit = static_cast<int>(((thread_id + 1) / static_cast<float>(thread_count)) * size[2]);
    }

    // Do final error checking on limits - make sure they are all within bounds
    // of the scalar input

    x_start = (x_start < 0) ? (0) : (x_start);
    y_start = (y_start < 0) ? (0) : (y_start);
    z_start = (z_start < 0) ? (0) : (z_start);

    x_limit = (x_limit > size[0]) ? (size[0]) : (x_limit);
    y_limit = (y_limit > size[1]) ? (size[1]) : (y_limit);
    z_limit = (z_limit > size[2]) ? (size[2]) : (z_limit);

    direction_encoder = estimator->GetDirectionEncoder();

    useClip = estimator->GetUseCylinderClip();
    clip = estimator->GetCircleLimits();

    auto data_ptr = vtk::DataArrayValueRange(data).begin();
    // Loop through all the data and compute the encoded normal and
    // gradient magnitude for each scalar location
    for (z = z_start; z < z_limit; z++)
    {
      for (y = y_start; y < y_limit; y++)
      {
        if (useClip)
        {
          xlow = ((clip[2 * y]) > x_start) ? (clip[2 * y]) : (x_start);
          xhigh = ((clip[2 * y + 1] + 1) < x_limit) ? (clip[2 * y + 1] + 1) : (x_limit);
        }
        else
        {
          xlow = x_start;
          xhigh = x_limit;
        }
        offset = z * zstep + y * ystep + xlow;

        // Set some pointers
        auto dptr = data_ptr + offset;
        nptr = estimator->EncodedNormals + offset;
        gptr = estimator->GradientMagnitudes + offset;

        for (x = xlow; x < xhigh; x++)
        {

          // Use a central difference method if possible,
          // otherwise use a forward or backward difference if
          // we are on the edge

          // Compute the X component
          if (x < estimator->SampleSpacingInVoxels)
          {
            if (zeroPad)
            {
              n[0] = -(static_cast<float>(*(dptr + xstep)));
            }
            else
            {
              n[0] = 2.0 * (static_cast<float>(*(dptr)) - static_cast<float>(*(dptr + xstep)));
            }
          }
          else if (x >= size[0] - estimator->SampleSpacingInVoxels)
          {
            if (zeroPad)
            {
              n[0] = static_cast<float>(*(dptr - xstep));
            }
            else
            {
              n[0] = 2.0 * (static_cast<float>(*(dptr - xstep)) - static_cast<float>(*(dptr)));
            }
          }
          else
          {
            n[0] = static_cast<float>(*(dptr - xstep)) - static_cast<float>(*(dptr + xstep));
          }

          // Compute the Y component
          if (y < estimator->SampleSpacingInVoxels)
          {
            if (zeroPad)
            {
              n[1] = -static_cast<float>(*(dptr + ystep));
            }
            else
            {
              n[1] = 2.0 * (static_cast<float>(*(dptr)) - static_cast<float>(*(dptr + ystep)));
            }
          }
          else if (y >= size[1] - estimator->SampleSpacingInVoxels)
          {
            if (zeroPad)
            {
              n[1] = static_cast<float>(*(dptr - ystep));
            }
            else
            {
              n[1] = 2.0 * (static_cast<float>(*(dptr - ystep)) - static_cast<float>(*(dptr)));
            }
          }
          else
          {
            n[1] = static_cast<float>(*(dptr - ystep)) - static_cast<float>(*(dptr + ystep));
          }

          // Compute the Z component
          if (z < estimator->SampleSpacingInVoxels)
          {
            if (zeroPad)
            {
              n[2] = -static_cast<float>(*(dptr + zstep));
            }
            else
            {
              n[2] = 2.0 * (static_cast<float>(*(dptr)) - static_cast<float>(*(dptr + zstep)));
            }
          }
          else if (z >= size[2] - estimator->SampleSpacingInVoxels)
          {
            if (zeroPad)
            {
              n[2] = static_cast<float>(*(dptr - zstep));
            }
            else
            {
              n[2] = 2.0 * (static_cast<float>(*(dptr - zstep)) - static_cast<float>(*(dptr)));
            }
          }
          else
          {
            n[2] = static_cast<float>(*(dptr - zstep)) - static_cast<float>(*(dptr + zstep));
          }

          // Take care of the aspect ratio of the data
          // Scaling in the vtkVolume is isotropic, so this is the
          // only place we have to worry about non-isotropic scaling.
          n[0] /= aspect[0];
          n[1] /= aspect[1];
          n[2] /= aspect[2];

          // Compute the gradient magnitude
          t = sqrt(static_cast<double>(n[0] * n[0] + n[1] * n[1] + n[2] * n[2]));

          if (computeGradientMagnitudes)
          {
            // Encode this into an 8 bit value
            gvalue = (t + bias) * scale;

            if (gvalue < 0.0)
            {
              *gptr = 0;
            }
            else if (gvalue > 255.0)
            {
              *gptr = 255;
            }
            else
            {
              *gptr = static_cast<unsigned char>(gvalue);
            }
            gptr++;
          }

          // Normalize the gradient direction
          if (t > zeroNormalThreshold)
          {
            n[0] /= t;
            n[1] /= t;
            n[2] /= t;
          }
          else
          {
            n[0] = n[1] = n[2] = 0.0;
          }

          // Convert the gradient direction into an encoded index value
          *nptr = direction_encoder->GetEncodedDirection(n);
          nptr++;
          dptr++;
        }
      }
    }
  }
};

// Construct a vtkFiniteDifferenceGradientEstimator
vtkFiniteDifferenceGradientEstimator::vtkFiniteDifferenceGradientEstimator()
{
  this->SampleSpacingInVoxels = 1;
}

// Destruct a vtkFiniteDifferenceGradientEstimator - free up any memory used
vtkFiniteDifferenceGradientEstimator::~vtkFiniteDifferenceGradientEstimator() = default;

static VTK_THREAD_RETURN_TYPE vtkSwitchOnDataType(void* arg)
{
  vtkFiniteDifferenceGradientEstimator* estimator;
  int thread_count;
  int thread_id;
  vtkDataArray* scalars;

  thread_id = ((vtkMultiThreader::ThreadInfo*)(arg))->ThreadID;
  thread_count = ((vtkMultiThreader::ThreadInfo*)(arg))->NumberOfThreads;
  estimator =
    (vtkFiniteDifferenceGradientEstimator*)(((vtkMultiThreader::ThreadInfo*)(arg))->UserData);
  scalars = estimator->InputData->GetPointData()->GetScalars();

  if (scalars == nullptr)
  {
    return VTK_THREAD_RETURN_VALUE;
  }

  // Find the data type of the Input and call the correct
  // templated function to actually compute the normals and magnitudes
  vtkComputeGradientsFunctor functor;
  if (!vtkArrayDispatch::Dispatch::Execute(scalars, functor, estimator, thread_id, thread_count))
  {
    functor(scalars, estimator, thread_id, thread_count);
  }

  return VTK_THREAD_RETURN_VALUE;
}

// This method is used to compute the encoded normal and the
// magnitude of the gradient for each voxel location in the
// Input.
void vtkFiniteDifferenceGradientEstimator::UpdateNormals()
{
  vtkDebugMacro(<< "Updating Normals!");
  this->Threader->SetNumberOfThreads(this->NumberOfThreads);

  this->Threader->SetSingleMethod(vtkSwitchOnDataType, this);

  this->Threader->SingleMethodExecute();
}

// Print the vtkFiniteDifferenceGradientEstimator
void vtkFiniteDifferenceGradientEstimator::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);

  os << indent << "Sample spacing in voxels: " << this->SampleSpacingInVoxels << endl;
}
VTK_ABI_NAMESPACE_END
