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


#include "itkWatershedBoundary.h"
#include "itkWatershedSegmentTable.h"
#include "itkEquivalencyTable.h"

namespace itk::watershed
{
/**
 * \class Segmenter
 *
 * This filter implements the first step in the N-d watershed segmentation
 * algorithm.  It produces a segmented, labeled image from a scalar-valued
 * image input. This filter is used in conjunction with
 * WatershedSegmentTreeGenerator and WatershedRelabeler to produce a final
 * segmented image volume. See itk::WatershedImageFilter for an overview of the
 * entire algorithm and notes on the terminology used in describing it.
 *
 * \par
 * The filter is designed to operate in streaming or non-streaming mode.  For
 * more information, see the itk::WatershedImageFilter documentation.
 *
 * \par Input
 * There is one input to this algorithm, a real-valued (scalar) itk::Image of
 * arbitrary dimension.  The input is assumed to represents a height function,
 * such as a gradient magnitude edge image.  The filter can process an image of
 * any dimension. Note that the terms "pixel" and "voxel" are
 * interchangeable in this and other watershed component class documentation.
 *
 * \par Outputs
 * There are three potential outputs of this algorithm described below.
 *
 * \par
 * The first output is a labeled image of IdentifierType integers.  This is an
 * initial segmentation and labeling that is fed into successive components of
 * the watershed algorithm.
 *
 * \par
 * The second output is a table of segment information,
 * itk::watershed::SegmentTable. This table is a record of each segment
 * numbered in the initial segmentation (output number one) with relevant
 * information needed in successive stages of the algorithm.
 *
 * \par
 * The third output is a data structure containing boundary pixel information,
 * itk::watershed::Boundary.  This data is only generated if the flag
 * DoBoundaryAnalysis is set to true and is only useful in streaming
 * applications.
 *
 * \par Parameters
 * Threshold is specified as a percentage (0.0 - 1.0) of the maximum height of
 * the image. This filter thresholds the input image to remove all values below
 * \f$ L = min + T * (max - min) \f$, where \f$ max, min \f$ are the maximum,
 * minimum values in the image and \f$ T \f$ is the threshold parameter
 * value. Values in the image less than \f$ L \f$ are raised to \f$ L \f$.
 *
 * \par
 * Thresholding minimum values in the image decreases the number of local
 * minima in the image and produces an initial segmentation with fewer
 * segments.  The assumption is that the "shallow" regions that this
 * thresholding eliminates are generally not of interest.
 *
 * \sa WatershedImageFilter
 * \ingroup WatershedSegmentation
 * \ingroup ITKWatersheds
 */
template <typename TInputImage>
class ITK_TEMPLATE_EXPORT Segmenter : public ProcessObject
{
public:
  /** Standard self type alias */
  using Self = Segmenter;

  /** Define image types and dimensionality  */
  using InputImageType = TInputImage;
  static constexpr unsigned int ImageDimension = TInputImage::ImageDimension;

  using OutputImageType = Image<IdentifierType, Self::ImageDimension>;
  using ImageRegionType = typename InputImageType::RegionType;
  using InputPixelType = typename InputImageType::PixelType;
  using BoundaryType = Boundary<InputPixelType, Self::ImageDimension>;
  using BoundaryIndexType = typename BoundaryType::IndexType;
  using BoundaryFlatHashValueType = typename BoundaryType::FlatHashValueType;
  using SegmentTableType = SegmentTable<InputPixelType>;
  using DataObjectPointer = DataObject::Pointer;

  /** Methods to implement smart pointers and work with the itk object factory
   */
  /** @ITKStartGrouping */
  using Superclass = ProcessObject;
  using Pointer = SmartPointer<Self>;
  using ConstPointer = SmartPointer<const Self>;
  itkNewMacro(Self);
  itkOverrideGetNameOfClassMacro(Segmenter);
  /** @ITKEndGrouping */

  /** Typedefs necessary on microsoft VC++ to avoid internal compiler errors */
  using InputImageTypePointer = typename InputImageType::Pointer;
  using OutputImageTypePointer = typename OutputImageType::Pointer;
  using SegmentTableTypePointer = typename SegmentTableType::Pointer;
  using BoundaryTypePointer = typename BoundaryType::Pointer;

  /** A constant used in the labeling algorithm.  */
  static constexpr IdentifierType NULL_LABEL = 0;

  /** A constant used in the labeling algorithm.  */
  static constexpr short NULL_FLOW = -1;

  /** Get/Set the input image.   */
  InputImageType *
  GetInputImage()
  {
    return itkDynamicCastInDebugMode<InputImageType *>(this->ProcessObject::GetInput(0));
  }

  void
  SetInputImage(InputImageType * img)
  {
    this->ProcessObject::SetNthInput(0, img);
  }

