// SPDX-FileCopyrightText: Copyright (c) Ken Martin, Will Schroeder, Bill Lorensen
// SPDX-License-Identifier: BSD-3-Clause

#include "vtkProcrustesAlignmentFilter.h"

#include "vtkExecutive.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkLandmarkTransform.h"
#include "vtkMath.h"
#include "vtkMultiBlockDataSet.h"
#include "vtkObjectFactory.h"
#include "vtkPolyData.h"
#include "vtkTransformPolyDataFilter.h"

VTK_ABI_NAMESPACE_BEGIN
vtkStandardNewMacro(vtkProcrustesAlignmentFilter);

//------------------------------------------------------------------------------
// protected
vtkProcrustesAlignmentFilter::vtkProcrustesAlignmentFilter()
{
  this->LandmarkTransform = vtkLandmarkTransform::New();
  this->StartFromCentroid = false;
  this->OutputPointsPrecision = DEFAULT_PRECISION;

  // The precision of the mean points is set in RequestData().
  this->MeanPoints = vtkPoints::New();
}

//------------------------------------------------------------------------------
// protected
vtkProcrustesAlignmentFilter::~vtkProcrustesAlignmentFilter()
{
  if (this->LandmarkTransform)
  {
    this->LandmarkTransform->Delete();
  }
  if (this->MeanPoints)
  {
    this->MeanPoints->Delete();
  }
}

//------------------------------------------------------------------------------
// Calculate the centroid of a point cloud
static inline void Centroid(vtkPoints* pd, double* cp)
{
  // Center point
  cp[0] = 0;
  cp[1] = 0;
  cp[2] = 0;

  int np = pd->GetNumberOfPoints();

  // Calculate center of shape
  for (int i = 0; i < np; i++)
  {
    double p[3];
    pd->GetPoint(i, p);
    cp[0] += p[0];
    cp[1] += p[1];
    cp[2] += p[2];
  }
  cp[0] /= np;
  cp[1] /= np;
  cp[2] /= np;
}

//------------------------------------------------------------------------------
// Calculate the centroid size of a point cloud
static inline double CentroidSize(vtkPoints* pd, double* cp)
{
  Centroid(pd, cp);

  double S = 0;
  for (int i = 0; i < pd->GetNumberOfPoints(); i++)
  {
    double p[3];
    pd->GetPoint(i, p);
    S += vtkMath::Distance2BetweenPoints(p, cp);
  }

  return sqrt(S);
}

//------------------------------------------------------------------------------
// Translation of point cloud. Could be done using transformations
static inline void TranslateShape(vtkPoints* pd, double* tp)
{
  for (int i = 0; i < pd->GetNumberOfPoints(); i++)
  {
    double p[3];
    pd->GetPoint(i, p);
    pd->SetPoint(i, p[0] + tp[0], p[1] + tp[1], p[2] + tp[2]);
  }
}

//------------------------------------------------------------------------------
// Scaling of point cloud. Could be done using transformations
static inline void ScaleShape(vtkPoints* pd, double S)
{
  for (int i = 0; i < pd->GetNumberOfPoints(); i++)
  {
    double p[3];
    pd->GetPoint(i, p);
    pd->SetPoint(i, p[0] * S, p[1] * S, p[2] * S);
  }
}

//------------------------------------------------------------------------------
// Normalise a point cloud to have centroid (0,0,0) and centroid size 1
static inline int NormaliseShape(vtkPoints* pd)
{
  double cp[3];
  double S = CentroidSize(pd, cp);
  if (S == 0)
    return 0;

  double tp[3];
  tp[0] = -cp[0];
  tp[1] = -cp[1];
  tp[2] = -cp[2];
  TranslateShape(pd, tp);
  ScaleShape(pd, 1 / S);
  return 1;
}

