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

#include "vtkArrayDispatch.h"
#include "vtkArrayDispatchDataSetArrayList.h"
#include "vtkArrayListTemplate.h" // For processing attribute data
#include "vtkCamera.h"
#include "vtkCellArray.h"
#include "vtkCommand.h"
#include "vtkCoordinate.h"
#include "vtkDataArrayRange.h"
#include "vtkFloatArray.h"
#include "vtkImageData.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkMatrix4x4.h"
#include "vtkObjectFactory.h"
#include "vtkPointData.h"
#include "vtkPolyData.h"
#include "vtkRenderWindow.h"
#include "vtkRenderer.h"
#include "vtkSMPTools.h"
#include "vtkStreamingDemandDrivenPipeline.h"

VTK_ABI_NAMESPACE_BEGIN
vtkStandardNewMacro(vtkDepthImageToPointCloud);
vtkCxxSetObjectMacro(vtkDepthImageToPointCloud, Camera, vtkCamera);

//------------------------------------------------------------------------------
// Helper classes to support efficient computing, and threaded execution.
namespace
{
struct MapPointsFunctor
{
  // Map input point id to output point id. This map is needed because of the
  // optionally capability to cull near and far points.
  template <typename TArray>
  void operator()(TArray* depthsArray, vtkIdType numPts, bool cullNear, bool cullFar,
    vtkIdType* map, vtkIdType& numOutPts)
  {
    auto depths = vtk::DataArrayValueRange<1>(depthsArray).begin();
    numOutPts = 0;
    float d;
    for (vtkIdType ptId = 0; ptId < numPts; ++depths, ++map, ++ptId)
    {
      d = static_cast<float>(*depths);
      if ((cullNear && d <= 0.0) || (cullFar && d >= 1.0))
      {
        *map = (-1);
      }
      else
      {
        *map = numOutPts++;
      }
    }
  }
};

// This class performs point by point transformation. The view matrix is
// used to transform each pixel. IMPORTANT NOTE: The transformation occurs
// by normalizing the image pixels into the (-1,1) view space (depth values
// are passed through). The process follows the vtkCoordinate class which is
// the standard for VTK rendering transformations. Subtle differences in
// whether the lower left pixel origin are at the center of the pixel
// versus the lower-left corner of the pixel will make slight differences
// in how pixels are transformed. (Similarly for the upper right pixel as
// well). This half pixel difference can cause transformation issues. Here
// we've played around with the scaling below to produce the best results
// in the current version of VTK.
template <typename TDepthsArray, typename TPointsArray>
struct MapDepthImage
{
  TDepthsArray* Depths;
  TPointsArray* Pts;
  const int* Dims;
  const double* Matrix;
  const vtkIdType* PtMap;

  MapDepthImage(TDepthsArray* depths, TPointsArray* pts, int dims[2], double* m, vtkIdType* ptMap)
    : Depths(depths)
    , Pts(pts)
    , Dims(dims)
    , Matrix(m)
    , PtMap(ptMap)
  {
  }

  void operator()(vtkIdType row, vtkIdType end)
  {
    double drow, result[4];
    vtkIdType offset = row * this->Dims[0];
    auto dptr = vtk::DataArrayValueRange<1>(this->Depths, offset).begin();
    const vtkIdType* mptr = this->PtMap + offset;
    auto points = vtk::DataArrayTupleRange<3>(this->Pts);
    for (; row < end; ++row)
    {
      drow = -1.0 + (2.0 * static_cast<double>(row) / static_cast<double>(this->Dims[1] - 1));
      // if pixel origin is pixel center use the two lines below
      // drow = -1.0 + 2.0*((static_cast<double>(row)+0.5) /
      //                    static_cast<double>(this->Dims[1]));
      for (vtkIdType i = 0; i < this->Dims[0]; ++dptr, ++mptr, ++i)
      {
        if (*mptr > (-1)) // if not masked
        {
          auto point = points[*mptr];
          result[0] = -1.0 + 2.0 * static_cast<double>(i) / static_cast<double>(this->Dims[0] - 1);
          // if pixel origin is pixel center use the two lines below
          // result[0] = -1.0 + 2.0*((static_cast<double>(i)+0.5) /
          //                         static_cast<double>(this->Dims[0]));
          result[1] = drow;
          result[2] = *dptr;
          result[3] = 1.0;
          vtkMatrix4x4::MultiplyPoint(this->Matrix, result, result);
          point[0] = result[0] / result[3]; // x
          point[1] = result[1] / result[3]; // y
          point[2] = result[2] / result[3]; // z
        }                                   // transform this point
      }
    }
  }
};

// Interface to vtkSMPTools. Threading over image rows. Also perform
// one time calculation/initialization for more efficient processing.
struct XFormPointsFunctor
{
  template <typename TDepthsArray, typename TPointsArray>
  void operator()(TDepthsArray* depthsArray, TPointsArray* ptsArray, vtkIdType* ptMap, int dims[2],
    vtkCamera* cam)
  {
    double m[16], aspect = static_cast<double>(dims[0]) / static_cast<double>(dims[1]);
    vtkMatrix4x4* matrix = cam->GetCompositeProjectionTransformMatrix(aspect, 0, 1);
    vtkMatrix4x4::Invert(*matrix->Element, m);

    MapDepthImage<TDepthsArray, TPointsArray> mapDepths(depthsArray, ptsArray, dims, m, ptMap);
    vtkSMPTools::For(0, dims[1], mapDepths);
  }
};

// Process the color scalars. It would be pretty easy to process all
// attribute types if this was ever desired.
struct MapScalars
{
  vtkIdType NumColors;
  vtkDataArray* InColors;
  ArrayList Colors;
  const vtkIdType* PtMap;
  vtkDataArray* OutColors;

