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

#include "vtkArrayDispatch.h"
#include "vtkArrayDispatchDataSetArrayList.h"
#include "vtkCellData.h"
#include "vtkDataArrayRange.h"
#include "vtkGenericCell.h"
#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkObjectFactory.h"
#include "vtkPixel.h"
#include "vtkPointData.h"
#include "vtkPolyData.h"
#include "vtkSMPThreadLocalObject.h"
#include "vtkSMPTools.h"
#include "vtkStreamingDemandDrivenPipeline.h"

VTK_ABI_NAMESPACE_BEGIN
vtkStandardNewMacro(vtkFitToHeightMapFilter);

namespace
{
//------------------------------------------------------------------------------
// The threaded core of the algorithm for projecting points.
template <typename TPointArray, typename TScalarArray>
struct FitPointsFunctor
{
  TPointArray* InPoints;
  TScalarArray* Scalars;
  TPointArray* OutPoints;
  double Dims[3];
  double Origin[3];
  double H[3];
  vtkFitToHeightMapFilter* Filter;

  FitPointsFunctor(TPointArray* inPts, TScalarArray* s, TPointArray* outPts, int dims[3],
    double o[3], double h[3], vtkFitToHeightMapFilter* filter)
    : InPoints(inPts)
    , Scalars(s)
    , OutPoints(outPts)
    , Filter(filter)
  {
    for (int i = 0; i < 3; ++i)
    {
      this->Dims[i] = static_cast<double>(dims[i]);
      this->Origin[i] = o[i];
      this->H[i] = h[i];
    }
  }

  void Initialize() {}

  void operator()(vtkIdType ptId, vtkIdType endPtId)
  {
    auto xi = vtk::DataArrayValueRange<3>(this->InPoints, 3 * ptId).begin();
    auto xo = vtk::DataArrayValueRange<3>(this->OutPoints, 3 * ptId).begin();
    const double* d = this->Dims;
    const double* o = this->Origin;
    const double* h = this->H;
    int ii, jj;
    double i, j, z, pc[2], ix, iy, w[4];
    auto scalars = vtk::DataArrayValueRange(this->Scalars).begin();
    int s0, s1, s2, s3;
    bool isFirst = vtkSMPTools::GetSingleThread();

    for (; ptId < endPtId; ++ptId, xi += 3, xo += 3)
    {
      if (isFirst)
      {
        this->Filter->CheckAbort();
      }
      if (this->Filter->GetAbortOutput())
      {
        break;
      }
      // Location in image
      i = (xi[0] - o[0]) / h[0];
      j = (xi[1] - o[1]) / h[1];

      // Clamp to image. i,j is integral index into image pixels. It has to be done carefully
      // to manage the parametric coordinates correctly.
      if (i < 0.0)
      {
        ix = 0.0;
        pc[0] = 0.0;
      }
      else if (i >= (d[0] - 1))
      {
        ix = d[0] - 2;
        pc[0] = 1.0;
      }
      else
      {
        pc[0] = modf(i, &ix);
      }

      if (j < 0.0)
      {
        iy = 0.0;
        pc[1] = 0.0;
      }
      else if (j >= (d[1] - 1))
      {
        iy = d[1] - 2;
        pc[1] = 1.0;
      }
      else
      {
        pc[1] = modf(j, &iy);
      }

      // Parametric coordinates for interpolation
      vtkPixel::InterpolationFunctions(pc, w);
      ii = static_cast<int>(ix);
      jj = static_cast<int>(iy);

      // Interpolate height from surrounding data values
      s0 = ii + jj * d[0];
      s1 = s0 + 1;
      s2 = s0 + d[0];
      s3 = s2 + 1;
      z = w[0] * scalars[s0] + w[1] * scalars[s1] + w[2] * scalars[s2] + w[3] * scalars[s3];

      // Set the output point coordinates with a new z-value.
      xo[0] = xi[0];
      xo[1] = xi[1];
      xo[2] = z;
    }
  }

