WeightedFastHoughTree.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
#pragma once
 
#include <tracking/trackFindingCDC/hough/trees/DynTree.h>
#include <tracking/trackFindingCDC/hough/baseelements/WithWeightedItems.h>
#include <tracking/trackFindingCDC/hough/baseelements/WithSharedMark.h>
 
#include <vector>
#include <memory>
#include <cassert>
#include <cfloat>
#include <cmath>
#include <algorithm>
#include <type_traits>
 
namespace Belle2 {
  namespace TrackFindingCDC {

    template<class T, class ADomain, class ADomainDivsion>
    class  WeightedParititioningDynTree :
      public DynTree<  WithWeightedItems<ADomain, T>, ADomainDivsion> {
    private:
      using Super = DynTree<  WithWeightedItems<ADomain, T>, ADomainDivsion>;
 
    public:
      WeightedParititioningDynTree(ADomain topDomain, ADomainDivsion domainDivsion) :
        Super(WithWeightedItems<ADomain, T>(std::move(topDomain)), std::move(domainDivsion))
      {}
    };

    template<class T, class ADomain, class ADomainDivsion>
    class WeightedFastHoughTree :
      public WeightedParititioningDynTree<WithSharedMark<T>, ADomain, ADomainDivsion> {
 
    private:
      using Super = WeightedParititioningDynTree<WithSharedMark<T>, ADomain, ADomainDivsion>;
 
    public:
      using WeightedParititioningDynTree<WithSharedMark<T>, ADomain, ADomainDivsion>::WeightedParititioningDynTree;

      using Node = typename Super::Node;
 
    public:
      template<class Ts>
      void seed(const Ts& items)
      {
        this->fell();
        Node& topNode = this->getTopNode();
        for (auto&& item : items) {
          m_marks.push_back(false);
          bool& markOfItem = m_marks.back();
          TrackingUtilities::Weight weight = DBL_MAX;
          topNode.insert(WithSharedMark<T>(T(item), &markOfItem), weight);
        }
      }

      template <class AItemInDomainMeasure>
      std::vector<std::pair<ADomain, std::vector<T>>>
      findHeavyLeavesDisjoint(AItemInDomainMeasure& weightItemInDomain,
                              int maxLevel,
                              double minWeight)
      {
        auto skipLowWeightNode = [minWeight](const Node * node) {
          return not(node->getWeight() >= minWeight);
        };
        return findLeavesDisjoint(weightItemInDomain, maxLevel, skipLowWeightNode);
      }

      template <class AItemInDomainMeasure, class ASkipNodePredicate>
      std::vector<std::pair<ADomain, std::vector<T>>>
      findLeavesDisjoint(AItemInDomainMeasure& weightItemInDomain,
                         int maxLevel,
                         ASkipNodePredicate& skipNode)
      {
        std::vector<std::pair<ADomain,  std::vector<T> > > found;
        auto isLeaf = [&found, &skipNode, maxLevel](Node * node) {
          // Skip the expansion and the filling of the children
          if (skipNode(node)) {
            return true;
          }
 
          // Node is a leaf at the maximum level
          // Save its content
          // Do not walk children
          if (node->getLevel() >= maxLevel) {
            const ADomain* domain = node;
            found.emplace_back(*domain, std::vector<T>(node->begin(), node->end()));
            for (WithSharedMark<T>& markableItem : *node) {
              markableItem.mark();
            }
            return true;
          }
 
          // Else to node has enough weight and is not at the lowest level
          // Signal that it is not a leaf
          // Continue to create and fill children.
          return false;
        };
        fillWalk(weightItemInDomain, isLeaf);
        return found;
      }

      template <class AItemInDomainMeasure>
      std::vector<std::pair<ADomain, std::vector<T>>>
      static findHeaviestLeafRepeated(AItemInDomainMeasure& weightItemInDomain,
                                      int maxLevel,
                                      const TrackingUtilities::Weight minWeight = NAN)
      {
        auto skipLowWeightNode = [minWeight](const Node * node) {
          return not(node->getWeight() >= minWeight);
        };
        return findHeaviestLeafRepeated(weightItemInDomain, maxLevel, skipLowWeightNode);
      }

      template <class AItemInDomainMeasure, class ASkipNodePredicate>
      std::vector<std::pair<ADomain, std::vector<T>>>
      findHeaviestLeafRepeated(AItemInDomainMeasure& weightItemInDomain,
                               int maxLevel,
                               ASkipNodePredicate& skipNode)
      {
        std::vector<std::pair<ADomain,  std::vector<T> > > found;
        Node* node = findHeaviestLeaf(weightItemInDomain, maxLevel, skipNode);
        while (node) {
          const ADomain* domain = node;
          found.emplace_back(*domain, std::vector<T>(node->begin(), node->end()));
          for (WithSharedMark<T>& markableItem : *node) {
            markableItem.mark();
          }
          node = findHeaviestLeaf(weightItemInDomain, maxLevel, skipNode);
        }
        return found;
      }

      template <class AItemInDomainMeasure, class ASkipNodePredicate>
      std::unique_ptr<std::pair<ADomain, std::vector<T>>>
      findHeaviestLeafSingle(AItemInDomainMeasure& weightItemInDomain,
                             int maxLevel,
                             ASkipNodePredicate& skipNode)
      {
        using Result = std::pair<ADomain,  std::vector<T> >;
        std::unique_ptr<Result> found = nullptr;
        Node* node = findHeaviestLeaf(weightItemInDomain, maxLevel, skipNode);
        if (node) {
          const ADomain* domain = node;
          found.reset(new Result(*domain, std::vector<T>(node->begin(), node->end())));
          for (WithSharedMark<T>& markableItem : *node) {
            markableItem.mark();
          }
        }
        return found;
      }

      template <class AItemInDomainMeasure, class ASkipNodePredicate>
      Node* findHeaviestLeaf(AItemInDomainMeasure& weightItemInDomain,
                             int maxLevel,
                             ASkipNodePredicate& skipNode)
      {
        Node* heaviestNode = nullptr;
        TrackingUtilities::Weight heighestWeigth = NAN;
        auto isLeaf = [&heaviestNode, &heighestWeigth, maxLevel, &skipNode](Node * node) {
          // Skip the expansion and the filling of the children
          if (skipNode(node)) {
            return true;
          }
 
          TrackingUtilities::Weight nodeWeight = node->getWeight();
          // Skip the expansion and filling of the children if the node has not enough weight
          if (not std::isnan(heighestWeigth) and not(nodeWeight > heighestWeigth)) {
            return true;
          }
 
          // Node is a leaf at the maximum level and is heavier than everything seen before.
          // Save its content
          // Do not walk children
          if (node->getLevel() >= maxLevel) {
            heaviestNode = node;
            heighestWeigth = nodeWeight;
            return true;
          }
          return false;
        };
        // The isLeaf predicate does not mark any items.
        const bool isLeafMarksItems = false;
        fillWalk(weightItemInDomain, isLeaf, isLeafMarksItems);
        return heaviestNode;
      }
 
    public:
      template<class AItemInDomainMeasure, class AIsLeafPredicate>
      void fillWalk(AItemInDomainMeasure& weightItemInDomain,
                    AIsLeafPredicate& isLeaf,
                    bool isLeafMarksItems = true)
      {
        auto walker = [&weightItemInDomain, &isLeaf](Node * node) {
          // Check if node is a leaf
          // Do not create children in this case
          if (isLeaf(node)) {
            // Do not walk children.
            return false;
          }
 
          // Node is not a leaf.
          // Check if it has children.
          // If children have not been created, create and fill them.
          typename Node::Children* children = node->getChildren();
          if (not children) {
            node->createChildren();
            children = node->getChildren();
            if constexpr(std::is_invocable_v<AItemInDomainMeasure&, const T&, Node*>) {
              // Weighting function does not modify the item: fill all children in a single pass
              // over the parent items without copies. Each child receives the items in the same order.
              for (const WithSharedMark<T>& markableItem : *node) {
                // Weighting function should not see the mark, but only the item itself.
                const T& item(markableItem);
                for (Node& childNode : *children) {
                  const TrackingUtilities::Weight weight = weightItemInDomain(item, &childNode);
                  if (not std::isnan(weight)) {
                    childNode.insert(markableItem, weight);
                  }
                }
              }
            } else {
              for (Node& childNode : *children) {
                assert(childNode.getChildren() == nullptr);
                assert(childNode.size() == 0);
                auto measure =
                  // cppcheck-suppress constParameterReference ; the item is unwrapped as a non-const reference below
                [&childNode, &weightItemInDomain](WithSharedMark<T>& markableItem) -> TrackingUtilities::Weight {
                  // Weighting function should not see the mark, but only the item itself.
                  T & item(markableItem);
                  return weightItemInDomain(item, &childNode);
                };
                childNode.insert(*node, measure);
              }
            }
          }
          // Continue to walk the children.
          return true;
        };
        walkHeighWeightFirst(walker, isLeafMarksItems);
      }

      template<class ATreeWalker>
      void walkHeighWeightFirst(ATreeWalker& walker, bool walkerMarksItems = true)
      {
        if (not walkerMarksItems and std::find(m_marks.begin(), m_marks.end(), true) == m_marks.end()) {
          auto unmarkedPriority = [](Node * node) -> float {
            return node->getWeight();
          };
          this->walk(walker, unmarkedPriority);
          return;
        }
 
        auto priority = [](Node * node) -> float {
          auto isMarked = [](const WithSharedMark<T>& markableItem) -> bool {
            return markableItem.isMarked();
          };
          node->eraseIf(isMarked);
          return node->getWeight();
        };
 
        this->walk(walker, priority);
      }

      // cppcheck-suppress duplInheritedMember ; intentionally hides the base class member, which it extends and then calls
      void fell()
      {
        this->getTopNode().clear();
        m_marks.clear();
        Super::fell();
      }

      // cppcheck-suppress duplInheritedMember ; intentionally hides the base class member, which it extends and then calls
      void raze()
      {
        this->fell();
        Super::raze();
        m_marks.shrink_to_fit();
      }
 
    private:
      std::deque<bool> m_marks;
      // Note: Have to use a deque here because std::vector<bool> is special
      // std::vector<bool> m_marks;
    };
  }
}