/*=========================================================================

  Program:   Visualization Toolkit
  Module:    vtkSampleFunction.cxx

  Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
  All rights reserved.
  See Copyright.txt or http://www.kitware.com/Copyright.htm for details.

     This software is distributed WITHOUT ANY WARRANTY; without even
     the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
     PURPOSE.  See the above copyright notice for more information.

=========================================================================*/
#include "vtkSampleFunction.h"

#include "vtkDoubleArray.h"
#include "vtkFloatArray.h"
#include "vtkGarbageCollector.h"
#include "vtkImageData.h"
#include "vtkImplicitFunction.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMath.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkSMPTools.h"
#include "vtkStreamingDemandDrivenPipeline.h"

vtkStandardNewMacro(vtkSampleFunction);
vtkCxxSetObjectMacro(vtkSampleFunction, ImplicitFunction, vtkImplicitFunction);

// The heart of the algorithm plus interface to the SMP tools.
template <class T>
class vtkSampleFunctionAlgorithm
{
public:
  vtkImplicitFunction* ImplicitFunction;
  T* Scalars;
  float* Normals;
  vtkIdType Extent[6];
  vtkIdType Dims[3];
  vtkIdType SliceSize;
  double Origin[3];
  double Spacing[3];
  double CapValue;

  // Constructor
  vtkSampleFunctionAlgorithm();

  // Interface between VTK and templated functions
  static void SampleAcrossImage(
    vtkSampleFunction* self, vtkImageData* output, int extent[6], T* scalars, float* normals);

  // Cap the boundaries with the specified cap value (if requested).
  void Cap();

  // Interface implicit function computation to SMP tools.
  template <class TT>
  class FunctionValueOp
  {
  public:
    FunctionValueOp(vtkSampleFunctionAlgorithm<TT>* algo) { this->Algo = algo; }
    vtkSampleFunctionAlgorithm* Algo;

    void operator()(vtkIdType k, vtkIdType end)
    {
      double x[3];
      vtkIdType* extent = this->Algo->Extent;
      vtkIdType i, j, jOffset, kOffset;
      for (; k < end; ++k)
      {
        x[2] = this->Algo->Origin[2] + k * this->Algo->Spacing[2];
        kOffset = (k - extent[4]) * this->Algo->SliceSize;
        for (j = extent[2]; j <= extent[3]; ++j)
        {
          x[1] = this->Algo->Origin[1] + j * this->Algo->Spacing[1];
          jOffset = (j - extent[2]) * this->Algo->Dims[0];
          for (i = extent[0]; i <= extent[1]; ++i)
          {
            x[0] = this->Algo->Origin[0] + i * this->Algo->Spacing[0];
            this->Algo->Scalars[(i - extent[0]) + jOffset + kOffset] =
              static_cast<TT>(this->Algo->ImplicitFunction->FunctionValue(x));
          }
        }
      }
    }
  };

  // Interface implicit function graadient computation to SMP tools.
  template <class TT>
  class FunctionGradientOp
  {
  public:
    FunctionGradientOp(vtkSampleFunctionAlgorithm<TT>* algo) { this->Algo = algo; }
    vtkSampleFunctionAlgorithm* Algo;

    void operator()(vtkIdType k, vtkIdType end)
    {
      double x[3], n[3];
      float* nPtr;
      vtkIdType* extent = this->Algo->Extent;
      vtkIdType i, j, jOffset, kOffset;
      for (; k < end; ++k)
      {
        x[2] = this->Algo->Origin[2] + k * this->Algo->Spacing[2];
        kOffset = (k - extent[4]) * this->Algo->SliceSize;
        for (j = extent[2]; j <= extent[3]; ++j)
        {
          x[1] = this->Algo->Origin[1] + j * this->Algo->Spacing[1];
          jOffset = (j - extent[2]) * this->Algo->Dims[0];
          for (i = extent[0]; i <= extent[1]; ++i)
          {
            x[0] = this->Algo->Origin[0] + i * this->Algo->Spacing[0];
            this->Algo->ImplicitFunction->FunctionGradient(x, n);
            vtkMath::Normalize(n);
            nPtr = this->Algo->Normals + 3 * ((i - extent[0]) + jOffset + kOffset);
            nPtr[0] = static_cast<TT>(-n[0]);
            nPtr[1] = static_cast<TT>(-n[1]);
            nPtr[2] = static_cast<TT>(-n[2]);
          } // i
        }   // j
      }     // k
    }
  };
};