  /** Get/Set the labeled output image.  The output image is always of
    IdentifierType integers. */
  OutputImageType *
  GetOutputImage()
  {
    return itkDynamicCastInDebugMode<OutputImageType *>(this->ProcessObject::GetOutput(0));
  }

  void
  SetOutputImage(OutputImageType * img)
  {
    this->ProcessObject::SetNthOutput(0, img);
  }

  /** Get/Set the segment table.  The segment table is a table of segmentation
   * information identifying each region produced by the labeling algorithm. */
  SegmentTableType *
  GetSegmentTable()
  {
    return itkDynamicCastInDebugMode<SegmentTableType *>(this->ProcessObject::GetOutput(1));
  }

  void
  SetSegmentTable(SegmentTableType * s)
  {
    this->ProcessObject::SetNthOutput(1, s);
  }

  /** Returns the boundary information data necessary only for data streaming
    applications.  */
  BoundaryType *
  GetBoundary()
  {
    return itkDynamicCastInDebugMode<BoundaryType *>(this->ProcessObject::GetOutput(2));
  }

  void
  SetBoundary(BoundaryType * b)
  {
    this->ProcessObject::SetNthOutput(2, b);
  }

  /** Standard non-threaded pipeline execution method. */
  void
  GenerateData() override;

  /** This method is necessary until the streaming mechanisms of the Itk
   * pipeline are full fleshed out.  It is only used for streaming
   * applications.  Calling this method gets/sets the image size of the
   * complete volume being streamed.  The member variables controlled by
   * this method will not be modified by the Itk pipeline and are necessary
   * for analysis of boundaries.   */
  /** @ITKStartGrouping */
  void
  SetLargestPossibleRegion(ImageRegionType reg)
  {
    if (reg == m_LargestPossibleRegion)
    {
      return;
    }
    m_LargestPossibleRegion = reg;
    this->Modified();
  }

  ImageRegionType
  GetLargestPossibleRegion() const
  {
    return m_LargestPossibleRegion;
  }
  /** @ITKEndGrouping */

  /** Helper function.  Other classes may have occasion to use this. Relabels
      an image according to a table of equivalencies. */
  static void RelabelImage(OutputImageTypePointer, ImageRegionType, EquivalencyTable::Pointer);

  /** Standard itk::ProcessObject subclass method. */
  using DataObjectPointerArraySizeType = ProcessObject::DataObjectPointerArraySizeType;
  using Superclass::MakeOutput;
  DataObjectPointer
  MakeOutput(DataObjectPointerArraySizeType idx) override;

  /** Gets/Sets the initial label (IdentifierType integer value) used
   * by the labeling algorithm.  Only necessary for streaming applications. */
  /** @ITKStartGrouping */
  itkSetMacro(CurrentLabel, IdentifierType);
  itkGetConstMacro(CurrentLabel, IdentifierType);
  /** @ITKEndGrouping */

  /** Gets/Sets the input threshold. Threshold is specified as a percentage
   * (0.0 - 1.0) of the maximum height of the image. This filter thresholds the
   * input image to remove all values below \f$ L = min + T * (max -  min) \f$,
   * where \f$ max, min \f$ are the maximum, minimum values in the image and \f$
   * T \f$ is the threshold parameter value. Values in the image less than \f$ L
   * \f$ are raised to \f$ L \f$. Thresholding minimum values in the image
   * decreases the number of local minima in the image and produces an initial
   * segmentation with fewer segments.  The assumption is that the "shallow"
   * regions that this thresholding eliminates are generally not of
   * interest. */
  /** @ITKStartGrouping */
  itkSetClampMacro(Threshold, double, 0.0, 1.0);
  itkGetConstMacro(Threshold, double);
  /** @ITKEndGrouping */

  /** Turns on special labeling of the boundaries for streaming applications.
   * The default value is FALSE, meaning that boundary analysis is turned
   * off.   */
  /** @ITKStartGrouping */
  itkSetMacro(DoBoundaryAnalysis, bool);
  itkGetConstMacro(DoBoundaryAnalysis, bool);
  /** @ITKEndGrouping */

  /** Determines whether the algorithm will sort the adjacencies in its
   * SegmentTable before returning.  Default is true.  This is an option only
   * useful for streaming applications where the sorting only needs to be done
   * after all iterations have taken place. */
  /** @ITKStartGrouping */
  itkGetConstMacro(SortEdgeLists, bool);
  itkSetMacro(SortEdgeLists, bool);
  /** @ITKEndGrouping */
protected:
  /** Structure storing information about image flat regions.
   * Flat regions are connected pixels of the same value.  */
  struct flat_region_t
  {
    IdentifierType * min_label_ptr;
    InputPixelType   bounds_min;
    //    InputPixelType  bounds_max; // <-- may not be necc.
    InputPixelType value;
    bool           is_on_boundary{ false };
    flat_region_t() = default;
  };

