/*=========================================================================
 *
 *  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
 *
 *         https://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 itkMaskedFFTNormalizedCorrelationImageFilter_h
#define itkMaskedFFTNormalizedCorrelationImageFilter_h

#include "itkImageToImageFilter.h"
#include "itkImage.h"

namespace itk
{
/**
 * \class MaskedFFTNormalizedCorrelationImageFilter
 * \brief Calculate masked normalized cross correlation using FFTs.
 *
 * This filter calculates the masked normalized cross correlation (NCC) of two
 * images under masks using FFTs instead of spatial correlation.
 * It is much faster than spatial correlation
 * for reasonably large structuring elements.
 * This filter is not equivalent to simply masking the images first and then
 * correlating them; the latter approach yields incorrect results because the
 * zeros in the images still affect the metric in the correlation process.
 * This filter implements the masked NCC correctly so that the masked-out regions are completely ignored.
 * The fundamental difference is described in detail in the references below.
 * If the masks are set to images of all ones, the result of this filter is
 * the same as standard NCC.
 *
 * Inputs:
 * Two images are required as inputs, fixedImage and movingImage,
 * and two are optional, fixedMask and movingMask.
 * In the context of correlation, inputs are often defined as: "image"
 * and "template".  In this filter, the fixedImage plays the role of the
 * image, and the movingImage plays the role of the template.
 * However, this filter is capable of correlating any two images and
 * is not restricted to small movingImages (templates).
 * In the fixedMask and movingMask, non-zero positive values
 * indicate locations of useful information in the corresponding image,
 * whereas zero and negative values indicate locations that should be
 * masked out (ignored).
 * Internally, the masks are converted to have values of only 0 and 1.
 * For each optional mask that is not set, the filter
 * internally creates an image of ones, which is equivalent to not
 * masking the image.  Thus, if both masks are not set, the result
 * will be equivalent to unmasked NCC.
 * For example, if only a mask for the fixed image is needed, the
 * movingMask can either not be set or can be set to an image of ones.
 *
 * Optional parameters:
 * The RequiredNumberOfOverlappingPixels enables the user to specify the minimum number of voxels
 * of the two masks that must overlap; any location in the correlation map that results
 * from fewer than this number of voxels will be set to zero.
 * Larger values zero-out pixels on a larger border around the correlation image.
 * Thus, larger values remove less stable computations but also limit the capture range.
 * If RequiredNumberOfOverlappingPixels is set to 0, the default, no zeroing will take place.
 *
 * The RequiredFractionOfOverlappingPixels enables the user to specify a fraction of the maximum
 * number of overlapping pixels that need to overlap; any location in the correlation map that
 * results from fewer than the product of this fraction and the internally computed maximum
 * number of overlapping pixels will be set to zero.
 * The value ranges between 0.0 and 1.0.
 * This is very useful when the user does does not know beforehand the maximum number of pixels
 * of the masks that will overlap.
 * For example, when the masks have strange shapes, it is difficult to predict
 * how the correlation of the masks will interact and what the maximum overlap will be.
 * It is also useful when the mask shapes or sizes change because it is relative to the internally
 * computed maximum of the overlap.
 * Larger values zero-out pixels on a larger border around the correlation image.
 * Thus, larger values remove less stable computations but also limit the capture range.
 * Experiments have shown that a value between 0.1 and 0.6 works well for images with significant
 * overlap and between 0.05 and 0.1 for images with little overlap (such as in stitching applications).
 * If RequiredFractionOfOverlappingPixels is set to 0, the default, no zeroing will take place.
 *
 * The user can either specify RequiredNumberOfOverlappingPixels or RequiredFractionOfOverlappingPixels
 * (or both or none).
 * Internally, the number of required pixels resulting from both of these methods is calculated
 * and the one that gives the largest number of pixels is chosen.
 * Since these both default to 0, if a user only sets one, the other is ignored.
 *
 * Image size:
 * fixedImage and movingImage need not be the same size, but fixedMask
 * must be the same size as fixedImage, and movingMask must be the same
 * size as movingImage.
 * Furthermore, whereas some algorithms require that the "template"
 * be smaller than the "image" because of errors in the regions where
 * the two are not fully overlapping, this filter has no such restriction.
 *
 * Image spacing:
 * Since the computations are done in the pixel domain, all input
 * images must have the same spacing.
 *
 * Outputs;
 * The output is an image of RealPixelType that is the masked NCC of
 * the two images and its values range from -1.0 to 1.0.
 * The size of this NCC image is, by definition,
 * size(fixedImage) + size(movingImage) - 1.
 *
 * Example filter usage:
   \code
   using FilterType = itk::MaskedFFTNormalizedCorrelationImageFilter< ShortImageType, DoubleImageType >;
   auto filter = FilterType::New();
   filter->SetFixedImage( fixedImage );
   filter->SetMovingImage( movingImage );
   filter->SetFixedImageMask( fixedMask );
   filter->SetMovingImageMask( movingMask );
   filter->SetRequiredNumberOfOverlappingPixels(20);
   filter->Update();
   \endcode
 *
 * \warning The pixel type of the output image must be of real type
 * (float or double). ConceptChecking is used to enforce the output pixel
 * type. You will get a compilation error if the pixel type of the
 * output image is not float or double.
 *
 * References:
 * 1) D. Padfield. "Masked object registration in the Fourier domain."
 * Transactions on Image Processing.
 * 2) D. Padfield. "Masked FFT registration". In Proc. Computer
 * Vision and Pattern Recognition, 2010.
 *
 * \author: Dirk Padfield, GE Global Research, padfield\@research.ge.com
 * \ingroup ITKConvolution
 *
 * \sphinx
 * \sphinxexample{Filtering/Convolution/NormalizedCorrelationUsingFFTWithMaskImages,Normalized Correlation Using FFT
 With Mask Images For Input Images}
 * \endsphinx
 */

