/** \class ExRootTreeReader
 *
 *  Class simplifying access to ROOT tree branches
 *
 *  $Date: 2009-02-02 11:32:01 $
 *  $Revision: 1.2 $
 *
 *
 *  \author S. Ovyn - UCL, Louvain-la-Neuve
 *
 */

#include "ExRootTreeReader.h"
#include "TBranchElement.h"
#include "TClonesArray.h"
#include "TChain.h"
#include <iostream>

using namespace std;

//------------------------------------------------------------------------------

ExRootTreeReader::ExRootTreeReader(TTree *tree) :
  fChain(tree), fCurrentTree(-1), fIsInitDone(kFALSE)
{

}

//------------------------------------------------------------------------------

ExRootTreeReader::~ExRootTreeReader()
{
  TBranchMap::iterator it_map;

  for(it_map = fBranchMap.begin(); it_map != fBranchMap.end(); ++it_map)
  {
    delete it_map->second.second;
  }
}

//------------------------------------------------------------------------------

Bool_t ExRootTreeReader::ReadEntry(Long64_t entry)
{
  // Read contents of entry.
  if(!fChain) return kFALSE;

  if(!fIsInitDone) Init();

  Int_t treeEntry = fChain->LoadTree(entry);
  if(treeEntry < 0) return kFALSE;
  
  if(fChain->IsA() == TChain::Class())
  {
    TChain *chain = static_cast<TChain*>(fChain);
    if(chain->GetTreeNumber() != fCurrentTree)
    {
      fCurrentTree = chain->GetTreeNumber();
      Notify();
    }
  }

  TBranchMap::iterator it_map;
  TBranch *branch;

  for(it_map = fBranchMap.begin(); it_map != fBranchMap.end(); ++it_map)
  {
    branch = it_map->second.first;
    if(branch)
    {
      branch->GetEntry(treeEntry);
    }
  }

  return kTRUE;
}

//------------------------------------------------------------------------------

TClonesArray *ExRootTreeReader::UseBranch(const TString &branchName)
{
  fIsInitDone = kFALSE;

  TClonesArray *array = 0;

  TBranchMap::iterator it_map = fBranchMap.find(branchName);

  if(it_map != fBranchMap.end())
  {
    array = it_map->second.second;
  }
  else
  {
    TBranch *branch = fChain->GetBranch(branchName);
    if(branch)
    {
      if(branch->IsA() == TBranchElement::Class())
      {
        TBranchElement *element = static_cast<TBranchElement*>(branch);
        const char *className = element->GetClonesName();
        Int_t size = element->GetMaximum();
        TClass *cl = gROOT->GetClass(className);
        if(cl)
        {
          array = new TClonesArray(cl, size);
          fBranchMap.insert(make_pair(branchName, make_pair(branch, array)));
        }
      }
    }
  }

  if(!array)
  {
    cout << "** WARNING: cannot access branch '" << branchName << "', return NULL pointer" << endl;
  }

  return array;
}

//------------------------------------------------------------------------------

void ExRootTreeReader::Init()
{
  TBranchMap::iterator it_map;

  for(it_map = fBranchMap.begin(); it_map != fBranchMap.end(); ++it_map)
  {
    fChain->SetBranchAddress(it_map->first, &(it_map->second.second));
  }

  Notify();

  fIsInitDone = kTRUE;
}

//------------------------------------------------------------------------------

void ExRootTreeReader::Notify()
{
  // Called when loading a new file.
  // Get branch pointers.
  if(!fChain) return;

  TBranchMap::iterator it_map;

  for(it_map = fBranchMap.begin(); it_map != fBranchMap.end(); ++it_map)
  {
    it_map->second.first = fChain->GetBranch(it_map->first);
    if(!it_map->second.first)
    {
      cout << "** WARNING: cannot get branch '" << it_map->first << "'" << endl;
    }
  }
}

//------------------------------------------------------------------------------