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 <vector>
#include <unordered_map>

#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