  MapScalars(vtkIdType num, vtkDataArray* colors, vtkIdType* ptMap)
    : NumColors(num)
    , InColors(colors)
    , PtMap(ptMap)
    , OutColors(nullptr)
  {
    vtkStdString outName = "DepthColors";
    this->OutColors = vtkArrayDownCast<vtkDataArray>(
      Colors.AddArrayPair(this->NumColors, this->InColors, outName, 0.0, false));
  }

  void operator()(vtkIdType id, vtkIdType end)
  {
    vtkIdType outId;
    for (; id < end; ++id)
    {
      if ((outId = this->PtMap[id]) > (-1))
      {
        this->Colors.Copy(id, outId);
      }
    }
  }
};

} // anonymous namespace

//================= Begin class proper =======================================
//------------------------------------------------------------------------------
vtkDepthImageToPointCloud::vtkDepthImageToPointCloud()
{
  this->Camera = nullptr;
  this->CullNearPoints = false;
  this->CullFarPoints = true;
  this->ProduceColorScalars = true;
  this->ProduceVertexCellArray = true;
  this->OutputPointsPrecision = vtkAlgorithm::DEFAULT_PRECISION;

  this->SetNumberOfInputPorts(2);
  this->SetNumberOfOutputPorts(1);
}

//------------------------------------------------------------------------------
vtkDepthImageToPointCloud::~vtkDepthImageToPointCloud()
{
  if (this->Camera)
  {
    this->Camera->UnRegister(this);
    this->Camera = nullptr;
  }
}

//------------------------------------------------------------------------------
vtkMTimeType vtkDepthImageToPointCloud::GetMTime()
{
  vtkCamera* cam = this->GetCamera();
  vtkMTimeType t1 = this->MTime.GetMTime();
  vtkMTimeType t2;

  if (!cam)
  {
    return t1;
  }

  // Check the camera
  t2 = cam->GetMTime();
  t1 = std::max(t2, t1);

  return t1;
}

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

  return 1;
}

//------------------------------------------------------------------------------
int vtkDepthImageToPointCloud::FillOutputPortInformation(int vtkNotUsed(port), vtkInformation* info)
{
  // now add our info
  info->Set(vtkDataObject::DATA_TYPE_NAME(), "vtkPolyData");
  return 1;
}

//------------------------------------------------------------------------------
int vtkDepthImageToPointCloud::RequestInformation(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** vtkNotUsed(inputVector), vtkInformationVector* vtkNotUsed(outputVector))
{
  return 1;
}

//------------------------------------------------------------------------------
int vtkDepthImageToPointCloud::RequestUpdateExtent(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** inputVector, vtkInformationVector* vtkNotUsed(outputVector))
{
  int inExt[6];
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);

  inInfo->Get(vtkStreamingDemandDrivenPipeline::WHOLE_EXTENT(), inExt);
  inInfo->Set(vtkStreamingDemandDrivenPipeline::UPDATE_EXTENT(), inExt, 6);

  // need to set the stencil update extent to the input extent
  if (this->GetNumberOfInputConnections(1) > 0)
  {
    vtkInformation* in2Info = inputVector[1]->GetInformationObject(0);
    in2Info->Set(vtkStreamingDemandDrivenPipeline::UPDATE_EXTENT(), inExt, 6);
  }

  return 1;
}

