/*=========================================================================
 *
 *  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 itkImageGaussianModelEstimator_hxx
#define itkImageGaussianModelEstimator_hxx

#include "itkImageGaussianModelEstimator.h"
#include "itkMath.h"
#include "itkNumericTraits.h"

namespace itk
{
template <typename TInputImage, typename TMembershipFunction, typename TTrainingImage>
ImageGaussianModelEstimator<TInputImage, TMembershipFunction, TTrainingImage>::ImageGaussianModelEstimator()

  = default;

template <typename TInputImage, typename TMembershipFunction, typename TTrainingImage>
ImageGaussianModelEstimator<TInputImage, TMembershipFunction, TTrainingImage>::~ImageGaussianModelEstimator()
{
  delete[] m_Covariance;
}

/**
 * PrintSelf
 */
template <typename TInputImage, typename TMembershipFunction, typename TTrainingImage>
void
ImageGaussianModelEstimator<TInputImage, TMembershipFunction, TTrainingImage>::PrintSelf(std::ostream & os,
                                                                                         Indent         indent) const
{
  os << indent << "                   " << std::endl;
  os << indent << "Gaussian Models generated from the training data." << std::endl;
  os << indent << "TrainingImage: ";
  os << m_TrainingImage.GetPointer() << std::endl;
  os << indent << "Results printed in the superclass " << std::endl;
  os << indent << "                   " << std::endl;

  Superclass::PrintSelf(os, indent);
} // end PrintSelf

/**
 * Generate data (start the model building process)
 */
template <typename TInputImage, typename TMembershipFunction, typename TTrainingImage>
void
ImageGaussianModelEstimator<TInputImage, TMembershipFunction, TTrainingImage>::GenerateData()
{
  this->EstimateModels();
} // end Generate data

// 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 TMembershipFunction, typename TTrainingImage>
void
ImageGaussianModelEstimator<TInputImage, TMembershipFunction, TTrainingImage>::EstimateModels()
{
  // Do some error checking
  InputImageConstPointer inputImage = this->GetInputImage();

  // Check if the training and input image dimensions are the same
  if ((int)(TInputImage::ImageDimension) != (int)(TTrainingImage::ImageDimension))
  {
    throw ExceptionObject(__FILE__, __LINE__, "Training and input image dimensions are not the same.", ITK_LOCATION);
  }

  InputImageSizeType inputImageSize = inputImage->GetBufferedRegion().GetSize();

  TrainingImageConstPointer trainingImage = this->GetTrainingImage();

  using TrainingImageSizeType = InputImageSizeType;
  TrainingImageSizeType trainingImageSize = trainingImage->GetBufferedRegion().GetSize();

  // Check if size of the two inputs are the same
  for (unsigned int i = 0; i < TInputImage::ImageDimension; i++)
  {
    if (inputImageSize[i] != trainingImageSize[i])
    {
      throw ExceptionObject(
        __FILE__, __LINE__, "Input image size is not the same as the training image size.", ITK_LOCATION);
    }
  }

  //-------------------------------------------------------------------
  // Set up the gaussian membership calculators
  //-------------------------------------------------------------------

  unsigned int numberOfModels = this->GetNumberOfModels();

  //-------------------------------------------------------------------
  // Call local function to estimate mean variances of the various
  // class labels in the training set.
  // The statistics class functions have not been used since all the
  // class statistics are calculated simultaneously here.
  //-------------------------------------------------------------------

  this->EstimateGaussianModelParameters();

  //-------------------------------------------------------------------
  // Populate the membership functions for all the classes
  //-------------------------------------------------------------------
  MembershipFunctionPointer                             membershipFunction;
  typename MembershipFunctionType::MeanVectorType       tmean;
  typename MembershipFunctionType::CovarianceMatrixType tcov;

  NumericTraits<typename MembershipFunctionType::MeanVectorType>::SetLength(tmean, VectorDimension);
  for (unsigned int classIndex = 0; classIndex < numberOfModels; classIndex++)
  {
    membershipFunction = TMembershipFunction::New();

    // Convert to the datatype used for the mean
    for (unsigned int i = 0; i < VectorDimension; ++i)
    {
      tmean[i] = m_Means.get(classIndex, i);
    }
    membershipFunction->SetMean(tmean);

    tcov = m_Covariance[classIndex]; // convert cov for membership fn
    membershipFunction->SetCovariance(tcov);

    this->AddMembershipFunction(membershipFunction);
  }
} // end train classifier

