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

#include "itkMultiResolutionImageRegistrationMethod.h"
#include "itkRecursiveMultiResolutionPyramidImageFilter.h"

namespace itk
{
/**
 * Constructor
 */
template <typename TFixedImage, typename TMovingImage>
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::MultiResolutionImageRegistrationMethod()
{
  this->SetNumberOfRequiredOutputs(1); // for the Transform

  m_FixedImage = nullptr;   // has to be provided by the user.
  m_MovingImage = nullptr;  // has to be provided by the user.
  m_Transform = nullptr;    // has to be provided by the user.
  m_Interpolator = nullptr; // has to be provided by the user.
  m_Metric = nullptr;       // has to be provided by the user.
  m_Optimizer = nullptr;    // has to be provided by the user.

  // Use MultiResolutionPyramidImageFilter as the default
  // image pyramids.
  m_FixedImagePyramid = FixedImagePyramidType::New();
  m_MovingImagePyramid = MovingImagePyramidType::New();

  m_NumberOfLevels = 1;
  m_CurrentLevel = 0;

  m_Stop = false;

  m_ScheduleSpecified = false;
  m_NumberOfLevelsSpecified = false;

  m_InitialTransformParameters = ParametersType(1);
  m_InitialTransformParametersOfNextLevel = ParametersType(1);
  m_LastTransformParameters = ParametersType(1);

  m_InitialTransformParameters.Fill(0.0f);
  m_InitialTransformParametersOfNextLevel.Fill(0.0f);
  m_LastTransformParameters.Fill(0.0f);

  TransformOutputPointer transformDecorator =
    itkDynamicCastInDebugMode<TransformOutputType *>(this->MakeOutput(0).GetPointer());

  this->ProcessObject::SetNthOutput(0, transformDecorator.GetPointer());
}

/*
 * Initialize by setting the interconnects between components.
 */
template <typename TFixedImage, typename TMovingImage>
void
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::Initialize()
{
  // Sanity checks
  if (!m_Metric)
  {
    itkExceptionMacro(<< "Metric is not present");
  }

  if (!m_Optimizer)
  {
    itkExceptionMacro(<< "Optimizer is not present");
  }

  if (!m_Transform)
  {
    itkExceptionMacro(<< "Transform is not present");
  }

  if (!m_Interpolator)
  {
    itkExceptionMacro(<< "Interpolator is not present");
  }

  // Setup the metric
  m_Metric->SetMovingImage(m_MovingImagePyramid->GetOutput(m_CurrentLevel));
  m_Metric->SetFixedImage(m_FixedImagePyramid->GetOutput(m_CurrentLevel));
  m_Metric->SetTransform(m_Transform);
  m_Metric->SetInterpolator(m_Interpolator);
  m_Metric->SetFixedImageRegion(m_FixedImageRegionPyramid[m_CurrentLevel]);
  m_Metric->Initialize();

  // Setup the optimizer
  m_Optimizer->SetCostFunction(m_Metric);
  m_Optimizer->SetInitialPosition(m_InitialTransformParametersOfNextLevel);

  //
  // Connect the transform to the Decorator.
  //
  auto * transformOutput = static_cast<TransformOutputType *>(this->ProcessObject::GetOutput(0));

  transformOutput->Set(m_Transform);
}

/*
 * Stop the Registration Process
 */
template <typename TFixedImage, typename TMovingImage>
void
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::StopRegistration()
{
  m_Stop = true;
}

/**
 * Set the schedules for the fixed and moving image pyramid
 */
template <typename TFixedImage, typename TMovingImage>
void
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::SetSchedules(
  const ScheduleType & fixedImagePyramidSchedule,
  const ScheduleType & movingImagePyramidSchedule)
{
  if (m_NumberOfLevelsSpecified)
  {
    itkExceptionMacro("SetSchedules should not be used "
                      << "if numberOfLevelves are specified using SetNumberOfLevels");
  }
  m_FixedImagePyramidSchedule = fixedImagePyramidSchedule;
  m_MovingImagePyramidSchedule = movingImagePyramidSchedule;
  m_ScheduleSpecified = true;

  // Set the number of levels based on the pyramid schedule specified
  if (m_FixedImagePyramidSchedule.rows() != m_MovingImagePyramidSchedule.rows())
  {
    itkExceptionMacro("The specified schedules contain unequal number of levels");
  }
  else
  {
    m_NumberOfLevels = m_FixedImagePyramidSchedule.rows();
  }

  this->Modified();
}

/**
 * Set the number of levels
 */
template <typename TFixedImage, typename TMovingImage>
void
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::SetNumberOfLevels(SizeValueType numberOfLevels)
{
  if (m_ScheduleSpecified)
  {
    itkExceptionMacro("SetNumberOfLevels should not be used "
                      << "if schedules have been specified using SetSchedules method ");
  }

  m_NumberOfLevels = numberOfLevels;
  m_NumberOfLevelsSpecified = true;
  this->Modified();
}

/**
 * Stop the Registration Process
 */
template <typename TFixedImage, typename TMovingImage>
void
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::PreparePyramids()
{
  if (!m_Transform)
  {
    itkExceptionMacro(<< "Transform is not present");
  }

  m_InitialTransformParametersOfNextLevel = m_InitialTransformParameters;

  if (m_InitialTransformParametersOfNextLevel.Size() != m_Transform->GetNumberOfParameters())
  {
    itkExceptionMacro(<< "Size mismatch between initial parameter and transform");
  }

  // Sanity checks
  if (!m_FixedImage)
  {
    itkExceptionMacro(<< "FixedImage is not present");
  }

  if (!m_MovingImage)
  {
    itkExceptionMacro(<< "MovingImage is not present");
  }

  if (!m_FixedImagePyramid)
  {
    itkExceptionMacro(<< "Fixed image pyramid is not present");
  }

  if (!m_MovingImagePyramid)
  {
    itkExceptionMacro(<< "Moving image pyramid is not present");
  }

  // Setup the fixed and moving image pyramid
  if (m_NumberOfLevelsSpecified)
  {
    m_FixedImagePyramid->SetNumberOfLevels(m_NumberOfLevels);
    m_MovingImagePyramid->SetNumberOfLevels(m_NumberOfLevels);
  }

  if (m_ScheduleSpecified)
  {
    m_FixedImagePyramid->SetNumberOfLevels(m_FixedImagePyramidSchedule.rows());
    m_FixedImagePyramid->SetSchedule(m_FixedImagePyramidSchedule);

    m_MovingImagePyramid->SetNumberOfLevels(m_MovingImagePyramidSchedule.rows());
    m_MovingImagePyramid->SetSchedule(m_MovingImagePyramidSchedule);
  }

  m_FixedImagePyramid->SetInput(m_FixedImage);
  m_FixedImagePyramid->UpdateLargestPossibleRegion();

  // Setup the moving image pyramid
  m_MovingImagePyramid->SetInput(m_MovingImage);
  m_MovingImagePyramid->UpdateLargestPossibleRegion();

  using SizeType = typename FixedImageRegionType::SizeType;
  using IndexType = typename FixedImageRegionType::IndexType;

  ScheduleType schedule = m_FixedImagePyramid->GetSchedule();
  itkDebugMacro(<< "FixedImage schedule: " << schedule);

  ScheduleType movingschedule = m_MovingImagePyramid->GetSchedule();
  itkDebugMacro(<< "MovingImage schedule: " << movingschedule);

  SizeType  inputSize = m_FixedImageRegion.GetSize();
  IndexType inputStart = m_FixedImageRegion.GetIndex();

  const SizeValueType numberOfLevels = m_FixedImagePyramid->GetNumberOfLevels();

  m_FixedImageRegionPyramid.reserve(numberOfLevels);
  m_FixedImageRegionPyramid.resize(numberOfLevels);

  // Compute the FixedImageRegion corresponding to each level of the
  // pyramid. This uses the same algorithm of the ShrinkImageFilter
  // since the regions should be compatible.
  for (unsigned int level = 0; level < numberOfLevels; level++)
  {
    SizeType  size;
    IndexType start;
    for (unsigned int dim = 0; dim < TFixedImage::ImageDimension; dim++)
    {
      const auto scaleFactor = static_cast<float>(schedule[level][dim]);

      size[dim] =
        static_cast<typename SizeType::SizeValueType>(std::floor(static_cast<float>(inputSize[dim]) / scaleFactor));
      if (size[dim] < 1)
      {
        size[dim] = 1;
      }

      start[dim] =
        static_cast<typename IndexType::IndexValueType>(std::ceil(static_cast<float>(inputStart[dim]) / scaleFactor));
    }
    m_FixedImageRegionPyramid[level].SetSize(size);
    m_FixedImageRegionPyramid[level].SetIndex(start);
  }
}

/*
 * PrintSelf
 */
template <typename TFixedImage, typename TMovingImage>
void
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::PrintSelf(std::ostream & os, Indent indent) const
{
  Superclass::PrintSelf(os, indent);
  os << indent << "Metric: " << m_Metric.GetPointer() << std::endl;
  os << indent << "Optimizer: " << m_Optimizer.GetPointer() << std::endl;
  os << indent << "Transform: " << m_Transform.GetPointer() << std::endl;
  os << indent << "Interpolator: " << m_Interpolator.GetPointer() << std::endl;
  os << indent << "FixedImage: " << m_FixedImage.GetPointer() << std::endl;
  os << indent << "MovingImage: " << m_MovingImage.GetPointer() << std::endl;
  os << indent << "FixedImagePyramid: ";
  os << m_FixedImagePyramid.GetPointer() << std::endl;
  os << indent << "MovingImagePyramid: ";
  os << m_MovingImagePyramid.GetPointer() << std::endl;

  os << indent << "NumberOfLevels: ";
  os << m_NumberOfLevels << std::endl;

  os << indent << "CurrentLevel: ";
  os << m_CurrentLevel << std::endl;

  os << indent << "InitialTransformParameters: ";
  os << m_InitialTransformParameters << std::endl;
  os << indent << "InitialTransformParametersOfNextLevel: ";
  os << m_InitialTransformParametersOfNextLevel << std::endl;
  os << indent << "LastTransformParameters: ";
  os << m_LastTransformParameters << std::endl;
  os << indent << "FixedImageRegion: ";
  os << m_FixedImageRegion << std::endl;
  for (unsigned int level = 0; level < m_FixedImageRegionPyramid.size(); level++)
  {
    os << indent << "FixedImageRegion at level " << level << ": ";
    os << m_FixedImageRegionPyramid[level] << std::endl;
  }
  os << indent << "FixedImagePyramidSchedule : " << std::endl;
  os << m_FixedImagePyramidSchedule << std::endl;
  os << indent << "MovingImagePyramidSchedule : " << std::endl;
  os << m_MovingImagePyramidSchedule << std::endl;
}

/*
 * Generate Data
 */
template <typename TFixedImage, typename TMovingImage>
void
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::GenerateData()
{
  m_Stop = false;

  this->PreparePyramids();

  for (m_CurrentLevel = 0; m_CurrentLevel < m_NumberOfLevels; m_CurrentLevel++)
  {
    // Invoke an iteration event.
    // This allows a UI to reset any of the components between
    // resolution level.
    this->InvokeEvent(MultiResolutionIterationEvent());

    // Check if there has been a stop request
    if (m_Stop)
    {
      break;
    }

    try
    {
      // initialize the interconnects between components
      this->Initialize();
    }
    catch (ExceptionObject & err)
    {
      m_LastTransformParameters = ParametersType(1);
      m_LastTransformParameters.Fill(0.0f);

      // pass exception to caller
      throw err;
    }

    try
    {
      // do the optimization
      m_Optimizer->StartOptimization();
    }
    catch (ExceptionObject & err)
    {
      // An error has occurred in the optimization.
      // Update the parameters
      m_LastTransformParameters = m_Optimizer->GetCurrentPosition();

      // Pass exception to caller
      throw err;
    }

    // get the results
    m_LastTransformParameters = m_Optimizer->GetCurrentPosition();
    m_Transform->SetParameters(m_LastTransformParameters);

    // setup the initial parameters for next level
    if (m_CurrentLevel < m_NumberOfLevels - 1)
    {
      m_InitialTransformParametersOfNextLevel = m_LastTransformParameters;
    }
  }
}

template <typename TFixedImage, typename TMovingImage>
ModifiedTimeType
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::GetMTime() const
{
  ModifiedTimeType mtime = Superclass::GetMTime();
  ModifiedTimeType m;

  // Some of the following should be removed once ivars are put in the
  // input and output lists

  if (m_Transform)
  {
    m = m_Transform->GetMTime();
    mtime = (m > mtime ? m : mtime);
  }

  if (m_Interpolator)
  {
    m = m_Interpolator->GetMTime();
    mtime = (m > mtime ? m : mtime);
  }

  if (m_Metric)
  {
    m = m_Metric->GetMTime();
    mtime = (m > mtime ? m : mtime);
  }

  if (m_Optimizer)
  {
    m = m_Optimizer->GetMTime();
    mtime = (m > mtime ? m : mtime);
  }

  if (m_FixedImage)
  {
    m = m_FixedImage->GetMTime();
    mtime = (m > mtime ? m : mtime);
  }

  if (m_MovingImage)
  {
    m = m_MovingImage->GetMTime();
    mtime = (m > mtime ? m : mtime);
  }

  return mtime;
}

/*
 *  Get Output
 */
template <typename TFixedImage, typename TMovingImage>
const typename MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::TransformOutputType *
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::GetOutput() const
{
  return static_cast<const TransformOutputType *>(this->ProcessObject::GetOutput(0));
}

template <typename TFixedImage, typename TMovingImage>
DataObject::Pointer
MultiResolutionImageRegistrationMethod<TFixedImage, TMovingImage>::MakeOutput(DataObjectPointerArraySizeType output)
{
  if (output > 0)
  {
    itkExceptionMacro("MakeOutput request for an output number larger than the expected number of outputs.");
  }
  return TransformOutputType::New().GetPointer();
}
} // end namespace itk

#endif