//------------------------------------------------------------------------------
int vtkDepthImageToPointCloud::RequestData(
  vtkInformation*, vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  // Get the input, make sure that it is valid
  int numInputs = 0;
  vtkInformation* info = inputVector[0]->GetInformationObject(0);
  vtkImageData* inData = vtkImageData::SafeDownCast(info->Get(vtkDataObject::DATA_OBJECT()));
  if (inData == nullptr)
  {
    vtkErrorMacro("At least one input image is required");
    return 0;
  }
  ++numInputs;

  vtkInformation* info2 = inputVector[1]->GetInformationObject(0);
  vtkImageData* inData2 = nullptr;
  if (info2)
  {
    inData2 = vtkImageData::SafeDownCast(info2->Get(vtkDataObject::DATA_OBJECT()));
    if (inData2 != nullptr)
    {
      ++numInputs;
    }
  }

  vtkCamera* cam = this->Camera;
  if (cam == nullptr)
  {
    vtkErrorMacro("Input camera required");
    return 0;
  }

  // At this point we have at least one input, possibly two. If one input, we
  // assume we either have 1) depth values or 2) color scalars + depth values
  // (if depth values are in an array called "ZBuffer".) If two inputs, then the
  // depth values are in input0 and the color scalars are in input1.
  vtkDataArray* depths = nullptr;
  vtkDataArray* colors = nullptr;
  if (numInputs == 2)
  {
    depths = inData->GetPointData()->GetScalars();
    colors = inData2->GetPointData()->GetScalars();
  }
  else if (numInputs == 1)
  {
    if ((depths = inData->GetPointData()->GetArray("ZBuffer")) != nullptr)
    {
      colors = inData->GetPointData()->GetScalars();
    }
    else
    {
      depths = inData->GetPointData()->GetScalars();
    }
  }
  else
  {
    vtkErrorMacro("Wrong number of inputs");
    return 0;
  }

  // Extract relevant information to generate output
  vtkInformation* outInfo = outputVector->GetInformationObject(0);
  vtkPolyData* outData = vtkPolyData::SafeDownCast(outInfo->Get(vtkDataObject::DATA_OBJECT()));

  // Determine the image extents
  const int* ext = inData->GetExtent();
  int dims[2];
  dims[0] = ext[1] - ext[0] + 1;
  dims[1] = ext[3] - ext[2] + 1;
  vtkIdType numPts = dims[0] * dims[1];

  // Estimate the total number of output points. Note that if we are culling
  // near and.or far points, then the number of output points is not known,
  // so a point mask is created.
  vtkIdType numOutPts = 0;
  vtkIdType* ptMap = new vtkIdType[numPts];
  using FloatArrays = vtkArrayDispatch::FilterArraysByValueType<vtkArrayDispatch::Arrays,
    vtkArrayDispatch::Reals>::Result;
  MapPointsFunctor mapfunctor;
  if (!vtkArrayDispatch::DispatchByArray<FloatArrays>::Execute(
        depths, mapfunctor, numPts, this->CullNearPoints, this->CullFarPoints, ptMap, numOutPts))
  {
    mapfunctor(depths, numPts, this->CullNearPoints, this->CullFarPoints, ptMap, numOutPts);
  }

  // Manage the requested output point precision
  int pointsType = VTK_DOUBLE;
  if (this->OutputPointsPrecision == vtkAlgorithm::SINGLE_PRECISION)
  {
    pointsType = VTK_FLOAT;
  }

  // Create the points array which represents the point cloud
  vtkSmartPointer<vtkPoints> points = vtkSmartPointer<vtkPoints>::New();
  points->SetDataType(pointsType);
  points->SetNumberOfPoints(numOutPts);
  outData->SetPoints(points);

  // Threaded over x-edges (rows). Each depth value is transformed into a
  // world point. Below there is a double allocation based on the depth type
  // and output point type.
  using Dispatcher =
    vtkArrayDispatch::Dispatch2ByArray<FloatArrays, vtkArrayDispatch::AOSPointArrays>;
  XFormPointsFunctor formFunctor;
  if (!Dispatcher::Execute(depths, points->GetData(), formFunctor, ptMap, dims, cam))
  {
    formFunctor(depths, points->GetData(), ptMap, dims, cam);
  }

  // Produce the output colors if requested. Another templated, threaded loop.
  if (colors && this->ProduceColorScalars)
  {
    vtkPointData* outPD = outData->GetPointData();
    MapScalars mapScalars(numOutPts, colors, ptMap);
    vtkSMPTools::For(0, numPts, mapScalars);
    outPD->SetScalars(mapScalars.OutColors);
  }

  // Clean up
  delete[] ptMap;

  // If requested, create an output vertex array
  if (this->ProduceVertexCellArray)
  {
    vtkSmartPointer<vtkCellArray> verts = vtkSmartPointer<vtkCellArray>::New();
    vtkIdType npts = points->GetNumberOfPoints();
    verts->InsertNextCell(npts);
    for (vtkIdType i = 0; i < npts; ++i)
    {
      verts->InsertCellPoint(i);
    }
    outData->SetVerts(verts);
  }

  return 1;
}

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

  if (this->Camera)
  {
    os << indent << "Camera:\n";
    this->Camera->PrintSelf(os, indent.GetNextIndent());
  }
  else
  {
    os << indent << "Camera: (none)\n";
  }

  os << indent << "Cull Near Points: " << (this->CullNearPoints ? "On\n" : "Off\n");

  os << indent << "Cull Far Points: " << (this->CullFarPoints ? "On\n" : "Off\n");

  os << indent << "Produce Color Scalars: " << (this->ProduceColorScalars ? "On\n" : "Off\n");

  os << indent
     << "Produce Vertex Cell Array: " << (this->ProduceVertexCellArray ? "On\n" : "Off\n");

  os << indent << "OutputPointsPrecision: " << this->OutputPointsPrecision << "\n";
}
VTK_ABI_NAMESPACE_END