  /** Table for storing flat region information.  */
  using flat_region_table_t = std::unordered_map<IdentifierType, flat_region_t>;

  struct connectivity_t
  {
    unsigned int                          size;
    unsigned int *                        index;
    typename InputImageType::OffsetType * direction;
  };

  /** Table for storing tables of edges.  This is convenient in
   * generating the segment table,  even though the edge tables
   * are stored as ordered lists.  An "edge" in this context
   * is synonymous with a segment "adjacency".   */
  using edge_table_t = std::map<IdentifierType, InputPixelType>;

  using edge_table_hash_t = std::unordered_map<IdentifierType, edge_table_t>;

  Segmenter();
  Segmenter(const Self &) {}
  ~Segmenter() override;
  void
  PrintSelf(std::ostream & os, Indent indent) const override;

  void
  operator=(const Self &)
  {}

  /** Constructs the connectivity list and the corresponding set of directional
   * Offset indices. */
  virtual void
  GenerateConnectivity();

  /** This method asks for an image region that is one pixel larger
   * at each boundary than the region being processed.  This single pixel
   * expansion represents an overlap with adjacent image chunks   */
  void
  GenerateInputRequestedRegion() override;

  void
  GenerateOutputRequestedRegion(DataObject * output) override;

  void
  UpdateOutputInformation() override;

  /**  Allocates boundary structure information and sets the
   * boundary data to null values.   */
  void
  InitializeBoundary();

  /** Performs a gradient descent connected component analysis
   * at the boundaries of the images that border other
   * image chunks.  Useful only in data streaming applications.   */
  void
  AnalyzeBoundaryFlow(InputImageTypePointer, flat_region_table_t &, InputPixelType);

  /** Fills boundary pixels with a specified value.  Used by labeling
   * methods to build a very high "wall" around the image so that
   * gradient descent does not need to watch boundaries.   */
  void BuildRetainingWall(InputImageTypePointer, ImageRegionType, InputPixelType);

  /** Labels all the local minima in the image.  Also identifies and labels
   * connected  "flat" regions.   */
  void
  LabelMinima(InputImageTypePointer, ImageRegionType, flat_region_table_t &, InputPixelType);

  /** Follows each unlabeled pixel in the image down its path of steepest
   * descent.  Each pixel along that path is identified with the local minima
   * already labeled at the end of the path.   */
  void GradientDescent(InputImageTypePointer, ImageRegionType);

  /** Associates each flat region with a local minimum and relabels
    accordingly.  */
  void
  DescendFlatRegions(flat_region_table_t &, ImageRegionType);

  /** Adds entries to the output segment table for all labeled segments in the
   * image.  */
  void UpdateSegmentTable(InputImageTypePointer, ImageRegionType);

  /** Traverses each boundary and fills in the data needed for joining
   * streamed chunks of an image volume.  Only necessary for streaming
   * applications.   */
  void
  CollectBoundaryInformation(flat_region_table_t &);

  /** Helper function.  Thresholds low values and copies values from one image
   * into another. The source and destination regions must match in size (not
   * enforced).  For integral types, the dynamic range of the image is
   * adjusted such that the maximum value in the image is always at
   * least one less than the maximum value allowed for that data type. */
  static void
  Threshold(InputImageTypePointer destination,
            InputImageTypePointer source,
            const ImageRegionType source_region,
            const ImageRegionType destination_region,
            InputPixelType        threshold);

  /** Helper function.  Finds the minimum and maximum values in an image. */
  static void
  MinMax(InputImageTypePointer img, ImageRegionType region, InputPixelType & min, InputPixelType & max);

  /** Helper function. Finds the minimum and maximum values in an image.   */
  static void
  MergeFlatRegions(flat_region_table_t &, EquivalencyTable::Pointer);

  /** Helper functions for filling in regions with values   */
  static void
  SetInputImageValues(InputImageTypePointer img, const ImageRegionType region, InputPixelType value);

  static void
  SetOutputImageValues(OutputImageTypePointer img, const ImageRegionType region, IdentifierType value);

  /** This is a debugging method.  Will be removed. 11/14/01 jc   */
  //  bool CheckLabeledBoundaries();

  /** Holds generalized connectivity information for connected component
   * labeling and gradient descent analysis in pixel neighborhoods.  */
  connectivity_t m_Connectivity{};

private:
  /** Helper, debug method.   */
  //  void PrintFlatRegions(flat_region_table_t &t);

  /** This is the actual data set size.  The pipeline will alter its
   *  LargestPossibleRegion, so we need to preserve it here explicitly for
   *  streaming applications*/
  ImageRegionType m_LargestPossibleRegion{};

  bool           m_SortEdgeLists{};
  bool           m_DoBoundaryAnalysis{};
  double         m_Threshold{};
  double         m_MaximumFloodLevel{};
  IdentifierType m_CurrentLabel{};
};
} // namespace itk::watershed

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

#endif
