/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/
#ifndef itkImagePCAShapeModelEstimator_hxx
#define itkImagePCAShapeModelEstimator_hxx

#include "itkImagePCAShapeModelEstimator.h"

namespace itk
{
template <typename TInputImage, typename TOutputImage>
ImagePCAShapeModelEstimator<TInputImage, TOutputImage>::ImagePCAShapeModelEstimator()
{
  m_EigenVectors.set_size(0, 0);
  m_EigenValues.set_size(0);

  m_NumberOfPrincipalComponentsRequired = 0;
  this->SetNumberOfPrincipalComponentsRequired(1);
}
/**
 * PrintSelf
 */
template <typename TInputImage, typename TOutputImage>
void
ImagePCAShapeModelEstimator<TInputImage, TOutputImage>::PrintSelf(std::ostream & os, Indent indent) const
{
  os << indent << "                   " << std::endl;
  os << indent << "Shape Models " << std::endl;
  os << indent << "Results printed in the superclass " << std::endl;
  os << indent << "                   " << std::endl;

  Superclass::PrintSelf(os, indent);

  itkDebugMacro(<< "                                    ");
  itkDebugMacro(<< "Results of the shape model algorithms");
  itkDebugMacro(<< "====================================");

  itkDebugMacro(<< "The eigen values new method are: ");

  itkDebugMacro(<< m_EigenValues);
  itkDebugMacro(<< m_EigenVectorNormalizedEnergy);

  itkDebugMacro(<< " ");
  itkDebugMacro(<< "==================   ");

  itkDebugMacro(<< "The eigen vectors new method are: ");

  for (unsigned int i = 0; i < m_EigenValues.size(); i++)
  {
    itkDebugMacro(<< m_EigenVectors.get_row(i));
  }

  itkDebugMacro(<< " ");
  itkDebugMacro(<< "+++++++++++++++++++++++++");

  // Print out ivars
  os << indent << "NumberOfPrincipalComponentsRequired: ";
  os << m_NumberOfPrincipalComponentsRequired << std::endl;
  os << indent << "NumberOfTrainingImages: ";
  os << m_NumberOfTrainingImages << std::endl;
} // end PrintSelf

/**
 * Enlarge the output requested region to the largest possible region.
 */
template <typename TInputImage, typename TOutputImage>
void
ImagePCAShapeModelEstimator<TInputImage, TOutputImage>::EnlargeOutputRequestedRegion(DataObject * itkNotUsed(output))
{
  // this filter requires the all of the output images to be in
  // the buffer
  for (unsigned int idx = 0; idx < this->GetNumberOfIndexedOutputs(); ++idx)
  {
    if (this->GetOutput(idx))
    {
      this->GetOutput(idx)->SetRequestedRegionToLargestPossibleRegion();
    }
  }
}

/**
 * Requires all of the inputs to be in the buffer up to the
 * LargestPossibleRegion of the first input.
 */
template <typename TInputImage, typename TOutputImage>
void
ImagePCAShapeModelEstimator<TInputImage, TOutputImage>::GenerateInputRequestedRegion()
{
  Superclass::GenerateInputRequestedRegion();

  if (this->GetInput(0))
  {
    // Set the requested region of the first input to largest possible region
    InputImagePointer input = const_cast<TInputImage *>(this->GetInput(0));
    input->SetRequestedRegionToLargestPossibleRegion();

    // Set the requested region of the remaining input to the largest possible
    // region of the first input
    unsigned int idx;
    for (idx = 1; idx < this->GetNumberOfIndexedInputs(); ++idx)
    {
      if (this->GetInput(idx))
      {
        typename TInputImage::RegionType requestedRegion = this->GetInput(0)->GetLargestPossibleRegion();

        typename TInputImage::RegionType largestRegion = this->GetInput(idx)->GetLargestPossibleRegion();

        if (!largestRegion.IsInside(requestedRegion))
        {
          itkExceptionMacro("LargestPossibleRegion of input "
                            << idx << " is not a superset of the LargestPossibleRegion of input 0");
        }

        InputImagePointer ptr = const_cast<TInputImage *>(this->GetInput(idx));
        ptr->SetRequestedRegion(requestedRegion);
      } // if ( this->GetIntput(idx))
    }   // for idx
  }     // if( this->GetInput(0) )
}

/**
 * Generate data (start the model building process)
 */
template <typename TInputImage, typename TOutputImage>
void
ImagePCAShapeModelEstimator<TInputImage, TOutputImage>::GenerateData()
{
  this->EstimateShapeModels();

  // Allocate memory for each output.
  auto numberOfOutputs = static_cast<unsigned int>(this->GetNumberOfIndexedOutputs());

  InputImagePointer input = const_cast<TInputImage *>(this->GetInput(0));
  unsigned int      j;
  for (j = 0; j < numberOfOutputs; j++)
  {
    OutputImagePointer output = this->GetOutput(j);
    output->SetBufferedRegion(output->GetRequestedRegion());
    output->Allocate();
  }

  // Fill the output images.
  VectorOfDoubleType m_OneEigenVector;
  using OutputIterator = ImageRegionIterator<OutputImageType>;

  // Fill the mean image first

  typename OutputImageType::RegionType region = this->GetOutput(0)->GetRequestedRegion();
  OutputIterator                       outIter(this->GetOutput(0), region);

  unsigned int i = 0;
  outIter.GoToBegin();
  while (!outIter.IsAtEnd())
  {
    outIter.Set(static_cast<typename OutputImageType::PixelType>(m_Means[i]));
    ++outIter;
    ++i;
  }

  // Now fill the principal component outputs
  unsigned int kthLargestPrincipalComp = m_NumberOfTrainingImages;
  unsigned int numberOfValidOutputs = std::min(numberOfOutputs, m_NumberOfTrainingImages + 1);

  for (j = 1; j < numberOfValidOutputs; j++)
  {
    // Extract one column vector at a time
    m_OneEigenVector = m_EigenVectors.get_column(kthLargestPrincipalComp - 1);

    region = this->GetOutput(j)->GetRequestedRegion();
    OutputIterator outIterJ(this->GetOutput(j), region);

    unsigned int idx = 0;
    outIterJ.GoToBegin();
    while (!outIterJ.IsAtEnd())
    {
      outIterJ.Set(static_cast<typename OutputImageType::PixelType>(m_OneEigenVector[idx]));
      ++outIterJ;
      ++idx;
    }

    // Decrement to get the next principal component
    --kthLargestPrincipalComp;
  }

  // Fill extraneous outputs with zero
  for (; j < numberOfOutputs; j++)
  {
    region = this->GetOutput(j)->GetRequestedRegion();
    OutputIterator outIterJ(this->GetOutput(j), region);

    outIterJ.GoToBegin();
    while (!outIterJ.IsAtEnd())
    {
      outIterJ.Set(0);
      ++outIterJ;
    }
  }

  // Delete eigenvectors at the end of generateData to free memory
  if (this->GetReleaseDataFlag())
  {
    m_EigenVectors.set_size(0, 0);
  }
} // end Generate data

/**
 * Set the number of required principal components
 */
template <typename TInputImage, typename TOutputImage>
void
ImagePCAShapeModelEstimator<TInputImage, TOutputImage>::SetNumberOfPrincipalComponentsRequired(unsigned int n)
{
  if (m_NumberOfPrincipalComponentsRequired != n)
  {
    m_NumberOfPrincipalComponentsRequired = n;

    this->Modified();

    // Modify the required number of outputs ( 1 extra for the mean image )
    this->SetNumberOfRequiredOutputs(m_NumberOfPrincipalComponentsRequired + 1);

    auto         numberOfOutputs = static_cast<unsigned int>(this->GetNumberOfIndexedOutputs());
    unsigned int idx;

    if (numberOfOutputs < m_NumberOfPrincipalComponentsRequired + 1)
    {
      // Make and add extra outputs
      for (idx = numberOfOutputs; idx <= m_NumberOfPrincipalComponentsRequired; idx++)
      {
        typename DataObject::Pointer output = this->MakeOutput(idx);
        this->SetNthOutput(idx, output.GetPointer());
      }
    }
    else if (numberOfOutputs > m_NumberOfPrincipalComponentsRequired + 1)
    {
      // Remove the extra outputs
      for (idx = numberOfOutputs - 1; idx >= m_NumberOfPrincipalComponentsRequired + 1; idx--)
      {
        this->RemoveOutput(idx);
      }
    }
  }
}

/**
 * Set the number of training images.
 */
template <typename TInputImage, typename TOutputImage>
void
ImagePCAShapeModelEstimator<TInputImage, TOutputImage>::SetNumberOfTrainingImages(unsigned int n)
{
  if (m_NumberOfTrainingImages != n)
  {
    m_NumberOfTrainingImages = n;

    this->Modified();

    // Modify the required number of inputs
    this->SetNumberOfRequiredInputs(m_NumberOfTrainingImages);
  }
}

/**-----------------------------------------------------------------
 * Takes a set of training images and returns the means
 * and variance of the various classes defined in the
 * training set.
 */
template <typename TInputImage, typename TOutputImage>
void
ImagePCAShapeModelEstimator<TInputImage, TOutputImage>::EstimateShapeModels()
{
  this->CalculateInnerProduct();

  this->EstimatePCAShapeModelParameters();
} // end EstimateShapeModels

/**
 * Calculate the inner product between the input training vector
 * where each image is treated as a vector of n-elements
 */
template <typename TInputImage, typename TOutputImage>
void
ImagePCAShapeModelEstimator<TInputImage, TOutputImage>::CalculateInnerProduct()
{
  // Get the pointers to the input images and initialize the iterators
  // We use dynamic_cast since inputs are stored as DataObjects.  The
  // ImageToImageFilter::GetInput(int) always returns a pointer to a
  // TInputImage1 so it cannot be used for the second input.

  InputImagePointerArray inputImagePointerArray(m_NumberOfTrainingImages);

  m_InputImageIteratorArray.resize(m_NumberOfTrainingImages);

  for (unsigned int i = 0; i < m_NumberOfTrainingImages; i++)
  {
    InputImageConstPointer inputImagePtr = dynamic_cast<const TInputImage *>(ProcessObject::GetInput(i));

    inputImagePointerArray[i] = inputImagePtr;

    InputImageConstIterator inputImageIt(inputImagePtr, inputImagePtr->GetBufferedRegion());

    m_InputImageIteratorArray[i] = inputImageIt;

    m_InputImageIteratorArray[i].GoToBegin();
  }

  //-------------------------------------------------------------------
  // Set up the matrix to hold the inner product and the means from the
  // training data
  //-------------------------------------------------------------------
  m_InputImageSize = (inputImagePointerArray[0])->GetBufferedRegion().GetSize();

  m_NumberOfPixels = 1;
  for (unsigned int i = 0; i < InputImageDimension; i++)
  {
    m_NumberOfPixels *= m_InputImageSize[i];
  }

  //-------------------------------------------------------------------------
  // Calculate the Means
  //-------------------------------------------------------------------------
  m_Means.set_size(m_NumberOfPixels);
  m_Means.fill(0);

  InputImageConstIterator tempImageItA;

  for (unsigned int img_number = 0; img_number < m_NumberOfTrainingImages; img_number++)
  {
    tempImageItA = m_InputImageIteratorArray[img_number];

    for (unsigned int band_x = 0; band_x < m_NumberOfPixels; band_x++)
    {
      m_Means[band_x] += tempImageItA.Get();
      ++tempImageItA;
    }
  } // end: looping through the image
  //-------------------------------------------------------------------------

  m_Means /= m_NumberOfTrainingImages;

  //-------------------------------------------------------------------------
  // Calculate the inner product
  //-------------------------------------------------------------------------
  m_InnerProduct.set_size(m_NumberOfTrainingImages, m_NumberOfTrainingImages);
  m_InnerProduct.fill(0);

  InputImageConstIterator tempImageItB;

  //-------------------------------------------------------------------------
  for (unsigned int band_x = 0; band_x < m_NumberOfTrainingImages; band_x++)
  {
    for (unsigned int band_y = 0; band_y <= band_x; band_y++)
    {
      // Pointer to the vector (in original matrix)
      tempImageItA = m_InputImageIteratorArray[band_x];

      // Pointer to the vector in the transposed matrix
      tempImageItB = m_InputImageIteratorArray[band_y];

      for (unsigned int pix_number = 0; pix_number < m_NumberOfPixels; pix_number++)
      {
        m_InnerProduct[band_x][band_y] +=
          (tempImageItA.Get() - m_Means[pix_number]) * (tempImageItB.Get() - m_Means[pix_number]);

        ++tempImageItA;
        ++tempImageItB;
      } // end: looping through the image
    }   // end: band_y loop
  }     // end: band_x loop

  //---------------------------------------------------------------------
  // Fill the rest of the inner product matrix and make it symmetric
  //---------------------------------------------------------------------

  for (unsigned int band_x = 0; band_x < (m_NumberOfTrainingImages - 1); band_x++)
  {
    for (unsigned int band_y = band_x + 1; band_y < m_NumberOfTrainingImages; band_y++)
    {
      m_InnerProduct[band_x][band_y] = m_InnerProduct[band_y][band_x];
    } // end band_y loop
  }   // end band_x loop

  if ((m_NumberOfTrainingImages - 1) != 0)
  {
    m_InnerProduct /= (m_NumberOfTrainingImages - 1);
  }
  else
  {
    m_InnerProduct.fill(0);
  }
} // end CalculateInnerProduct

/*-----------------------------------------------------------------
 *Estimage shape models using PCA.
 *-----------------------------------------------------------------
 */
template <typename TInputImage, typename TOutputImage>
void
ImagePCAShapeModelEstimator<TInputImage, TOutputImage>::EstimatePCAShapeModelParameters()
{
  MatrixOfDoubleType identityMatrix(m_NumberOfTrainingImages, m_NumberOfTrainingImages);

  identityMatrix.set_identity();

  vnl_generalized_eigensystem eigenVectors_eigenValues(m_InnerProduct, identityMatrix);

  MatrixOfDoubleType eigenVectorsOfInnerProductMatrix = eigenVectors_eigenValues.V;

  //--------------------------------------------------------------------
  // Calculate the principal shape variations
  //
  // m_EigenVectors capture the principal shape variantions
  // m_EigenValues capture the relative weight of each variation
  // Multiply original image vetors with the eigenVectorsOfInnerProductMatrix
  // to derive the principal shapes.
  //--------------------------------------------------------------------

  m_EigenVectors.set_size(m_NumberOfPixels, m_NumberOfTrainingImages);
  m_EigenVectors.fill(0);

  double                  pix_value;
  InputImageConstIterator tempImageItA;

  for (unsigned int img_number = 0; img_number < m_NumberOfTrainingImages; img_number++)
  {
    tempImageItA = m_InputImageIteratorArray[img_number];
    for (unsigned int pix_number = 0; pix_number < m_NumberOfPixels; pix_number++)
    {
      pix_value = tempImageItA.Get();
      for (unsigned int vec_number = 0; vec_number < m_NumberOfTrainingImages; vec_number++)
      {
        m_EigenVectors[pix_number][vec_number] +=
          (pix_value * eigenVectorsOfInnerProductMatrix[img_number][vec_number]);
      }
      ++tempImageItA;
    }
  }

  m_EigenVectors.normalize_columns();

  m_EigenValues.set_size(m_NumberOfTrainingImages);

  // Extract the diagonal elements into the Eigen value vector
  m_EigenValues = (eigenVectors_eigenValues.D).diagonal();

  // Flip the eigen values since the eigen vectors output
  // is ordered in decending order of their corresponding eigen values.
  m_EigenValues.flip();

  //--------------------------------------------------------------------
  // Normalize the eigen values
  //--------------------------------------------------------------------

  m_EigenVectorNormalizedEnergy = m_EigenValues;
  m_EigenVectorNormalizedEnergy.normalize();
} // end EstimatePCAShapeModelParameters

//-----------------------------------------------------------------
} // namespace itk

#endif
