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

#include "itkWatershedSegmentTable.h"
#include "itkWatershedSegmentTree.h"
#include "itkEquivalencyTable.h"

#include <algorithm>
#include <utility>

namespace itk
{
namespace watershed
{
/**
 * \class SegmentTreeGenerator
 * This filter is a process object class that implements a step in the N-d
 * watershed segmentation algorithm.  It produces an ordered list (SegmentTree)
 * of binary merges between segments at increasing saliency levels. This
 * filter is used in conjunction with WatershedSegmenter and WatershedRelabeler
 * to process an initial segmentation into a final relabeled volume.  See
 * itk::WatershedImageFilter for an overview of watershed segmentation.
 *
 * \par Inputs
 * There are two inputs to this filter, the second is optional.
 * (1) A table of segments (SegmentTable) derived from a labeled image.
 * (2) Optionally, a table of predetermined equivalencies among segments
 * (EquivalencyTable), necessary for streaming applications where
 * segments may be joined across streamed chunk boundaries.  The flag
 * Merge must be set to true to enable this functionality.
 *
 * \par Outputs
 * The output of this filter is a list of binary merges of segments at
 * increasing saliency.  This is the data structure
 * itk::watershed::SegmentTree referred to as a "merge tree" in the
 * itk::WatershedImageFilter documentation.
 *
 * \par Parameters
 * There are two parameters to this filter described below.
 *
 * \par
 * FloodLevel is specified as a percentage (0.0 - 1.0) of the maximum possible
 * saliency value in the initial image from which the segment table was
 * derived.  A value of 0.0 calculates no merges.  A value of 1.0 calculates
 * all of the potential merges that can occur as the FloodLevel is increased to
 * the maximum saliency value.  Typically, there is no need to calculate merges
 * past about 40% of the total depth.  Because this is the most computationally
 * intensive piece of the watershed segmentation algorithm, it is a good idea
 * to tune this parameter as low as possible on larger volumes.
 *
 * \par
 * Merge is a boolean flag indicating whether or not to pre-merge the segments
 * marked as equivalent in the EquivalencyTable.  This is only useful for
 * streaming applications and is turned off by default.  (TRUE == merge, FALSE
 * == do not merge).
 * \sa itk::WatershedImageFilter
 * \ingroup WatershedSegmentation
 * \ingroup ITKWatersheds
 */
template <typename TScalar>
class ITK_TEMPLATE_EXPORT SegmentTreeGenerator : public ProcessObject
{
public:
  ITK_DISALLOW_COPY_AND_MOVE(SegmentTreeGenerator);

  /**  Standard itk smart pointer declarations    */
  using Self = SegmentTreeGenerator;
  using Superclass = ProcessObject;
  using Pointer = SmartPointer<Self>;
  using ConstPointer = SmartPointer<const Self>;

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

  /** Convenient type definitions */
  using ScalarType = TScalar;
  using SegmentTableType = SegmentTable<ScalarType>;
  using SegmentTreeType = SegmentTree<ScalarType>;
  using EquivalencyTableType = EquivalencyTable;
  using OneWayEquivalencyTableType = OneWayEquivalencyTable;
  using DataObjectPointer = DataObject::Pointer;

  /** Typedefs to avoid internal compiler error bug on Microsoft VC++ */
  using SegmentTableTypePointer = typename SegmentTableType::Pointer;
  using OneWayEquivalencyTableTypePointer = typename OneWayEquivalencyTableType::Pointer;
  using SegmentTreeTypePointer = typename SegmentTreeType::Pointer;

  /** Get/Set the input table of segments to process */
  SegmentTableType *
  GetInputSegmentTable()
  {
    return static_cast<SegmentTableType *>(this->ProcessObject::GetInput(0));
  }

  void
  SetInputSegmentTable(SegmentTableType * st)
  {
    // Reset the highest calculated flood level if we are given a
    // different input image.
    if (st != this->GetInput(0))
    {
      m_HighestCalculatedFloodLevel = 0.0;
    }
    this->ProcessObject::SetNthInput(0, st);
  }

