SingleTreeChain.h

Go to the documentation of this file.

#ifndef NAIA_SINGLETREECHAIN_H
#define NAIA_SINGLETREECHAIN_H

// ROOT headers
#include "TChain.h"
#include "TTree.h"

// NAIA headers
#include "Chain/Event.h"
#include "Chain/SkimTreeHandle.h"
#include "Containers/FileInfo.h"
#include "Containers/RTIInfo.h"

namespace NAIA {

class SingleTreeChain {
  friend class NtpSelector;
  friend class SkimTreeHandle<SingleTreeChain>;

public:
  enum class AccessMode { Read, Write };

  SingleTreeChain(AccessMode mode = AccessMode::Read);

  const FileInfo &GetEventFileInfo();

  const MCFileInfo &GetEventMCFileInfo();

  const RTIInfo &GetEventRTIInfo();

  Event &GetEvent(unsigned long long iEv);

  Event &GetEventWithIndex(unsigned int run, unsigned int eventno);

  TChain *GetFileInfoTree();

  TChain *GetRTITree();

  void SetupBranches(bool isMC = false);

  SkimTreeHandle<SingleTreeChain> CreateSkimTree(const std::string &filename, const std::string &exclBranches);

  int Add(const std::string &filePath);

  void Clear() { m_event.Clear(); }

  int Fill();

  int FillRTI();

  int FillFileInfo();

  unsigned long int GetEntries();

  bool IsMC() { return m_data.chain->GetBranch(MCTruthBase::BranchName.c_str()); };

  void Print();

  int Write();

  void SetDirectory(TDirectory *directory);

  void SetEntryList(TEntryList *entryList, Option_t *option);

  class EventItr {
    friend bool operator==(const EventItr &lhs, const EventItr &rhs);

  public:
    EventItr(unsigned long long entry, SingleTreeChain *chain) : m_entry{entry}, m_chain{chain} {};

    // prefix increment operator
    EventItr &operator++() {
      m_entry++;
      return *this;
    }
    // postfix increment operator
    EventItr operator++(int) {
      EventItr old = *this;
      operator++();
      return old;
    }
    // prefix decrement operator
    EventItr &operator--() {
      m_entry--;
      return *this;
    }
    // postfix decrement operator
    EventItr operator--(int) {
      EventItr old = *this;
      operator--();
      return old;
    }

    Event &operator*() { return m_chain->GetEvent(m_entry); }

  private:
    unsigned long long m_entry = 0;
    SingleTreeChain *m_chain;
  };

  EventItr begin() { return {0, this}; }
  EventItr end() { return {GetEntries() - 1, this}; }

private:
  unsigned long int m_iEv = 0;                
  AccessMode m_accessMode = AccessMode::Read; 

  FileInfo m_fileInfo;     
  MCFileInfo m_fileInfoMC; 
  RTIInfo m_rtiInfo;       
  Event m_event;           

  FileInfo *m_fileInfoPtr = &m_fileInfo;       
  MCFileInfo *m_fileInfoMCPtr = &m_fileInfoMC; 
  RTIInfo *m_rtiInfoPtr = &m_rtiInfo;          

  // in c++17 we would use a std::variant
  union TC {
    TTree *tree;
    TChain *chain;
  };

  TC m_data{}; 
  TC m_rti{};  
  TC m_file{}; 
};

inline bool operator==(const SingleTreeChain::EventItr &lhs, const SingleTreeChain::EventItr &rhs) {
  return (lhs.m_entry == rhs.m_entry && lhs.m_chain == rhs.m_chain);
};

inline bool operator!=(const SingleTreeChain::EventItr &lhs, const SingleTreeChain::EventItr &rhs) {
  return !(lhs == rhs);
};

inline SingleTreeChain::EventItr begin(SingleTreeChain &chain) { return {0, &chain}; }
inline SingleTreeChain::EventItr end(SingleTreeChain &chain) { return {chain.GetEntries() - 1, &chain}; }

} // namespace NAIA

#endif