//------------------------------------------------------------------------------
// Initialized mainly to eliminate compiler warnings.
template <class T>
vtkSampleFunctionAlgorithm<T>::vtkSampleFunctionAlgorithm()
  : Scalars(nullptr)
  , Normals(nullptr)
{
  for (int i = 0; i < 3; ++i)
  {
    this->Extent[2 * i] = this->Extent[2 * i + 1] = 0;
    this->Dims[i] = 0;
    this->Origin[i] = this->Spacing[i] = 0.0;
  }
  this->SliceSize = 0;
  this->CapValue = 0.0;
}

//------------------------------------------------------------------------------
// Templated class is glue between VTK and templated algorithms.
template <class T>
void vtkSampleFunctionAlgorithm<T>::SampleAcrossImage(
  vtkSampleFunction* self, vtkImageData* output, int extent[6], T* scalars, float* normals)
{
  // Populate data into local storage
  vtkSampleFunctionAlgorithm<T> algo;
  algo.ImplicitFunction = self->GetImplicitFunction();
  algo.Scalars = scalars;
  algo.Normals = normals;
  for (int i = 0; i < 3; ++i)
  {
    algo.Extent[2 * i] = extent[2 * i];
    algo.Extent[2 * i + 1] = extent[2 * i + 1];
    algo.Dims[i] = extent[2 * i + 1] - extent[2 * i] + 1;
  }
  algo.SliceSize = algo.Dims[0] * algo.Dims[1];
  output->GetOrigin(algo.Origin);
  output->GetSpacing(algo.Spacing);
  algo.CapValue = self->GetCapValue();

  // Okay now generate samples using SMP tools
  FunctionValueOp<T> values(&algo);
  vtkSMPTools::For(extent[4], extent[5] + 1, values);

  // If requested, generate normals
  if (algo.Normals)
  {
    FunctionGradientOp<T> gradient(&algo);
    vtkSMPTools::For(extent[4], extent[5] + 1, gradient);
  }

  // If requested, cap boundaries
  if (self->GetCapping())
  {
    algo.Cap();
  }
}

//------------------------------------------------------------------------------
// Cap the boundaries of the volume if requested.
template <class T>
void vtkSampleFunctionAlgorithm<T>::Cap()
{
  vtkIdType i, j, k;
  vtkIdType idx;

  // i-j planes
  // k = this->Extent[4];
  for (j = this->Extent[2]; j <= this->Extent[3]; j++)
  {
    for (i = this->Extent[0]; i <= this->Extent[1]; i++)
    {
      this->Scalars[i + j * this->Dims[0]] = this->CapValue;
    }
  }

  idx = this->Extent[5] * this->SliceSize;
  for (j = this->Extent[2]; j <= this->Extent[3]; j++)
  {
    for (i = this->Extent[0]; i <= this->Extent[1]; i++)
    {
      this->Scalars[idx + i + j * this->Dims[0]] = this->CapValue;
    }
  }

  // j-k planes
  // i = this->Extent[0];
  for (k = this->Extent[4]; k <= this->Extent[5]; k++)
  {
    for (j = this->Extent[2]; j <= this->Extent[3]; j++)
    {
      this->Scalars[j * this->Dims[0] + k * this->SliceSize] = this->CapValue;
    }
  }

  i = this->Extent[1];
  for (k = this->Extent[4]; k <= this->Extent[5]; k++)
  {
    for (j = this->Extent[2]; j <= this->Extent[3]; j++)
    {
      this->Scalars[i + j * this->Dims[0] + k * this->SliceSize] = this->CapValue;
    }
  }

  // i-k planes
  // j = this->Extent[2];
  for (k = this->Extent[4]; k <= this->Extent[5]; k++)
  {
    for (i = this->Extent[0]; i <= this->Extent[1]; i++)
    {
      this->Scalars[i + k * this->SliceSize] = this->CapValue;
    }
  }

  j = this->Extent[3];
  idx = j * this->Dims[0];
  for (k = this->Extent[4]; k <= this->Extent[5]; k++)
  {
    for (i = this->Extent[0]; i <= this->Extent[1]; i++)
    {
      this->Scalars[idx + i + k * this->SliceSize] = this->CapValue;
    }
  }
}