template <typename TInputImage, typename TOutputImage, typename TMaskImage = TInputImage>
class ITK_TEMPLATE_EXPORT MaskedFFTNormalizedCorrelationImageFilter
  : public ImageToImageFilter<TInputImage, TOutputImage>
{
public:
  ITK_DISALLOW_COPY_AND_MOVE(MaskedFFTNormalizedCorrelationImageFilter);

  /** Standard class type aliases. */
  using Self = MaskedFFTNormalizedCorrelationImageFilter;
  using Superclass = ImageToImageFilter<TInputImage, TOutputImage>;
  using Pointer = SmartPointer<Self>;
  using ConstPointer = SmartPointer<const Self>;

  /** Method for creation through the object factory. */
  itkNewMacro(Self);

  /** \see LightObject::GetNameOfClass() */
  itkOverrideGetNameOfClassMacro(MaskedFFTNormalizedCorrelationImageFilter);

  /** Extract some information from the image types.  Dimensionality
   * of the two images is assumed to be the same. */
  static constexpr unsigned int ImageDimension = TOutputImage::ImageDimension;

  /** Extract some information from the image types. */
  using InputImageType = TInputImage;
  using InputRegionType = typename InputImageType::RegionType;
  using InputImagePointer = typename InputImageType::Pointer;
  using InputImageConstPointer = typename InputImageType::ConstPointer;
  using InputSizeType = typename InputImageType::SizeType;
  using SizeValueType = itk::SizeValueType;

  using OutputImageType = TOutputImage;
  using OutputImagePointer = typename OutputImageType::Pointer;
  using OutputPixelType = typename OutputImageType::PixelType;

  using RealPixelType = OutputPixelType;
  using RealImageType = Image<RealPixelType, ImageDimension>;
  using RealImagePointer = typename RealImageType::Pointer;
  using RealIndexType = typename RealImageType::IndexType;
  using RealSizeType = typename RealImageType::SizeType;
  using RealRegionType = typename RealImageType::RegionType;
  using RealPointType = typename RealImageType::PointType;

  using MaskImageType = TMaskImage;
  using MaskImagePointer = typename MaskImageType::Pointer;

  using FFTImageType = Image<std::complex<RealPixelType>, ImageDimension>;
  using FFTImagePointer = typename FFTImageType::Pointer;

  /** Set and get the fixed image */
  itkSetInputMacro(FixedImage, InputImageType);
  itkGetInputMacro(FixedImage, InputImageType);

  /** Set and get the moving image */
  itkSetInputMacro(MovingImage, InputImageType);
  itkGetInputMacro(MovingImage, InputImageType);

  /** Set and get the fixed mask */
  itkSetInputMacro(FixedImageMask, MaskImageType);
  itkGetInputMacro(FixedImageMask, MaskImageType);

  /** Set and get the moving mask */
  itkSetInputMacro(MovingImageMask, MaskImageType);
  itkGetInputMacro(MovingImageMask, MaskImageType);

  /** Set and get the required number of overlapping pixels */
  itkSetMacro(RequiredNumberOfOverlappingPixels, SizeValueType);
  itkGetMacro(RequiredNumberOfOverlappingPixels, SizeValueType);

  /** Set and get the required fraction of overlapping pixels */
  itkGetMacro(RequiredFractionOfOverlappingPixels, RealPixelType);
  itkSetClampMacro(RequiredFractionOfOverlappingPixels, RealPixelType, 0.0f, 1.0f);

  /** Get the maximum number of overlapping pixels. */
  itkGetMacro(MaximumNumberOfOverlappingPixels, SizeValueType);

#ifdef ITK_USE_CONCEPT_CHECKING
  // Begin concept checking
  itkConceptMacro(OutputPixelTypeIsFloatingPointCheck, (Concept::IsFloatingPoint<OutputPixelType>));
  // End concept checking
#endif

protected:
  MaskedFFTNormalizedCorrelationImageFilter()
  {
    // #0 "FixedImage" required
    Self::SetPrimaryInputName("FixedImage");

    // #1 "MaskImage" required
    Self::AddRequiredInputName("MovingImage", 1);

    // #2 "FixedImageMask" optional
    Self::AddOptionalInputName("FixedImageMask", 2);

    // #3 "MaskImageMask" optional
    Self::AddOptionalInputName("MovingImageMask", 3);

    m_RequiredNumberOfOverlappingPixels = 0;
    m_RequiredFractionOfOverlappingPixels = 0;
    m_MaximumNumberOfOverlappingPixels = 0;
    m_AccumulatedProgress = 0.0;
  }
  ~MaskedFFTNormalizedCorrelationImageFilter() override = default;
  void
  PrintSelf(std::ostream & os, Indent indent) const override;

  /** Overlap the VerifyInputInformation method */
  void
  VerifyInputInformation() ITKv5_CONST override;

  /** Standard pipeline method.*/
  void
  GenerateData() override;

  /** This filter needs a different input requested region than the output
   * requested region.  As such, it needs to provide an
   * implementation for GenerateInputRequestedRegion() in order to inform the
   * pipeline execution model.
   * \sa ProcessObject::GenerateInputRequestedRegion() */
  void
  GenerateInputRequestedRegion() override;

  /** Since the output of this filter is a different
   * size than the input, it must provide an implementation of
   * GenerateOutputInformation.
   * \sa ProcessObject::GenerateOutputRequestedRegion() */
  void
  GenerateOutputInformation() override;

  void
  EnlargeOutputRequestedRegion(DataObject * output) override;

  typename TMaskImage::Pointer
  PreProcessMask(const InputImageType * inputImage, const MaskImageType * inputMask);

  typename TInputImage::Pointer
  PreProcessImage(const InputImageType * inputImage, const MaskImageType * inputMask);

  template <typename LocalInputImageType>
  typename LocalInputImageType::Pointer
  RotateImage(LocalInputImageType * inputImage);

  template <typename LocalInputImageType, typename LocalOutputImageType>
  typename LocalOutputImageType::Pointer
  CalculateForwardFFT(LocalInputImageType * inputImage, InputSizeType & FFTImageSize);

  template <typename LocalInputImageType, typename LocalOutputImageType>
  typename LocalOutputImageType::Pointer
  CalculateInverseFFT(LocalInputImageType * inputImage, RealSizeType & combinedImageSize);

  // Helper math methods.
  template <typename LocalInputImageType, typename LocalOutputImageType>
  typename LocalOutputImageType::Pointer
  ElementProduct(LocalInputImageType * inputImage1, LocalInputImageType * inputImage2);

  template <typename LocalInputImageType>
  typename LocalInputImageType::Pointer
  ElementQuotient(LocalInputImageType * inputImage1, LocalInputImageType * inputImage2);

  template <typename LocalInputImageType>
  typename LocalInputImageType::Pointer
  ElementSubtraction(LocalInputImageType * inputImage1, LocalInputImageType * inputImage2);

  template <typename LocalInputImageType>
  typename LocalInputImageType::Pointer
  ElementPositive(LocalInputImageType * inputImage);

  template <typename LocalInputImageType, typename LocalOutputImageType>
  typename LocalOutputImageType::Pointer
  ElementRound(LocalInputImageType * inputImage);

  // This function factorizes the image size uses factors of 2, 3, and
  // 5.  After this factorization, if there are any remaining values,
  // the function returns this value.
  int
  FactorizeNumber(int n);

  // Find the closest valid dimension above the desired dimension.  This
  // will be a combination of 2s, 3s, and 5s.
  int
  FindClosestValidDimension(int n);

  template <typename LocalInputImageType>
  double
  CalculatePrecisionTolerance(LocalInputImageType * inputImage);

private:
  /** Larger values zero-out pixels on a larger border around the correlation image.
   * Thus, larger values remove less stable computations but also limit the capture range.
   * The default is set to 0. */
  SizeValueType m_RequiredNumberOfOverlappingPixels{};
  /** Similar to m_RequiredNumberOfOverlappingPixels except that the m_RequiredFractionOfOverlappingPixels is multiplied
   * by the m_MaximumNumberOfOverlappingPixels to determine the requiredNumberOfOverlappingPixels. The default is 0. */
  RealPixelType m_RequiredFractionOfOverlappingPixels{};
  /** This is computed internally */
  SizeValueType m_MaximumNumberOfOverlappingPixels{};

  /** This is used for the progress reporter */
  const unsigned int m_TotalForwardAndInverseFFTs{ 12 };
  /** The total accumulated progress */
  float m_AccumulatedProgress{};
};
} // end namespace itk

#ifndef ITK_MANUAL_INSTANTIATION
#  include "itkMaskedFFTNormalizedCorrelationImageFilter.hxx"
#endif

#endif
