QuadTreeProcessor.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/legendre/quadtree/QuadTreeNode.h>
#include <tracking/trackFindingCDC/legendre/quadtree/QuadTreeItem.h>
 
#include <tracking/trackingUtilities/utilities/Algorithms.h>
 
#include <framework/logging/Logger.h>
 
#include <algorithm>
#include <array>
#include <functional>
#include <memory>
#include <map>
#include <vector>
#include <deque>
#include <utility>
 
namespace Belle2 {
  namespace TrackingUtilities {
    class CDCWireHit;
  }
  namespace TrackFindingCDC {

    template<typename AX, typename AY, class AData>
    class QuadTreeProcessor {
 
    private:
      static const int c_maxNChildren = 16;
 
    public:
      using Item = QuadTreeItem<AData>;

      using QuadTree = QuadTreeNode<AX, AY, Item>;

      using XSpan = typename QuadTree::XSpan;

      using YSpan = typename QuadTree::YSpan;

      using XYSpans = std::pair<XSpan, YSpan>;

      using QuadTreeChildren = typename QuadTree::Children;

      using CandidateReceiver = std::function<void(const std::vector<AData*>&, QuadTree*)>;
 
    public:
      QuadTreeProcessor(int lastLevel,
                        int seedLevel,
                        const XYSpans& xySpans,
                        bool debugOutput = false)
        : m_quadTree{std::make_unique<QuadTree>(xySpans.first, xySpans.second, 0, nullptr)}
        , m_lastLevel(lastLevel)
        , m_seedLevel(seedLevel)
        , m_debugOutput(debugOutput)
        , m_debugOutputMap()
      {
      }

      virtual ~QuadTreeProcessor()
      {
        clear();
      }

      void clear()
      {
        m_seededTrees.clear();
        m_quadTree->clearChildren();
        m_quadTree->clearItems();
        m_items.clear();
      }

      void seed(const std::vector<AData*>& datas)
      {
        // Create the items
        for (AData* data : datas) {
          m_items.emplace_back(data);
        }
 
        // Creating the seed level
        long nSeedBins = pow(2, m_seedLevel);
        m_seededTrees.reserve(nSeedBins * nSeedBins);
 
        // Expand the first levels to the seed sectors
        m_seededTrees.push_back(m_quadTree.get());
        std::vector<QuadTree*> nextSeededTrees;
 
        for (int level = 0; level < m_seedLevel; ++level) {
          for (QuadTree* node : m_seededTrees) {
            if (node->getChildren().empty()) {
              this->createChildren(node, node->getChildren());
            }
            for (QuadTree& child :  node->getChildren()) {
              nextSeededTrees.push_back(&child);
            }
          }
          std::swap(nextSeededTrees, m_seededTrees);
          nextSeededTrees.clear();
        }
 
        // Fill the seed level with the items
        m_itemPtrs.clear();
        m_itemPtrs.reserve(m_items.size());
        for (Item& item : m_items) {
          m_itemPtrs.push_back(&item);
        }
 
        for (QuadTree* seededTree : m_seededTrees) {
          seededTree->reserveItems(m_items.size());
        }
 
        this->insertItemsInNodes(m_seededTrees, m_itemPtrs);
      }
 
    public:
      std::vector<AData*> getAssignedItems()
      {
        std::vector<const TrackingUtilities::CDCWireHit*> result;
        for (QuadTree* seededTree : m_seededTrees) {
          for (Item* item : seededTree->getItems()) {
            result.push_back(item->getPointer());
          }
        }
        std::sort(result.begin(), result.end());
        result.erase(std::unique(result.begin(), result.end()), result.end());
        return result;
      }
 
    public:
      void fill(const CandidateReceiver& candidateReceiver, int nHitsThreshold)
      {
        fill(candidateReceiver, nHitsThreshold, std::numeric_limits<AY>::max());
      }

      void fill(const CandidateReceiver& candidateReceiver, int nHitsThreshold, float yLimit)
      {
        std::vector<QuadTree*> quadTrees = m_seededTrees;
        std::sort(quadTrees.begin(), quadTrees.end(), [](const QuadTree * quadTree1, const QuadTree * quadTree2) {
          return quadTree1->getNItems() > quadTree2->getNItems();
        });
 
        for (QuadTree* tree : quadTrees) {
          erase_remove_if(tree->getItems(), [](Item * hit) { return hit->isUsed(); });
          fillGivenTree(tree, candidateReceiver, nHitsThreshold, yLimit);
        }
      }
 