  void Reduce() {}
}; // FitPoints

struct FitPointsWorker
{
  template <typename TPointArray, typename TScalarArray>
  void operator()(TPointArray* inPts, TScalarArray* scalars, vtkDataArray* outPts, int dims[3],
    double origin[3], double h[3], vtkFitToHeightMapFilter* filter)
  {
    FitPointsFunctor<TPointArray, TScalarArray> fit(
      inPts, scalars, TPointArray::FastDownCast(outPts), dims, origin, h, filter);
    vtkSMPTools::For(0, inPts->GetNumberOfTuples(), fit);
  }
};

//------------------------------------------------------------------------------
// The threaded core of the algorithm when projecting cells.
template <typename TScalarArray>
struct FitCellsFunctor
{
  TScalarArray* Scalars;
  int Strategy;
  vtkPolyData* Mesh;
  double* CellHeights;
  double Dims[3];
  double Origin[3];
  double H[3];

  // Avoid repeated allocation
  vtkSMPThreadLocalObject<vtkGenericCell> Cell;
  vtkSMPThreadLocalObject<vtkIdList> Prims;
  vtkSMPThreadLocalObject<vtkPoints> PrimPts;
  vtkFitToHeightMapFilter* Filter;

  FitCellsFunctor(TScalarArray* s, int strat, vtkPolyData* mesh, double* cellHts, int dims[3],
    double o[3], double h[3], vtkFitToHeightMapFilter* filter)
    : Scalars(s)
    , Strategy(strat)
    , Mesh(mesh)
    , CellHeights(cellHts)
    , Filter(filter)
  {
    for (int i = 0; i < 3; ++i)
    {
      this->Dims[i] = static_cast<double>(dims[i]);
      this->Origin[i] = o[i];
      this->H[i] = h[i];
    }
  }

  void Initialize()
  {
    vtkGenericCell*& cell = this->Cell.Local();
    cell->PointIds->Allocate(128);
    cell->Points->Allocate(128);

    vtkIdList*& prims = this->Prims.Local();
    prims->Allocate(128); // allocate some memory

    vtkPoints*& primPts = this->PrimPts.Local();
    primPts->Allocate(128); // allocate some memory
  }

  void operator()(vtkIdType cellId, vtkIdType endCellId)
  {
    double* cellHts = this->CellHeights + cellId;
    auto scalars = vtk::DataArrayValueRange(this->Scalars).begin();
    vtkGenericCell*& cell = this->Cell.Local();
    vtkIdList*& prims = this->Prims.Local();
    vtkPoints*& primPts = this->PrimPts.Local();
    const double* d = this->Dims;
    const double* o = this->Origin;
    const double* h = this->H;
    vtkIdType p, pi, numPrims;
    int cellDim, ii, jj;
    double i, j, z, pc[2], ix, iy, w[4];
    double x0[3], center[2];
    int s0, s1, s2, s3;
    double sum, min, max;
    bool isFirst = vtkSMPTools::GetSingleThread();

    // Process all cells of different types and dimensions
    for (; cellId < endCellId; ++cellId)
    {
      if (isFirst)
      {
        this->Filter->CheckAbort();
      }
      if (this->Filter->GetAbortOutput())
      {
        break;
      }
      this->Mesh->GetCell(cellId, cell);
      cellDim = cell->GetCellDimension();

      cell->Triangulate(0, prims, primPts);
      numPrims = prims->GetNumberOfIds() / (cellDim + 1);

      // Loop over each primitive from the tessellation
      min = VTK_FLOAT_MAX;
      max = VTK_FLOAT_MIN;
      sum = 0.0;
      for (p = 0; p < numPrims; p++)
      {
        // Compute center of primitive
        center[0] = center[1] = 0.0;
        for (pi = 0; pi <= cellDim; ++pi)
        {
          primPts->GetPoint((cellDim + 1) * p + pi, x0);
          center[0] += x0[0];
          center[1] += x0[1];
        }
        center[0] /= static_cast<double>(cellDim + 1);
        center[1] /= static_cast<double>(cellDim + 1);

        // Location in image
        i = (center[0] - o[0]) / h[0];
        j = (center[1] - o[1]) / h[1];

        // Clamp to image. i,j is integral index into image pixels. It has to be done carefully
        // to manage the parametric coordinates correctly.
        if (i < 0.0)
        {
          ix = 0.0;
          pc[0] = 0.0;
        }
        else if (i >= (d[0] - 1))
        {
          ix = d[0] - 2;
          pc[0] = 1.0;
        }
        else
        {
          pc[0] = modf(i, &ix);
        }

        if (j < 0.0)
        {
          iy = 0.0;
          pc[1] = 0.0;
        }
        else if (j >= (d[1] - 1))
        {
          iy = d[1] - 2;
          pc[1] = 1.0;
        }
        else
        {
          pc[1] = modf(j, &iy);
        }

        // Parametric coordinates for interpolation
        vtkPixel::InterpolationFunctions(pc, w);
        ii = static_cast<int>(ix);
        jj = static_cast<int>(iy);

        // Interpolate height from surrounding data values
        s0 = ii + jj * d[0];
        s1 = s0 + 1;
        s2 = s0 + d[0];
        s3 = s2 + 1;
        z = w[0] * scalars[s0] + w[1] * scalars[s1] + w[2] * scalars[s2] + w[3] * scalars[s3];

        // Gather information for final determination of height z
        min = std::min(z, min);
        max = std::max(z, max);
        sum += z; // to compute average
      }           // for all tessellated primitives

      // Now set the cell height
      if (this->Strategy == vtkFitToHeightMapFilter::CELL_AVERAGE_HEIGHT)
      {
        z = fabs(sum / static_cast<double>(numPrims));
      }
      else if (this->Strategy == vtkFitToHeightMapFilter::CELL_MINIMUM_HEIGHT)
      {
        z = min;
      }
      else // if ( this->Strategy == vtkFitToHeightMapFilter::CELL_MAXIMUM_HEIGHT )
      {
        z = max;
      }

      *cellHts++ = z;
    }
  }

