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

#include "itkBinaryThresholdImageFilter.h"
#include "itkConnectedComponentImageFilter.h"
#include "itkExtractImageFilter.h"
#include "itkImageToImageFilter.h"
#include <unordered_map>

namespace itk
{
/** \class MorphologicalContourInterpolator
 *
 *  \brief Interpolates contours between slices. Based on a paper by Albu et al.
 *
 *  \par Inputs and Outputs
 *  This is an image-to-image filter. The dimensionality is 3D or higher.
 *  Input contains an image with some slices segmented, usually manually.
 *  The output has all in-between slices also segmented.
 *
 *  \par Parameters
 *  Slices are detected at positions where a pixel exists with same labeled
 *  neighbors in slice and only clear (0) neighbors perpendicular to the slice.
 *  If default behaviour is unwanted, contour indices can be set
 *  by enabling UseCustomSlicePositions and calling SetLabeledSliceIndices.
 *
 *  Filter can be restricted to run along only one axis, and/or to interpolate
 *  just one label.
 *
 *  Since optimal alignment between slices would require exhaustive search,
 *  the default is to use heuristic (breadth first search starting from centroid).
 *
 *  There is also an alternative algorithm based on distance transform approach.
 *  It is slightly faster, but it can jump across a twisty shape (not geodesic).
 *
 *  Reference:
 *  Albu AB, Beugeling T, Laurendeau D. "A morphology-based approach for
 *  interslice interpolation of anatomical slices from volumetric images."
 *  IEEE Trans Biomed Eng. 2008 Aug;55(8):2022-38. DOI:10.1109/TBME.2008.921158
 *
 *  Acknowledgement:
 *  This work is supported by NIH grant R01 EB014346, "Continued development
 *  and maintenance of the ITK-SNAP 3D image segmentation software."
 *
 * \ingroup MorphologicalContourInterpolation
 */
template <typename TImage>
class MorphologicalContourInterpolator : public ImageToImageFilter<TImage, TImage>
{
  template <typename T>
  friend class MorphologicalContourInterpolatorParallelInvoker;

public:
  ITK_DISALLOW_COPY_AND_MOVE(MorphologicalContourInterpolator);

  /** Standard class type alias. */
  using Self = MorphologicalContourInterpolator;
  using Superclass = ImageToImageFilter<TImage, TImage>;
  using Pointer = SmartPointer<Self>;
  using SliceType = Image<typename TImage::PixelType, TImage::ImageDimension - 1>;

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

  /** Run-time type information (and related methods). */
  itkOverrideGetNameOfClassMacro(MorphologicalContourInterpolator);

  /** Interpolate only this label. Interpolates all labels if set to 0 (default). */
  itkSetMacro(Label, typename TImage::PixelType);

  /** Which label is interpolated. 0 means all labels (default). */
  itkGetMacro(Label, typename TImage::PixelType);

  /** Which label is interpolated. 0 means all labels (default). */
  itkGetConstMacro(Label, typename TImage::PixelType);

  /** Interpolate only along this axis. Interpolates along all axes if set to -1 (default). */
  itkSetMacro(Axis, int);

  /** Axis of interpolation. -1 means interpolation along all axes (default). */
  itkGetMacro(Axis, int);

  /** Axis of interpolation. -1 means interpolation along all axes (default). */
  itkGetConstMacro(Axis, int);

  /** Heuristic alignment of regions for interpolation is faster than optimal alignment.
   *   Default is heuristic. */
  itkSetMacro(HeuristicAlignment, bool);

  /** Heuristic alignment of regions for interpolation is faster than optimal alignment.
   *   Default is heuristic. */
  itkGetMacro(HeuristicAlignment, bool);

  /** Heuristic alignment of regions for interpolation is faster than optimal alignment.
   *   Default is heuristic. */
  itkGetConstMacro(HeuristicAlignment, bool);