//------------------------------------------------------------------------------
// protected
int vtkProcrustesAlignmentFilter::RequestData(vtkInformation* vtkNotUsed(request),
  vtkInformationVector** inputVector, vtkInformationVector* outputVector)
{
  // get the info objects
  vtkInformation* inInfo = inputVector[0]->GetInformationObject(0);
  vtkInformation* outInfo = outputVector->GetInformationObject(0);

  // get the input and output
  vtkMultiBlockDataSet* mbInput =
    vtkMultiBlockDataSet::SafeDownCast(inInfo->Get(vtkDataObject::DATA_OBJECT()));
  const int N_SETS = mbInput->GetNumberOfBlocks();
  if (N_SETS == 0)
  {
    return 1;
  }

  int i, v;
  vtkPointSet* input = nullptr;
  for (i = 0; i < N_SETS; i++)
  {
    input = vtkPointSet::SafeDownCast(mbInput->GetBlock(i));
    if (input)
    {
      break;
    }
  }

  if (!input)
  {
    return 1;
  }

  vtkMultiBlockDataSet* output =
    vtkMultiBlockDataSet::SafeDownCast(outInfo->Get(vtkDataObject::DATA_OBJECT()));

  vtkDebugMacro(<< "Execute()");

  vtkPointSet* tmpInput;

  // copy the inputs across
  // (really actually only the points need to be deep copied since the rest stays the same)
  for (i = 0; i < N_SETS; i++)
  {
    tmpInput = vtkPointSet::SafeDownCast(mbInput->GetBlock(i));
    vtkPointSet* outputBlock = nullptr;
    if (tmpInput)
    {
      outputBlock = tmpInput->NewInstance();
      outputBlock->DeepCopy(tmpInput);

      // Set the desired precision for the points in the output. If
      // this->OutputPointsPrecision == vtkAlgorithm::DEFAULT_PRECISION
      // then the precision of the points in the output is correctly
      // set during the deep copy of tmpInput.
      if (this->OutputPointsPrecision == vtkAlgorithm::SINGLE_PRECISION)
      {
        // Only create another new vtkPoints if the output precision is
        // different from the input.
        if (tmpInput->GetPoints() && tmpInput->GetPoints()->GetDataType() != VTK_FLOAT)
        {
          vtkPoints* newPoints = vtkPoints::New();
          newPoints->SetDataType(VTK_FLOAT);
          newPoints->DeepCopy(tmpInput->GetPoints());
          outputBlock->SetPoints(newPoints);
          newPoints->Delete();
        }
      }
      else if (this->OutputPointsPrecision == vtkAlgorithm::DOUBLE_PRECISION)
      {
        // Only create another new vtkPoints if the output precision is
        // different from the input.
        if (tmpInput->GetPoints() && tmpInput->GetPoints()->GetDataType() != VTK_DOUBLE)
        {
          vtkPoints* newPoints = vtkPoints::New();
          newPoints->SetDataType(VTK_DOUBLE);
          newPoints->DeepCopy(tmpInput->GetPoints());
          outputBlock->SetPoints(newPoints);
          newPoints->Delete();
        }
      }
    }
    output->SetBlock(i, outputBlock);
    if (outputBlock)
    {
      outputBlock->Delete();
    }
  }

  // the number of points is determined by the first input (they must all be the same)
  const int N_POINTS = input->GetNumberOfPoints();

  vtkDebugMacro(<< "N_POINTS is " << N_POINTS);

  if (N_POINTS == 0)
  {
    vtkErrorMacro(<< "No points!");
    return 1;
  }

  // all the inputs must have the same number of points to consider executing

  for (i = 1; i < N_SETS; i++)
  {
    tmpInput = vtkPointSet::SafeDownCast(mbInput->GetBlock(i));
    if (!tmpInput)
    {
      continue;
    }
    if (tmpInput->GetNumberOfPoints() != N_POINTS)
    {
      vtkErrorMacro(<< "The inputs have different numbers of points!");
      return 1;
    }
  }

  // Set the desired precision for the mean points.
  if (this->OutputPointsPrecision == vtkAlgorithm::DEFAULT_PRECISION)
  {
    // The points in distinct blocks may be of differing precisions.
    this->MeanPoints->SetDataType(VTK_FLOAT);
    for (i = 0; i < N_SETS; i++)
    {
      tmpInput = vtkPointSet::SafeDownCast(mbInput->GetBlock(i));

      // Set the desired precision to VTK_DOUBLE if the precision of the
      // mean points in any of the blocks is VTK_DOUBLE.
      if (tmpInput && tmpInput->GetPoints()->GetDataType() == VTK_DOUBLE)
      {
        this->MeanPoints->SetDataType(VTK_DOUBLE);
        break;
      }
    }
  }
  else if (this->OutputPointsPrecision == vtkAlgorithm::SINGLE_PRECISION)
  {
    this->MeanPoints->SetDataType(VTK_FLOAT);
  }
  else if (this->OutputPointsPrecision == vtkAlgorithm::DOUBLE_PRECISION)
  {
    this->MeanPoints->SetDataType(VTK_DOUBLE);
  }

  this->MeanPoints->DeepCopy(input->GetPoints());
  // our initial estimate of the mean comes from the first example in the set

  // Move to the mutual centroid of the data if requested.
  if (this->GetStartFromCentroid())
  {
    double meanCentroid[3];
    double firstCentroid[3];
    Centroid(MeanPoints, firstCentroid);
    meanCentroid[0] = firstCentroid[0];
    meanCentroid[1] = firstCentroid[1];
    meanCentroid[2] = firstCentroid[2];

    for (i = 1; i < N_SETS; i++)
    {
      tmpInput = vtkPointSet::SafeDownCast(mbInput->GetBlock(i));
      if (!tmpInput)
      {
        continue;
      }
      double localCentroid[3];
      Centroid(tmpInput->GetPoints(), localCentroid);
      meanCentroid[0] += localCentroid[0];
      meanCentroid[1] += localCentroid[1];
      meanCentroid[2] += localCentroid[2];
    }
    meanCentroid[0] /= N_SETS;
    meanCentroid[1] /= N_SETS;
    meanCentroid[2] /= N_SETS;

    double translate[3];
    translate[0] = meanCentroid[0] - firstCentroid[0];
    translate[1] = meanCentroid[1] - firstCentroid[1];
    translate[2] = meanCentroid[2] - firstCentroid[2];

    TranslateShape(MeanPoints, translate);
  }

  // we keep a record of the first mean to fix the orientation and scale
  // (which are otherwise undefined and the loop will not converge)
  vtkPoints* first_mean = vtkPoints::New();
  first_mean->SetDataType(this->MeanPoints->GetDataType());
  first_mean->DeepCopy(MeanPoints);

  // If the similarity transform is used, the mean shape must be normalised
  // to avoid shrinking
  if (this->LandmarkTransform->GetMode() == VTK_LANDMARK_SIMILARITY)
  {
    if (!NormaliseShape(MeanPoints))
    {
      vtkErrorMacro(<< "Centroid size zero");
      return 1;
    }
    if (!NormaliseShape(first_mean))
    {
      vtkErrorMacro(<< "Centroid size zero");
      return 1;
    }
  }

  // storage for the new mean that is being calculated
  vtkPoints* new_mean = vtkPoints::New();
  new_mean->SetDataType(this->MeanPoints->GetDataType());
  new_mean->SetNumberOfPoints(N_POINTS);

  // compute mean and align all the shapes to it, until convergence
  int converged = 0; // bool converged=false
  int iterations = 0;
  constexpr int MAX_ITERATIONS = 5;
  double difference;
  double point[3], p[3], p2[3];
  double outPoint[3];
  do
  { // (while not converged)

    // align each pointset with the mean
    for (i = 0; i < N_SETS; i++)
    {
      vtkPointSet* block = vtkPointSet::SafeDownCast(output->GetBlock(i));
      if (!block)
      {
        continue;
      }
      this->LandmarkTransform->SetSourceLandmarks(block->GetPoints());
      this->LandmarkTransform->SetTargetLandmarks(this->MeanPoints);
      this->LandmarkTransform->Update();
      for (v = 0; v < N_POINTS; v++)
      {
        this->LandmarkTransform->InternalTransformPoint(block->GetPoint(v), outPoint);
        block->GetPoints()->SetPoint(v, outPoint);
      }
    }

    // compute the new mean (just average the point locations)
    for (v = 0; v < N_POINTS; v++)
    {
      point[0] = 0.0F;
      point[1] = 0.0F;
      point[2] = 0.0F;
      for (i = 0; i < N_SETS; i++)
      {
        vtkPointSet* block = vtkPointSet::SafeDownCast(output->GetBlock(i));
        block->GetPoint(v, p);
        point[0] += p[0];
        point[1] += p[1];
        point[2] += p[2];
      }
      p[0] = point[0] / (double)N_SETS;
      p[1] = point[1] / (double)N_SETS;
      p[2] = point[2] / (double)N_SETS;
      new_mean->SetPoint(v, p);
    }

    // align the new mean with the fixed mean if the transform
    // is similarity or rigidbody. It is not yet decided what to do with affine
    if (this->LandmarkTransform->GetMode() == VTK_LANDMARK_SIMILARITY ||
      this->LandmarkTransform->GetMode() == VTK_LANDMARK_RIGIDBODY)
    {
      this->LandmarkTransform->SetSourceLandmarks(new_mean);
      this->LandmarkTransform->SetTargetLandmarks(first_mean);
      this->LandmarkTransform->Update();
      for (v = 0; v < N_POINTS; v++)
      {
        this->LandmarkTransform->InternalTransformPoint(new_mean->GetPoint(v), outPoint);
        new_mean->SetPoint(v, outPoint);
      }
    }

    // If the similarity transform is used, the mean shape must be normalised
    // to avoid shrinking
    if (this->LandmarkTransform->GetMode() == VTK_LANDMARK_SIMILARITY)
    {
      if (!NormaliseShape(new_mean))
      {
        vtkErrorMacro(<< "Centroid size zero");
        return 1;
      }
    }

    // the new mean becomes our mean
    // compute the difference between the two
    difference = 0.0F;
    for (v = 0; v < N_POINTS; v++)
    {
      new_mean->GetPoint(v, p);
      MeanPoints->GetPoint(v, p2);
      difference += vtkMath::Distance2BetweenPoints(p, p2);
      MeanPoints->SetPoint(v, p);
    }

    // test for convergence
    iterations++;
    vtkDebugMacro(<< "Difference after " << iterations << " iteration(s) is: " << difference);
    if (difference < 1e-6 || iterations >= MAX_ITERATIONS)
    {
      converged = 1; // true
    }

    // The convergence test is that the sum of the distances between the
    // points on mean(t) and mean(t-1) is less than a very small number.
    // Procrustes shouldn't need more than 2 or 3 iterations but things could go wrong
    // so we impose an iteration limit to avoid getting stuck in an infinite loop.

  } while (!converged);

  if (iterations >= MAX_ITERATIONS)
  {
    vtkDebugMacro(<< "Procrustes did not converge in  " << MAX_ITERATIONS
                  << " iterations! Objects may not be aligned. Difference = " << difference);
    // we don't throw an Error here since the shapes most probably *are* aligned, but the
    // numerical precision is worse than our convergence test anticipated.
  }
  else
  {
    vtkDebugMacro(<< "Procrustes required " << iterations << " iterations to converge to "
                  << difference);
  }

  // clean up
  first_mean->Delete();
  new_mean->Delete();

  return 1;
}

//------------------------------------------------------------------------------
// public
void vtkProcrustesAlignmentFilter::PrintSelf(ostream& os, vtkIndent indent)
{
  this->Superclass::PrintSelf(os, indent);
  this->LandmarkTransform->PrintSelf(os, indent.GetNextIndent());
  this->MeanPoints->PrintSelf(os, indent.GetNextIndent());
  os << indent << "Start From Centroid: " << (this->StartFromCentroid ? "On\n" : "Off\n");
  os << indent << "Output Points Precision: " << this->OutputPointsPrecision << "\n";
}
VTK_ABI_NAMESPACE_END