  /** Get/Set input table of equivalencies to pre-merge before
   * running the tree generator algorithm.  Only useful for
   * streaming applications */
  void
  SetInputEquivalencyTable(EquivalencyTableType * eq)
  {
    this->ProcessObject::SetNthInput(1, eq);
  }

  EquivalencyTableType *
  GetInputEquivalencyTable()
  {
    return static_cast<EquivalencyTableType *>(this->ProcessObject::GetInput(1));
  }

  /** Get/Set the output data */
  SegmentTreeType *
  GetOutputSegmentTree()
  {
    return static_cast<SegmentTreeType *>(this->ProcessObject::GetOutput(0));
  }

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

  /** Get/Set a boolean flag indicating whether or not to pre-merge the
    segments marked  as equivalent in the EquivalencyTable.  This is only
    useful for streaming  applications and is turned off by default.  (TRUE ==
    merge, FALSE == do not merge).       */
  itkSetMacro(Merge, bool);
  itkGetConstMacro(Merge, bool);

  /** Get/Set FloodLevel.  FloodLevel is specified as a percentage (0.0 - 1.0)
   of the maximum possible saliency value in the initial image from which the
   segment table was derived. A value of 0.0 calculates no merges.  A value of
   1.0 calculates all of the potential merges that can occur as the FloodLevel
   is increased to the  maximum saliency value.    */
  void
  SetFloodLevel(double);

  itkGetConstMacro(FloodLevel, double);

  /** Get/Set HighestCalculatedFloodLevel.  HighestCalculatedFloodLevel keeps
   * track of the highest level this filter has computed.  It is
   * used to prevent unnecessary re-execution of the filter. */
  itkSetMacro(HighestCalculatedFloodLevel, double);
  itkGetConstMacro(HighestCalculatedFloodLevel, double);

  /** Get/Set a flag that prevents the filter from copying its input
   * segment table before executing.  This can be enabled to conserve
   * memory, especially in streaming applications where memory is a
   * concern. If enabled, the input to this filter must always be
   * re-executed on updates. Default is false. */
  itkSetMacro(ConsumeInput, bool);
  itkGetConstMacro(ConsumeInput, bool);

  /** Performs a merge of two segments in a SegmentTable according
   * to criteria specific to this algorithm.   */
  static void
  MergeSegments(SegmentTableTypePointer, OneWayEquivalencyTableTypePointer, const IdentifierType, const IdentifierType);

  /** This method should not be used.  It will be removed in future versions
   * of this filter. */
  static void
  PruneMergeSegments(SegmentTableTypePointer,
                     OneWayEquivalencyTableTypePointer,
                     const IdentifierType,
                     const IdentifierType,
                     ScalarType);

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

protected:
  SegmentTreeGenerator();
  ~SegmentTreeGenerator() override = default;

  void
  PrintSelf(std::ostream & os, Indent indent) const override;

  /** Generates an initial list of all potential merges in
   * the segment table.   */
  void CompileMergeList(SegmentTableTypePointer, SegmentTreeTypePointer);

  /** Compiles a list of all the actual segment merges up to the specified
   * flood level, recomputing new potential merges as it goes.   */
  void ExtractMergeHierarchy(SegmentTableTypePointer, SegmentTreeTypePointer);

  void
  MergeEquivalencies();

  /** Methods required by the itk pipeline */
  void
  GenerateOutputRequestedRegion(DataObject * output) override;

  void
  GenerateInputRequestedRegion() override;

private:
  bool   m_Merge{ false };
  double m_FloodLevel{ 0.0 };
  bool   m_ConsumeInput{ false };

  using HashMapType = std::unordered_map<IdentifierType, bool>;

  OneWayEquivalencyTableType::Pointer m_MergedSegmentsTable{};

  /** This value keeps track of the highest level this filter has been
   *  calculated.  m_FloodLevel can be manipulated anywhere below this
   *  level without re-executing the filter, preventing unnecessary
   *  updates. */
  double m_HighestCalculatedFloodLevel{ 0.0 };
};
} // end namespace watershed
} // end namespace itk

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

#endif