InteractionTree.h

/*
 * InteractionTree.h
 *
 *  Created on: 16 Mar 2017
 *      Author: Valerio Formato
 */

/*
 * A simple templated class to map particle interaction in a tree structure
 * with linked nodes. Particle/position hits can be stored per-track and sorted
 * by an arbitrary function.
 */

#include <algorithm>
#include <unordered_map>
#include <vector>

#ifndef _INTERACTIONTREE_
#define _INTERACTIONTREE_

template <class T> class TreeNode {
public:
  TreeNode();
  TreeNode(int ID);
  ~TreeNode();

  int trackID;
  std::vector<TreeNode *> children;
  std::vector<T> hits;
};

template <class T> class InteractionTree {
public:
  InteractionTree();
  ~InteractionTree();

  TreeNode<T> *GetRootNode() { return hashMap[1]; }
  void AddHit(int trackID, T hit);
  void AddNode(int trackID, int parentID);
  void SortHits(bool (*f)(T, T));
  // void Init();
  void Clear() { hashMap.clear(); };
  void Dump();

  // TreeNode* node;
  std::unordered_map<int, TreeNode<T> *> hashMap;

private:
  typedef typename std::unordered_map<int, TreeNode<T> *>::iterator UOMIterator;
};

template <class T> TreeNode<T>::TreeNode() {}

template <class T> TreeNode<T>::TreeNode(int ID) : trackID(ID) {}

template <class T> TreeNode<T>::~TreeNode() {
  children.clear();
  std::vector<T>().swap(hits); // we don't want to call the destructor of T
}

template <class T> InteractionTree<T>::InteractionTree() {}

template <class T> InteractionTree<T>::~InteractionTree() { Clear(); }

template <class T> void InteractionTree<T>::AddHit(int trackID, T hit) {

  if (hashMap.find(trackID) == hashMap.end())
    AddNode(trackID, hit->parentID);

  hashMap[trackID]->hits.push_back(hit);
}

template <class T> void InteractionTree<T>::AddNode(int trackID, int parentID) {

  std::pair<UOMIterator, bool> mapResult;

  if (hashMap.find(parentID) == hashMap.end()) {
    TreeNode<T> *parentnode = new TreeNode<T>(parentID);
    mapResult = hashMap.insert(std::make_pair(parentID, parentnode)); // let's create an empty parent...
  }

  TreeNode<T> *childnode = new TreeNode<T>(trackID);
  hashMap[parentID]->children.push_back(childnode);
  hashMap.insert(std::make_pair(trackID, childnode));
}

template <class T> void InteractionTree<T>::SortHits(bool (*order)(T, T)) {

  for (std::pair<int, TreeNode<T> *> element : hashMap) {
    std::sort(element.second->hits.begin(), element.second->hits.end(), order);
  }
}

template <class T> void InteractionTree<T>::Dump() {

  for (std::pair<int, TreeNode<T> *> element : hashMap) {
    std::cout << "trackID " << element.first << " has " << element.second->children.size() << " children" << std::endl;
    for (int ic = 0; ic < element.second->children.size(); ic++) {
      std::cout << " - " << element.second->children[ic]->trackID << std::endl;
      for (int ihit = 0; ihit < element.second->children[ic]->hits.size(); ihit++) {
        std::cout << " --- " << element.second->children[ic]->hits[ihit]->particlePdg << "  "
                  << element.second->children[ic]->hits[ihit]->time << "  "
                  << element.second->children[ic]->hits[ihit]->eDep << "  ";
        std::cout << "(" << element.second->children[ic]->hits[ihit]->entrancePoint[0] << ", ";
        std::cout << element.second->children[ic]->hits[ihit]->entrancePoint[1] << ", ";
        std::cout << element.second->children[ic]->hits[ihit]->entrancePoint[2] << ") -> ";
        std::cout << "(" << element.second->children[ic]->hits[ihit]->exitPoint[0] << ", ";
        std::cout << element.second->children[ic]->hits[ihit]->exitPoint[1] << ", ";
        std::cout << element.second->children[ic]->hits[ihit]->exitPoint[2] << ")";
        std::cout << std::endl;
      }
    }
  }
}

#endif