  void Reduce() {}
}; // FitCells

struct FitCellsWorker
{
  template <typename TScalarArray>
  void operator()(TScalarArray* scalars, int strategy, vtkPolyData* mesh, double* cellHts,
    int dims[3], double origin[3], double h[3], vtkFitToHeightMapFilter* filter)
  {
    FitCellsFunctor<TScalarArray> functor(
      scalars, strategy, mesh, cellHts, dims, origin, h, filter);
    vtkSMPTools::For(0, mesh->GetNumberOfCells(), functor);
  }
};
} // anonymous namespace

//------------------------------------------------------------------------------
// Construct object. Two inputs are mandatory.
vtkFitToHeightMapFilter::vtkFitToHeightMapFilter()
{
  this->SetNumberOfInputPorts(2);

  this->FittingStrategy = vtkFitToHeightMapFilter::POINT_PROJECTION;
  this->UseHeightMapOffset = true;
  this->Offset = 0.0;
}

//------------------------------------------------------------------------------
vtkFitToHeightMapFilter::~vtkFitToHeightMapFilter() = default;

//------------------------------------------------------------------------------
int vtkFitToHeightMapFilter::RequestData(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);
  vtkInformation* in2Info = inputVector[1]->GetInformationObject(0);
  vtkInformation* outInfo = outputVector->GetInformationObject(0);

  vtkDebugMacro(<< "Executing fit to height map");

  // get the two inputs and output
  vtkPolyData* input = vtkPolyData::SafeDownCast(inInfo->Get(vtkDataObject::DATA_OBJECT()));
  vtkImageData* image = vtkImageData::SafeDownCast(in2Info->Get(vtkDataObject::DATA_OBJECT()));
  vtkPolyData* output = vtkPolyData::SafeDownCast(outInfo->Get(vtkDataObject::DATA_OBJECT()));

  // Check input
  if (input == nullptr || image == nullptr || output == nullptr)
  {
    vtkErrorMacro(<< "Bad input");
    return 1;
  }
  vtkIdType numPts = input->GetNumberOfPoints();
  vtkIdType numCells = input->GetNumberOfCells();

  if (numPts < 1 || numCells < 1)
  {
    vtkDebugMacro(<< "Empty input");
    return 1;
  }

  // Only process real-type points
  vtkPoints* inPts = input->GetPoints();

  // Looking for a xy-oriented image
  int dims[3];
  double origin[3], h[3];
  image->GetDimensions(dims);
  image->GetOrigin(origin);
  image->GetSpacing(h);

  if (dims[0] < 2 || dims[1] < 2 || dims[2] != 1)
  {
    vtkErrorMacro(<< "Filter operates on 2D x-y images");
    return 1;
  }

  // Finally warn if the image data does not fully contain the
  // input polydata.
  double inputBds[6], imageBds[6];
  input->GetBounds(inputBds);
  image->GetBounds(imageBds);
  if (inputBds[0] < imageBds[0] || inputBds[1] > imageBds[1] || inputBds[2] < imageBds[2] ||
    inputBds[3] > imageBds[3])
  {
    vtkWarningMacro(<< "Input polydata extends beyond height map");
  }
  this->Offset = (this->UseHeightMapOffset ? imageBds[4] : 0.0);

  // Okay we are ready to rock and roll...
  output->CopyStructure(input);

  vtkNew<vtkPoints> newPts;
  newPts->SetDataType(inPts->GetDataType());
  newPts->SetNumberOfPoints(numPts);

  vtkPointData* pd = input->GetPointData();
  vtkPointData* outputPD = output->GetPointData();
  outputPD->CopyNormalsOff(); // normals are almost certainly messed up
  outputPD->PassData(pd);

  vtkCellData* cd = input->GetCellData();
  vtkCellData* outputCD = output->GetCellData();
  outputCD->PassData(cd);

  // We need random access to cells
  if (output->NeedToBuildCells())
  {
    output->BuildCells();
  }

  // This performs the projection of the points or cells.
  // We either process points or cells depending on strategy
  if (this->FittingStrategy <= vtkFitToHeightMapFilter::POINT_AVERAGE_HEIGHT)
  {
    using Dispatcher =
      vtkArrayDispatch::Dispatch2ByArray<vtkArrayDispatch::PointArrays, vtkArrayDispatch::Arrays>;
    FitPointsWorker worker;
    if (!Dispatcher::Execute(inPts->GetData(), image->GetPointData()->GetScalars(), worker,
          newPts->GetData(), dims, origin, h, this))
    {
      worker(inPts->GetData(), image->GetPointData()->GetScalars(), newPts->GetData(), dims, origin,
        h, this);
    }
    // Now final rollup and adjustment of points
    this->AdjustPoints(output, numCells, newPts);

  } // points strategies

  else // We are processing cells
  {
    vtkNew<vtkAOSDataArrayTemplate<double>> cellHtsArray;
    cellHtsArray->SetNumberOfValues(numCells);

    FitCellsWorker worker;
    if (!vtkArrayDispatch::Dispatch::Execute(image->GetPointData()->GetScalars(), worker,
          this->FittingStrategy, output, cellHtsArray->GetPointer(0), dims, origin, h, this))
    {
      worker(image->GetPointData()->GetScalars(), this->FittingStrategy, output,
        cellHtsArray->GetPointer(0), dims, origin, h, this);
    }

    // Now final rollup and adjustment of points
    this->AdjustCells(output, numCells, cellHtsArray->GetPointer(0), inPts, newPts);

  } // cells strategies

  // Clean up and get out. Replace the output's shallow-copied points with
  // the new, projected points.
  output->SetPoints(newPts);

  return 1;
}