    private:
      void fillGivenTree(QuadTree* node,
                         const CandidateReceiver& candidateReceiver,
                         int nItemsThreshold,
                         AY yLimit)
      {
        if (node->getNItems() < nItemsThreshold) {
          return;
        }
 
        if ((node->getYMin() > yLimit) or (-node->getYMax() > yLimit)) {
          return;
        }
 
        if (isLeaf(node)) {
          callResultFunction(node, candidateReceiver);
          return;
        }
 
        if (node->getChildren().empty()) {
          this->createChildren(node, node->getChildren());
        }
 
        if (!node->checkFilled()) {
          fillChildren(node, node->getItems());
          node->setFilled();
        }
 
        // Kept on the stack - this function is called millions of times per event
        std::array<QuadTree*, c_maxNChildren> children;
        int nChildren = 0;
        for (QuadTree& child : node->getChildren()) {
          B2ASSERT("More children than the quad tree processor supports", nChildren < c_maxNChildren);
          children[nChildren++] = &child;
        }
        const auto compareNItems = [](const QuadTree * lhs, const QuadTree * rhs) {
          return lhs->getNItems() < rhs->getNItems();
        };
 
        // Explicitly count down the children
        for (int nRemaining = nChildren; nRemaining > 0; --nRemaining) {
          auto itHeaviestChild =
            std::max_element(children.begin(), children.begin() + nRemaining, compareNItems);
          QuadTree* heaviestChild = *itHeaviestChild;
          // Drop the heaviest child from the list keeping the order of the remaining ones
          std::move(itHeaviestChild + 1, children.begin() + nRemaining, itHeaviestChild);
          // After we have processed some children we need to get rid of the already used hits in all the children,
          // because this can change the number of items drastically
          erase_remove_if(heaviestChild->getItems(), [&](Item * hit) { return hit->isUsed(); });
          this->fillGivenTree(heaviestChild, candidateReceiver, nItemsThreshold, yLimit);
        }
      }

      void createChildren(QuadTree* node, QuadTreeChildren& m_children) const
      {
        m_children.reserve(node->getXNbins() * node->getYNbins());
        for (int i = 0; i < node->getXNbins(); ++i) {
          for (int j = 0; j < node->getYNbins(); ++j) {
            const XYSpans& xySpans = createChild(node, i, j);
            const XSpan& xSpan = xySpans.first;
            const YSpan& ySpan = xySpans.second;
            m_children.push_back(QuadTree(xSpan, ySpan, node->getLevel() + 1, node));
          }
        }
      }

      void fillChildren(QuadTree* node, const std::vector<Item*>& items)
      {
        // An item can be inserted into each child at most once
        const size_t neededSize = items.size();
        m_childPtrs.clear();
        m_childPtrs.reserve(node->getChildren().size());
        for (QuadTree& child : node->getChildren()) {
          child.reserveItems(neededSize);
          m_childPtrs.push_back(&child);
        }
 
        this->insertItemsInNodes(m_childPtrs, items);
 
        afterFillDebugHook(node->getChildren());
      }

      static void callResultFunction(QuadTree* node, const CandidateReceiver& candidateReceiver)
      {
        const std::vector<Item*>& foundItems = node->getItems();
        std::vector<AData*> candidate;
        candidate.reserve(foundItems.size());
 
        for (Item* item : foundItems) {
          item->setUsedFlag();
          candidate.push_back(item->getPointer());
        }
 
        candidateReceiver(candidate, node);
      }
 
    protected: // Section of specialisable functions
      virtual XYSpans createChild(QuadTree* node, int iX, int iY) const
      {
        AX xMin = node->getXLowerBound(iX);
        AX xMax = node->getXUpperBound(iX);
        AY yMin = node->getYLowerBound(iY);
        AY yMax = node->getYUpperBound(iY);
        return XYSpans({xMin, xMax}, {yMin, yMax});
      }

      virtual bool isInNode(QuadTree* node, AData* item) const = 0;

      virtual void insertItemsInNodes(const std::vector<QuadTree*>& nodes,
                                      const std::vector<Item*>& items)
      {
        for (Item* item : items) {
          if (item->isUsed()) continue;
 
          for (QuadTree* node : nodes) {
            if (isInNode(node, item->getPointer())) {
              node->insertItem(item);
            }
          }
        }
      }

      virtual bool isLeaf(QuadTree* node) const
      {
        if (node->getLevel() >= m_lastLevel) {
          return true;
        } else {
          return false;
        }
      }

      int getLastLevel() const
      {
        return m_lastLevel;
      }
 
    public: // debug stuff
      virtual void afterFillDebugHook(QuadTreeChildren& children)
      {
        if (not m_debugOutput) return;
        for (const QuadTree& childNode : children) {
          if (childNode.getLevel() != getLastLevel()) continue; // Only write the lowest level
          //m_debugOutputMap[ {childNode.getXMean(), childNode.getYMean()}] = childNode.getItems();
        }
      }

      const std::map<std::pair<AX, AY>, std::vector<Item*>>& getDebugInformation() const
      {
        return m_debugOutputMap;
      }
 
    protected:
      std::unique_ptr<QuadTree> m_quadTree;

      std::deque<Item> m_items;

      std::vector<Item*> m_itemPtrs;

      std::vector<QuadTree*> m_childPtrs;

      std::vector<QuadTree*> m_seededTrees;
 
    private:
      int m_lastLevel;

      int m_seedLevel;

      bool m_debugOutput;

      std::map<std::pair<AX, AY>, std::vector<Item*>> m_debugOutputMap;
    };
  }
}