//------------------------------------------------------------------------------
// Okay define the VTK class proper
vtkSampleFunction::vtkSampleFunction()
{
  this->ModelBounds[0] = -1.0;
  this->ModelBounds[1] = 1.0;
  this->ModelBounds[2] = -1.0;
  this->ModelBounds[3] = 1.0;
  this->ModelBounds[4] = -1.0;
  this->ModelBounds[5] = 1.0;

  this->SampleDimensions[0] = 50;
  this->SampleDimensions[1] = 50;
  this->SampleDimensions[2] = 50;

  this->Capping = 0;
  this->CapValue = VTK_DOUBLE_MAX;

  this->ImplicitFunction = nullptr;

  this->ComputeNormals = 1;
  this->OutputScalarType = VTK_DOUBLE;

  this->ScalarArrayName = nullptr;
  this->SetScalarArrayName("scalars");

  this->NormalArrayName = nullptr;
  this->SetNormalArrayName("normals");

  this->SetNumberOfInputPorts(0);
}

//------------------------------------------------------------------------------
vtkSampleFunction::~vtkSampleFunction()
{
  this->SetImplicitFunction(nullptr);
  this->SetScalarArrayName(nullptr);
  this->SetNormalArrayName(nullptr);
}

//------------------------------------------------------------------------------
// Specify the dimensions of the data on which to sample.
void vtkSampleFunction::SetSampleDimensions(int i, int j, int k)
{
  int dim[3];

  dim[0] = i;
  dim[1] = j;
  dim[2] = k;

  this->SetSampleDimensions(dim);
}

//------------------------------------------------------------------------------
// Specify the dimensions of the data on which to sample.
void vtkSampleFunction::SetSampleDimensions(int dim[3])
{
  vtkDebugMacro(<< " setting SampleDimensions to (" << dim[0] << "," << dim[1] << "," << dim[2]
                << ")");

  if (dim[0] != this->SampleDimensions[0] || dim[1] != this->SampleDimensions[1] ||
    dim[2] != this->SampleDimensions[2])
  {
    for (int i = 0; i < 3; i++)
    {
      this->SampleDimensions[i] = (dim[i] > 0 ? dim[i] : 1);
    }
    this->Modified();
  }
}

//------------------------------------------------------------------------------
// Set the bounds of the model.
void vtkSampleFunction::SetModelBounds(const double bounds[6])
{
  this->SetModelBounds(bounds[0], bounds[1], bounds[2], bounds[3], bounds[4], bounds[5]);
}

//------------------------------------------------------------------------------
void vtkSampleFunction::SetModelBounds(
  double xMin, double xMax, double yMin, double yMax, double zMin, double zMax)
{
  vtkDebugMacro(<< " setting ModelBounds to ("
                << "(" << xMin << "," << xMax << "), "
                << "(" << yMin << "," << yMax << "), "
                << "(" << zMin << "," << zMax << "), ");
  if ((xMin > xMax) || (yMin > yMax) || (zMin > zMax))
  {
    vtkErrorMacro("Invalid bounds: "
      << "(" << xMin << "," << xMax << "), "
      << "(" << yMin << "," << yMax << "), "
      << "(" << zMin << "," << zMax << ")"
      << " Bound mins cannot be larger that bound maxs");
    return;
  }
  if (xMin != this->ModelBounds[0] || xMax != this->ModelBounds[1] ||
    yMin != this->ModelBounds[2] || yMax != this->ModelBounds[3] || zMin != this->ModelBounds[4] ||
    zMax != this->ModelBounds[5])
  {
    this->ModelBounds[0] = xMin;
    this->ModelBounds[1] = xMax;
    this->ModelBounds[2] = yMin;
    this->ModelBounds[3] = yMax;
    this->ModelBounds[4] = zMin;
    this->ModelBounds[5] = zMax;
    this->Modified();
  }
}

//------------------------------------------------------------------------------
int vtkSampleFunction::RequestInformation(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** vtkNotUsed(inputVector), vtkInformationVector* outputVector)
{
  // get the info objects
  vtkInformation* outInfo = outputVector->GetInformationObject(0);

  int i;
  double ar[3], origin[3];

  int wExt[6];
  wExt[0] = 0;
  wExt[2] = 0;
  wExt[4] = 0;
  wExt[1] = this->SampleDimensions[0] - 1;
  wExt[3] = this->SampleDimensions[1] - 1;
  wExt[5] = this->SampleDimensions[2] - 1;

  outInfo->Set(vtkStreamingDemandDrivenPipeline::WHOLE_EXTENT(), wExt, 6);

  for (i = 0; i < 3; i++)
  {
    origin[i] = this->ModelBounds[2 * i];
    if (this->SampleDimensions[i] <= 1)
    {
      ar[i] = 1;
    }
    else
    {
      ar[i] =
        (this->ModelBounds[2 * i + 1] - this->ModelBounds[2 * i]) / (this->SampleDimensions[i] - 1);
    }
  }
  outInfo->Set(vtkDataObject::ORIGIN(), origin, 3);
  outInfo->Set(vtkDataObject::SPACING(), ar, 3);

  vtkDataObject::SetPointDataActiveScalarInfo(outInfo, this->OutputScalarType, 1);

  outInfo->Set(vtkAlgorithm::CAN_PRODUCE_SUB_EXTENT(), 1);

  return 1;
}