//------------------------------------------------------------------------------
// Based on the fitting strategy, adjust the point coordinates.
void vtkFitToHeightMapFilter::AdjustPoints(
  vtkPolyData* output, vtkIdType numCells, vtkPoints* newPts)
{
  vtkIdType cellId;
  vtkIdType npts;
  const vtkIdType* ptIds;
  vtkIdType pId;
  vtkIdType i;
  double sum, min, max, p0[3], z;
  vtkIdType numHits;

  // Nothing to do except adjust offset if point projection
  if (this->FittingStrategy == vtkFitToHeightMapFilter::POINT_PROJECTION)
  {
    if (this->UseHeightMapOffset)
    {
      npts = newPts->GetNumberOfPoints();
      for (pId = 0; pId < npts; ++pId)
      {
        newPts->GetPoint(pId, p0);
        newPts->SetPoint(pId, p0[0], p0[1], p0[2] + this->Offset);
      }
    }
    return;
  }

  // Otherwise fancier point adjustment
  for (cellId = 0; cellId < numCells; ++cellId)
  {
    if (this->CheckAbort())
    {
      break;
    }
    output->GetCellPoints(cellId, npts, ptIds);

    // Gather information about cell
    min = VTK_FLOAT_MAX;
    max = VTK_FLOAT_MIN;
    sum = 0.0;
    numHits = 0;
    for (i = 0; i < npts; ++i)
    {
      pId = ptIds[i];
      numHits++;
      newPts->GetPoint(pId, p0);
      z = p0[2];

      min = std::min(z, min);
      max = std::max(z, max);
      sum += z; // to compute average
    }           // over primitive points

    // Adjust points as specified.
    if (numHits > 0)
    {
      if (this->FittingStrategy == vtkFitToHeightMapFilter::POINT_AVERAGE_HEIGHT)
      {
        z = fabs(sum / static_cast<double>(numHits));
      }
      else if (this->FittingStrategy == vtkFitToHeightMapFilter::POINT_MINIMUM_HEIGHT)
      {
        z = min;
      }
      else // if ( this->FittingStrategy == vtkFitToHeightMapFilter::POINT_MAXIMUM_HEIGHT )
      {
        z = max;
      }

      for (i = 0; i < npts; ++i)
      {
        pId = ptIds[i];
        newPts->GetPoint(pId, p0);
        newPts->SetPoint(pId, p0[0], p0[1], z + this->Offset);
      }
    } // if valid polygon

  } // for all cells
}