  /** Using distance transform instead of repeated dilations to calculate
   *   the median contour is slightly faster, but produces lower quality interpolations.
   *   Default is OFF (that is, use repeated dilations). */
  itkSetMacro(UseDistanceTransform, bool);

  /** Using distance transform instead of repeated dilations to calculate
   *   the median contour is slightly faster, but produces lower quality interpolations.
   *   Default is OFF (that is, use repeated dilations). */
  itkGetMacro(UseDistanceTransform, bool);

  /** Using distance transform instead of repeated dilations to calculate
   *   the median contour is slightly faster, but produces lower quality interpolations.
   *   Default is OFF (that is, use repeated dilations). */
  itkGetConstMacro(UseDistanceTransform, bool);

  /** Use custom slice positions (not slice auto-detection).
   *   SetLabeledSliceIndices has to be called prior to Update(). */
  itkSetMacro(UseCustomSlicePositions, bool);

  /** Use custom slice positions (not slice auto-detection).
   *   SetLabeledSliceIndices has to be called prior to Update(). */
  itkGetMacro(UseCustomSlicePositions, bool);

  /** Use custom slice positions (not slice auto-detection).
   *   SetLabeledSliceIndices has to be called prior to Update(). */
  itkGetConstMacro(UseCustomSlicePositions, bool);

  /** Perform extrapolation for branch extremities.
   * Branch extremities are defined as regions having no overlap with any region in the next slice.
   * Extrapolation helps generate smooth surface closings.
   *    Default is ON
   * */
  itkSetMacro(UseExtrapolation, bool);

  /** Use ball instead of default cross structuring element for repeated dilations. */
  void
  SetUseBallStructuringElement(bool useBall)
  {
    if (useBall != m_UseBallStructuringElement)
    {
      m_UseBallStructuringElement = useBall;
      m_ConnectedComponents->SetFullyConnected(useBall);
      this->Modified();
    }
  }

  /** Use ball instead of default cross structuring element for repeated dilations. */
  itkGetMacro(UseBallStructuringElement, bool);

  /** Use ball instead of default cross structuring element for repeated dilations. */
  itkGetConstMacro(UseBallStructuringElement, bool);

  /** If there is a pixel whose all 4-way neighbors belong the the same label
  except along one axis, and along that axis its neighbors are 0 (background),
  then that axis should be interpolated along. Interpolation is possible
  along more than one axis. Updates LabeledSliceIndices.*/
  void
  DetermineSliceOrientations();

  /** An std::set of slice indices which need to be interpolated. */
  using SliceSetType = std::set<typename TImage::IndexValueType>;

  /** Clears all custom slice positions. */
  void
  ClearLabeledSliceIndices()
  {
    m_LabeledSlices.clear();
    m_LabeledSlices.resize(TImage::ImageDimension);
    this->Modified();
  }

  /** If default slice detection is not wanted, slice indices
   *   between which interpolation is done can be set using this method. */
  void
  SetLabeledSliceIndices(unsigned int                                         axis,
                         typename TImage::PixelType                           label,
                         const std::vector<typename TImage::IndexValueType> & indices)
  {
    SliceSetType sliceSet;
    sliceSet.insert(indices.begin(), indices.end());
    m_LabeledSlices[axis][label] = sliceSet;
    this->Modified();
  }

  /** If default slice detection is not wanted, slice indices
   *   between which interpolation is done can be set using this method. */
  void
  SetLabeledSliceIndices(unsigned int axis, typename TImage::PixelType label, const SliceSetType & indices)
  {
    m_LabeledSlices[axis][label] = indices;
    this->Modified();
  }

  /** Slice indices between which interpolation is done. */
  SliceSetType
  GetLabeledSliceIndices(unsigned int axis, typename TImage::PixelType label)
  {
    return m_LabeledSlices[axis][label];
  }

  // each label gets a set of slices in which it is present
  using LabeledSlicesType = std::unordered_map<typename TImage::PixelType, SliceSetType>;
  using SliceIndicesType = std::vector<LabeledSlicesType>;

