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

  Program:   Visualization Toolkit
  Module:    $RCSfile: vtkProcrustesAlignmentFilter.cxx,v $

  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 "vtkProcrustesAlignmentFilter.h"
#include "vtkExecutive.h"
#include "vtkInformation.h"
#include "vtkInformationVector.h"
#include "vtkObjectFactory.h"
#include "vtkLandmarkTransform.h"
#include "vtkTransformPolyDataFilter.h"
#include "vtkPolyData.h"
#include "vtkMath.h"

vtkCxxRevisionMacro(vtkProcrustesAlignmentFilter, "$Revision: 1.20 $");
vtkStandardNewMacro(vtkProcrustesAlignmentFilter);

//----------------------------------------------------------------------------
// protected
vtkProcrustesAlignmentFilter::vtkProcrustesAlignmentFilter()
{
  this->LandmarkTransform = vtkLandmarkTransform::New();

  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 ouptut
  vtkPointSet *input = vtkPointSet::SafeDownCast(
    inInfo->Get(vtkDataObject::DATA_OBJECT()));
  vtkPointSet *output = vtkPointSet::SafeDownCast(
    outInfo->Get(vtkDataObject::DATA_OBJECT()));

  vtkDebugMacro(<<"Execute()");

  int i,v;

  const int N_SETS = this->GetNumberOfInputConnections(0);

  // copy the inputs across
  // (really actually only the points need to be deep copied since the rest stays the same)
  vtkInformation *tmpInfo;
  vtkPointSet *tmpInput;
  output->DeepCopy(input);
  for(i=1;i<N_SETS;i++)
    {
    tmpInfo = inputVector[0]->GetInformationObject(i);
    tmpInput = 0;
    if (tmpInfo)
      {
      tmpInput =
        vtkPointSet::SafeDownCast(tmpInfo->Get(vtkDataObject::DATA_OBJECT()));
      }
    this->GetOutput(i)->DeepCopy(tmpInput);
    }

  // 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++) 
    {
    tmpInfo = inputVector[0]->GetInformationObject(i);
    tmpInput = 0;
    if (tmpInfo)
      {
      tmpInput =
        vtkPointSet::SafeDownCast(tmpInfo->Get(vtkDataObject::DATA_OBJECT()));
      }
    else
      {
      continue;
      }
    if(tmpInput->GetNumberOfPoints() != N_POINTS)
      {
      vtkErrorMacro(<<"The inputs have different numbers of points!");
      return 1;
      }
    }

//  vtkPoints *mean_points = vtkPoints::New();
  MeanPoints->DeepCopy(input->GetPoints());
  // our initial estimate of the mean comes from the first example in the set

  // 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->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->SetNumberOfPoints(N_POINTS);

  // compute mean and align all the shapes to it, until convergence
  int converged=0; // bool converged=false
  int iterations=0;
  const 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++)
      {
      this->LandmarkTransform->SetSourceLandmarks(this->GetOutput(i)->GetPoints());
      this->LandmarkTransform->SetTargetLandmarks(MeanPoints);
      this->LandmarkTransform->Update();
      for(v=0;v<N_POINTS;v++)
        {
        this->LandmarkTransform->InternalTransformPoint(
          this->GetOutput(i)->GetPoint(v), outPoint);
        this->GetOutput(i)->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++)
        {
        this->GetOutput(i)->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::SetNumberOfInputs(int n)
{ 
  this->SetNumberOfInputConnections(0, n);
  this->SetNumberOfOutputPorts(n);

  // initialise the outputs
  for(int i=0;i<n;i++)
    {
    vtkPoints *points = vtkPoints::New();
    vtkPolyData *ps = vtkPolyData::New();
    ps->SetPoints(points);
    points->Delete();
    this->GetExecutive()->SetOutputData(i,ps);
    ps->Delete();
    }

  // is this the right thing to be doing here? if we don't initialise the outputs here
  // then the filter crashes but vtkPolyData may not be the type of the inputs
}

//----------------------------------------------------------------------------
void vtkProcrustesAlignmentFilter::SetInput(int idx, vtkPointSet *p)
{
  this->SetNthInputConnection(0, idx, p ? p->GetProducerPort() : 0);
}

//----------------------------------------------------------------------------
vtkPointSet* vtkProcrustesAlignmentFilter::GetInput(int idx)
{
  return vtkPointSet::SafeDownCast(
    this->GetExecutive()->GetInputData(0, idx));
}

//----------------------------------------------------------------------------
int vtkProcrustesAlignmentFilter::FillInputPortInformation(
  int port, vtkInformation *info)
{
  int retval = this->Superclass::FillInputPortInformation(port, info);
  info->Set(vtkAlgorithm::INPUT_IS_REPEATABLE(), 1);
  return retval;
}

//----------------------------------------------------------------------------
// 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());
}