template <typename TInputImage, typename TMembershipFunction, typename TTrainingImage>
void
ImageGaussianModelEstimator<TInputImage, TMembershipFunction, TTrainingImage>::EstimateGaussianModelParameters()
{
  // Set the iterators and the pixel type definition for the input image
  InputImageConstPointer  inputImage = this->GetInputImage();
  InputImageConstIterator inIt(inputImage, inputImage->GetBufferedRegion());

  //-------------------------------------------------------------------

  //-------------------------------------------------------------------
  // Set the iterators and the pixel type definition for the training image
  TrainingImageConstPointer trainingImage = this->GetTrainingImage();

  TrainingImageConstIterator trainingImageIt(trainingImage, trainingImage->GetBufferedRegion());

  //-------------------------------------------------------------------

  unsigned int numberOfModels = (this->GetNumberOfModels());

  //-------------------------------------------------------------------
  // Set up the matrices to hold the means and the covariance for the
  // training data

  m_Means.set_size(numberOfModels, VectorDimension);
  m_Means.fill(0);

  m_NumberOfSamples.set_size(numberOfModels, 1);
  m_NumberOfSamples.fill(0);

  // delete previous allocation first
  delete[] m_Covariance;
  // Number of covariance matrices are equal to the number of classes
  m_Covariance = (MatrixType *)new MatrixType[numberOfModels];

  for (unsigned int i = 0; i < numberOfModels; i++)
  {
    m_Covariance[i].set_size(VectorDimension, VectorDimension);
    m_Covariance[i].fill(0);
  }

  for (inIt.GoToBegin(); !inIt.IsAtEnd(); ++inIt, ++trainingImageIt)
  {
    auto classIndex = (unsigned int)trainingImageIt.Get();

    // Training data assumed =1 band; also the class indices go
    // from 1, 2, ..., n while the corresponding memory goes from
    // 0, 1, ..., n-1.

    // Ensure that the training data is labelled appropriately
    if (classIndex > numberOfModels)
    {
      throw ExceptionObject(__FILE__, __LINE__);
    }

    if (classIndex > 0)
    {
      m_NumberOfSamples[classIndex][0] += 1;
      InputImagePixelType inImgVec = inIt.Get();

      for (unsigned int band_x = 0; band_x < VectorDimension; band_x++)
      {
        m_Means[classIndex][band_x] += inImgVec[band_x];
        for (unsigned int band_y = 0; band_y <= band_x; band_y++)
        {
          m_Covariance[classIndex][band_x][band_y] += inImgVec[band_x] * inImgVec[band_y];
        }
      }
    }
  } // end for

  // Loop through the classes to calculate the means and covariance
  for (unsigned int classIndex = 0; classIndex < numberOfModels; classIndex++)
  {
    if (Math::NotAlmostEquals(m_NumberOfSamples[classIndex][0], 0.0))
    {
      for (unsigned int i = 0; i < VectorDimension; i++)
      {
        m_Means[classIndex][i] /= m_NumberOfSamples[classIndex][0];
      }
    } // end if

    else
    {
      for (unsigned int i = 0; i < VectorDimension; i++)
      {
        m_Means[classIndex][i] = 0;
      }
    } // end else

    if (Math::NotAlmostEquals((m_NumberOfSamples[classIndex][0] - 1), 0.0))
    {
      for (unsigned int band_x = 0; band_x < VectorDimension; band_x++)
      {
        for (unsigned int band_y = 0; band_y <= band_x; band_y++)
        {
          m_Covariance[classIndex][band_x][band_y] /= (m_NumberOfSamples[classIndex][0] - 1);
        } // end for band_y loop
      }   // end for band_x loop
    }     // end if

    else
    {
      for (unsigned int band_x = 0; band_x < VectorDimension; band_x++)
      {
        for (unsigned int band_y = 0; band_y <= band_x; band_y++)
        {
          m_Covariance[classIndex][band_x][band_y] = 0;
        }
      }
    } // end else

    MatrixType tempMeanSq;
    tempMeanSq.set_size(VectorDimension, VectorDimension);
    tempMeanSq.fill(0);

    for (unsigned int band_x = 0; band_x < VectorDimension; band_x++)
    {
      for (unsigned int band_y = 0; band_y <= band_x; band_y++)
      {
        tempMeanSq[band_x][band_y] = m_Means[classIndex][band_x] * m_Means[classIndex][band_y];
      }
    } // end for band_x loop

    if (Math::NotAlmostEquals((m_NumberOfSamples[classIndex][0] - 1), 0.0))
    {
      tempMeanSq *= (m_NumberOfSamples[classIndex][0] / (m_NumberOfSamples[classIndex][0] - 1));
    }
    m_Covariance[classIndex] -= tempMeanSq;

    // Fill the rest of the covairance matrix and make it symmetric
    if (m_NumberOfSamples[classIndex][0] > 0)
    {
      auto lastInX = (unsigned int)(VectorDimension - 1);
      auto upperY = (unsigned int)VectorDimension;
      for (unsigned int band_x = 0; band_x < lastInX; band_x++)
      {
        for (unsigned int band_y = band_x + 1; band_y < upperY; band_y++)
        {
          m_Covariance[classIndex][band_x][band_y] = m_Covariance[classIndex][band_y][band_x];
        } // end band_y loop
      }   // end band_x loop
    }     // end if loop
  }       // end class index loop
} // end EstimateGaussianModelParameters
} // namespace itk

#endif