  /** Slice indices between which interpolation is done. */
  SliceIndicesType
  GetLabeledSliceIndices()
  {
    return m_LabeledSlices;
  }

protected:
  MorphologicalContourInterpolator();
  ~MorphologicalContourInterpolator() override = default;
  typename TImage::PixelType m_Label;
  int                        m_Axis{ -1 };
  bool                       m_HeuristicAlignment{ true };
  bool                       m_UseDistanceTransform{ true };
  bool                       m_UseBallStructuringElement{ false };
  bool                       m_UseCustomSlicePositions{ false };
  bool                       m_UseExtrapolation{ true };
  IdentifierType             m_MinAlignIters; // minimum number of iterations in align method
  IdentifierType             m_MaxAlignIters; // maximum number of iterations in align method
  IdentifierType             m_ThreadCount;   // for thread local instances
  SliceIndicesType           m_LabeledSlices; // one for each axis

  /** Derived image type alias. */
  using BoolImageType = Image<bool, TImage::ImageDimension>;
  using FloatSliceType = Image<float, TImage::ImageDimension - 1>;
  using BoolSliceType = Image<bool, TImage::ImageDimension - 1>;

  /** Are these two slices equal? */
  bool
  ImagesEqual(typename BoolSliceType::Pointer & a, typename BoolSliceType::Pointer & b);

  /** Does the real work. */
  void
  GenerateData() override;

  /** Determines correspondances between two slices and calls apropriate methods. */
  void
  InterpolateBetweenTwo(int                             axis,
                        TImage *                        out,
                        typename TImage::PixelType      label,
                        typename TImage::IndexValueType i,
                        typename TImage::IndexValueType j,
                        typename SliceType::Pointer &   iconn,
                        typename SliceType::Pointer &   jconn);

  /** If interpolation is done along more than one axis,
  the interpolations are merged using a modified "or" rule:
  -if all interpolated images have 0 for a given pixel, the output is 0
  -if just one image has a non-zero label, then that label is chosen
  -if more than one image has a non-zero label, last written label is chosen */
  void
  InterpolateAlong(int axis, TImage * out, float startProgress, float endProgress);

  /** Slice i has a region, slice j does not */
  void
  Extrapolate(int                             axis,
              TImage *                        out,
              typename TImage::PixelType      label,
              typename TImage::IndexValueType i,
              typename TImage::IndexValueType j,
              typename SliceType::Pointer &   iConn,
              typename TImage::PixelType      iRegionId);

  /** Creates a signed distance field image. */
  typename FloatSliceType::Pointer
  MaurerDM(typename BoolSliceType::Pointer & inImage);

  /** A sequence of conditional dilations starting with begin and reaching end.
  begin and end must cover the same region (size and index the same) */
  std::vector<typename BoolSliceType::Pointer>
  GenerateDilationSequence(typename BoolSliceType::Pointer & begin, typename BoolSliceType::Pointer & end);

  /** Finds an interpolating mask for these two aligned masks */
  typename BoolSliceType::Pointer
  FindMedianImageDilations(typename BoolSliceType::Pointer & intersection,
                           typename BoolSliceType::Pointer & iMask,
                           typename BoolSliceType::Pointer & jMask);

  /** Finds an interpolating mask for these two aligned masks */
  typename BoolSliceType::Pointer
  FindMedianImageDistances(typename BoolSliceType::Pointer & intersection,
                           typename BoolSliceType::Pointer & iMask,
                           typename BoolSliceType::Pointer & jMask);

  /** Build transition sequence and pick the median */
  void
  Interpolate1to1(int                                   axis,
                  TImage *                              out,
                  typename TImage::PixelType            label,
                  typename TImage::IndexValueType       i,
                  typename TImage::IndexValueType       j,
                  typename SliceType::Pointer &         iConn,
                  typename TImage::PixelType            iRegionId,
                  typename SliceType::Pointer &         jConn,
                  typename TImage::PixelType            jRegionId,
                  const typename SliceType::IndexType & translation,
                  bool                                  recursive);