//------------------------------------------------------------------------------
// Based on the fitting strategy, adjust the points based on cell height
// information.
void vtkFitToHeightMapFilter::AdjustCells(
  vtkPolyData* output, vtkIdType numCells, double* cellHts, vtkPoints* inPts, vtkPoints* newPts)
{
  vtkIdType cellId;
  vtkIdType npts;
  const vtkIdType* ptIds;
  vtkIdType pId;
  vtkIdType i;
  double z, p0[3];

  for (cellId = 0; cellId < numCells; ++cellId)
  {
    if (this->CheckAbort())
    {
      break;
    }
    z = cellHts[cellId];
    output->GetCellPoints(cellId, npts, ptIds);

    for (i = 0; i < npts; ++i)
    {
      pId = ptIds[i];
      inPts->GetPoint(pId, p0);
      newPts->SetPoint(pId, p0[0], p0[1], z + this->Offset);
    }
  } // for all cells
}

//------------------------------------------------------------------------------
// Specify the height map connection
void vtkFitToHeightMapFilter::SetHeightMapConnection(vtkAlgorithmOutput* algOutput)
{
  this->SetInputConnection(1, algOutput);
}

//------------------------------------------------------------------------------
// Specify the height map data
void vtkFitToHeightMapFilter::SetHeightMapData(vtkImageData* id)
{
  this->SetInputData(1, id);
}

//------------------------------------------------------------------------------
// Get a pointer to a source object at a specified table location.
vtkImageData* vtkFitToHeightMapFilter::GetHeightMap()
{
  return vtkImageData::SafeDownCast(this->GetExecutive()->GetInputData(1, 0));
}

//------------------------------------------------------------------------------
vtkImageData* vtkFitToHeightMapFilter::GetHeightMap(vtkInformationVector* sourceInfo)
{
  vtkInformation* info = sourceInfo->GetInformationObject(1);
  if (!info)
  {
    return nullptr;
  }
  return vtkImageData::SafeDownCast(info->Get(vtkDataObject::DATA_OBJECT()));
}

//------------------------------------------------------------------------------
int vtkFitToHeightMapFilter::FillInputPortInformation(int port, vtkInformation* info)
{
  if (port == 0)
  {
    info->Set(vtkAlgorithm::INPUT_IS_REPEATABLE(), 0);
    info->Set(vtkAlgorithm::INPUT_IS_OPTIONAL(), 0);
    info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkPolyData");
  }
  else if (port == 1)
  {
    info->Set(vtkAlgorithm::INPUT_IS_REPEATABLE(), 0);
    info->Set(vtkAlgorithm::INPUT_IS_OPTIONAL(), 0);
    info->Set(vtkAlgorithm::INPUT_REQUIRED_DATA_TYPE(), "vtkImageData");
  }

  return 1;
}

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

  os << indent << "Fitting Strategy: " << this->FittingStrategy << "\n";
  os << indent << "Use Height Map Offset: " << (this->UseHeightMapOffset ? "On\n" : "Off\n");
}
VTK_ABI_NAMESPACE_END