//------------------------------------------------------------------------------
// Produce the data.
void vtkSampleFunction::ExecuteDataWithInformation(vtkDataObject* outp, vtkInformation* outInfo)
{
  vtkFloatArray* newNormals = nullptr;
  float* normals = nullptr;

  vtkImageData* output = this->GetOutput();
  int* extent = this->GetExecutive()->GetOutputInformation(0)->Get(
    vtkStreamingDemandDrivenPipeline::UPDATE_EXTENT());

  output->SetExtent(extent);
  output = this->AllocateOutputData(outp, outInfo);
  vtkDataArray* newScalars = output->GetPointData()->GetScalars();
  vtkIdType numPts = newScalars->GetNumberOfTuples();

  vtkDebugMacro(<< "Sampling implicit function");

  // Initialize self; create output objects
  //
  if (!this->ImplicitFunction)
  {
    vtkErrorMacro(<< "No implicit function specified");
    return;
  }

  if (this->ComputeNormals)
  {
    newNormals = vtkFloatArray::New();
    newNormals->SetNumberOfComponents(3);
    newNormals->SetNumberOfTuples(numPts);
    normals = newNormals->WritePointer(0, numPts);
  }

  void* ptr = output->GetArrayPointerForExtent(newScalars, extent);
  switch (newScalars->GetDataType())
  {
    vtkTemplateMacro(vtkSampleFunctionAlgorithm<VTK_TT>::SampleAcrossImage(
      this, output, extent, (VTK_TT*)ptr, normals));
  }

  newScalars->SetName(this->ScalarArrayName);

  // Update self
  //
  if (newNormals)
  {
    // For an unknown reason yet, if the following line is not commented out,
    // it will make ImplicitSum, TestBoxFunction and TestDiscreteMarchingCubes
    // to fail.
    newNormals->SetName(this->NormalArrayName);
    output->GetPointData()->SetNormals(newNormals);
    newNormals->Delete();
  }
}

//------------------------------------------------------------------------------
vtkMTimeType vtkSampleFunction::GetMTime()
{
  vtkMTimeType mTime = this->Superclass::GetMTime();
  vtkMTimeType impFuncMTime;

  if (this->ImplicitFunction != nullptr)
  {
    impFuncMTime = this->ImplicitFunction->GetMTime();
    mTime = (impFuncMTime > mTime ? impFuncMTime : mTime);
  }

  return mTime;
}

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

  os << indent << "Sample Dimensions: (" << this->SampleDimensions[0] << ", "
     << this->SampleDimensions[1] << ", " << this->SampleDimensions[2] << ")\n";
  os << indent << "ModelBounds: \n";
  os << indent << "  Xmin,Xmax: (" << this->ModelBounds[0] << ", " << this->ModelBounds[1] << ")\n";
  os << indent << "  Ymin,Ymax: (" << this->ModelBounds[2] << ", " << this->ModelBounds[3] << ")\n";
  os << indent << "  Zmin,Zmax: (" << this->ModelBounds[4] << ", " << this->ModelBounds[5] << ")\n";

  os << indent << "OutputScalarType: " << this->OutputScalarType << "\n";

  if (this->ImplicitFunction)
  {
    os << indent << "Implicit Function: " << this->ImplicitFunction << "\n";
  }
  else
  {
    os << indent << "No Implicit function defined\n";
  }

  os << indent << "Capping: " << (this->Capping ? "On\n" : "Off\n");
  os << indent << "Cap Value: " << this->CapValue << "\n";

  os << indent << "Compute Normals: " << (this->ComputeNormals ? "On\n" : "Off\n");

  os << indent << "ScalarArrayName: ";
  if (this->ScalarArrayName != nullptr)
  {
    os << this->ScalarArrayName << endl;
  }
  else
  {
    os << "(none)" << endl;
  }

  os << indent << "NormalArrayName: ";
  if (this->NormalArrayName != nullptr)
  {
    os << this->NormalArrayName << endl;
  }
  else
  {
    os << "(none)" << endl;
  }
}

//------------------------------------------------------------------------------
void vtkSampleFunction::ReportReferences(vtkGarbageCollector* collector)
{
  this->Superclass::ReportReferences(collector);
  vtkGarbageCollectorReport(collector, this->ImplicitFunction, "ImplicitFunction");
}