  using PixelList = std::vector<typename TImage::PixelType>;

  /** Splits the bigger region and does N 1-to-1 interpolations */
  void
  Interpolate1toN(int                                   axis,
                  TImage *                              out,
                  typename TImage::PixelType            label,
                  typename TImage::IndexValueType       i,
                  typename TImage::IndexValueType       j,
                  typename SliceType::Pointer &         iConn,
                  typename TImage::PixelType            iRegionId,
                  typename SliceType::Pointer &         jConn,
                  const PixelList &                     jRegionIds,
                  const typename SliceType::IndexType & translation);

  /** Crates a translated copy of part of the image which fits in the newRegion. */
  typename SliceType::Pointer
  TranslateImage(typename SliceType::Pointer &         image,
                 const typename SliceType::IndexType & translation,
                 typename SliceType::RegionType        newRegion);

  /** The returns cardingal of the symmetric distance between images.
  The images must cover the same region. */
  IdentifierType
  CardinalSymmetricDifference(typename BoolSliceType::Pointer & shape1, typename BoolSliceType::Pointer & shape2);

  /** Copied from ImageSource and changed to allocate a cleared buffer. */
  void
  AllocateOutputs() override;

  /** Returns the centroid of given regions */
  typename SliceType::IndexType
  Centroid(typename SliceType::Pointer & conn, const PixelList & regionIds);

  /** Calculates maximum intersection region for both slices given a translation.
  Both inputs are modified so that jRegion is a translated version of iRegion. */
  void
  IntersectionRegions(const typename SliceType::IndexType & translation,
                      typename SliceType::RegionType &      iRegion,
                      typename SliceType::RegionType &      jRegion);

  /** Returns number of intersecting pixels */
  IdentifierType
  Intersection(typename SliceType::Pointer &         iConn,
               typename TImage::PixelType            iRegionId,
               typename SliceType::Pointer &         jConn,
               const PixelList &                     jRegionIds,
               const typename SliceType::IndexType & translation);

  /** How much j needs to be translated to best align with i */
  typename SliceType::IndexType
  Align(typename SliceType::Pointer & iConn,
        typename TImage::PixelType    iRegionId,
        typename SliceType::Pointer & jConn,
        const PixelList &             jRegionIds);

  using BoundingBoxesType = std::unordered_map<typename TImage::PixelType, typename TImage::RegionType>;
  BoundingBoxesType m_BoundingBoxes; // bounding box for each label

  /** Calculates a bounding box of non-zero pixels. */
  typename SliceType::RegionType
  BoundingBox(itk::SmartPointer<SliceType> image);

  /** Expands a region to incorporate the provided index.
   *   Assumes both a valid region and a valid index.
   *   It can be invoked with 2D or 3D region, hence the additional template parameter. */
  template <typename T2>
  void
  ExpandRegion(typename T2::RegionType & region, const typename T2::IndexType & index);

  /** Connected components of a specified region. */
  typename SliceType::Pointer
  RegionedConnectedComponents(const typename TImage::RegionType & region,
                              typename TImage::PixelType          label,
                              IdentifierType &                    objectCount);

  /** Seed and mask must cover the same region (size and index the same). */
  typename BoolSliceType::Pointer
  Dilate1(typename BoolSliceType::Pointer & seed, typename BoolSliceType::Pointer & mask);

  using RoiType = ExtractImageFilter<TImage, SliceType>;
  typename RoiType::Pointer m_RoI;

  using BinarizerType = BinaryThresholdImageFilter<SliceType, BoolSliceType>;
  typename BinarizerType::Pointer m_Binarizer;

  using ConnectedComponentsType = ConnectedComponentImageFilter<BoolSliceType, SliceType>;
  typename ConnectedComponentsType::Pointer m_ConnectedComponents;
};
} // namespace itk

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

#endif // itkMorphologicalContourInterpolator_h
