source: trunk/MagicSoft/Mars/mfileio/MReadTree.cc@ 1471

/* ======================================================================== *\
!
! *
! * This file is part of MARS, the MAGIC Analysis and Reconstruction
! * Software. It is distributed to you in the hope that it can be a useful
! * and timesaving tool in analysing Data of imaging Cerenkov telescopes.
! * It is distributed WITHOUT ANY WARRANTY.
! *
! * Permission to use, copy, modify and distribute this software and its
! * documentation for any purpose is hereby granted without fee,
! * provided that the above copyright notice appear in all copies and
! * that both that copyright notice and this permission notice appear
! * in supporting documentation. It is provided "as is" without express
! * or implied warranty.
! *
!
!
!   Author(s): Thomas Bretz  12/2000 <mailto:tbretz@uni-sw.gwdg.de>
!
!   Copyright: MAGIC Software Development, 2000-2002
!
!
\* ======================================================================== */

/////////////////////////////////////////////////////////////////////////////
//                                                                         //
// MReadTree                                                               //
//                                                                         //
// This tasks opens all branches in a specified tree and creates the       //
// corresponding parameter containers if not already existing in the       //
// parameter list.                                                         //
//                                                                         //
// The Process function reads one events from the tree. To go through the  //
// events of one tree make sure that the event number is increased from    //
// outside. It makes also possible to go back by decreasing the number.    //
//                                                                         //
// If you don't want to start reading the first event you have to call     //
// MReadTree::SetEventNum after instantiating your MReadTree-object.       //
//                                                                         //
// To make reading much faster (up to a factor of 10 to 20) you can        //
// ensure that only the data you are really processing is enabled by       //
// calling MReadTree::UseLeaf.                                             //
//                                                                         //
// If the chain switches from one file to another file all                 //
// TObject::Notify() functions are called of TObject objects which were    //
// added to the Notifier list view MReadTree::AddNotify. If MReadTree      //
// is the owner (viw MReadTree::SetOwner) all this objects are deleted     //
// by the destructor of MReadTree                                          //
//                                                                         //
/////////////////////////////////////////////////////////////////////////////
#include "MReadTree.h"

#include <fstream.h>

#include <TFile.h>           // TFile::GetName
#include <TChain.h>
#include <TSystem.h>         // gSystem->ExpandPath
#include <TGProgressBar.h>
#include <TChainElement.h>
#include <TOrdCollection.h>

#include "MLog.h"
#include "MLogManip.h"

#include "MTime.h"
#include "MFilter.h"
#include "MParList.h"
#include "MTaskList.h"

ClassImp(MReadTree);

class MChain : public TChain
{
private:
    Bool_t fNotified;

public:
    MChain() : TChain(), fNotified(kFALSE) {}
    MChain(const char *name, const char *title="") : TChain(name, title), fNotified(kFALSE) {}

    void ResetTree() { fTree = 0; }

    virtual Bool_t Notify() { fNotified = kTRUE; return kTRUE; }
    virtual void   SetNotify(TObject *obj) { fNotify = obj; fNotified = kFALSE; }

    Int_t LoadTree(Int_t entry)
    {
        //
        // This is the code from TChain::LoadTree but skips the
        // notification in LoadTree. If LoadTree raises the notification
        // a flag is set and the notification is done by hand. This
        // is done to be able to catch the return value from Notify. If
        // it has not been successfull -15 is returned.
        // This is to support return values from Notify()/Reinit()
        //
        TObject *notify = GetNotify();

        SetNotify(this);

        Int_t rc = TChain::LoadTree(entry);

        if (rc >= 0 && fNotified && notify)
            if (!notify->Notify())
                rc = -15;

        SetNotify(notify);

        return rc;
    }
};

// --------------------------------------------------------------------------
//
//  Default constructor. It creates an TChain instance which represents the
//  the Tree you want to read and adds the given file (if you gave one).
//  More files can be added using MReadTree::AddFile.
//  Also an empty veto list is created. This list is used if you want to
//  veto (disable or "don't enable") a branch in the tree, it vetos also
//  the creation of the corresponding object.
//  An empty list of TObjects are also created. This objects are called
//  at any time the TChain starts to read from another file.
//
MReadTree::MReadTree(const char *tname, const char *fname,
                     const char *name, const char *title)
    : fNumEntry(0), fBranchChoosing(kFALSE), fAutoEnable(kTRUE), fProgress(NULL)
{
    fName  = name  ? name  : "MReadTree";
    fTitle = title ? title : "Task to loop over all events in one single tree";

    fVetoList = new TList;
    fVetoList->SetOwner();

    fNotify = new TList;

    //
    // open the input stream
    //
    fChain = new MChain(tname);

    // root 3.02:
    // In TChain::Addfile remove the limitation that the file name must contain
    // the string ".root". ".root" is necessary only in case one wants to specify
    // a Tree in a subdirectory of a Root file with eg, the format:

    if (fname)
        fChain->Add(fname);
}

// --------------------------------------------------------------------------
//
// Destructor. It deletes the TChain and veto list object
//
MReadTree::~MReadTree()
{
    //
    // Delete all the pointers to pointers to the objects where the
    // branche data gets stored.
    //
    TIter Next(fChain->GetStatus());

    TChainElement *element = NULL;
    while ((element=(TChainElement*)Next()))
        delete (MParContainer**)element->GetBaddress();

    //
    // Delete the chain and the veto list
    //
#if ROOT_VERSION_CODE < ROOT_VERSION(3,03,00)
    if (fChain->GetFile())
        delete fChain->GetFile();
#endif
    delete fChain;

    delete fNotify;
    delete fVetoList;
}

// --------------------------------------------------------------------------
//
//  If the owner flag is set all TObjects which are scheduled via
//  AddNotify are deleted by the destructor of MReadTree
//
void MReadTree::SetOwner(Bool_t flag)
{
    flag ? fNotify->SetBit(kIsOwner) : fNotify->ResetBit(kIsOwner);
}

// --------------------------------------------------------------------------
//
//  This function is called each time MReadTree changes the file to read
//  from. It calls all TObject::Notify() functions which are scheduled
//  via AddNotify.
//
Bool_t MReadTree::Notify()
{
    *fLog << inf << GetDescriptor() << ": Notify '" << fChain->GetName();
    *fLog << "' (before processing event #" << GetEventNum()-1 << ")" << endl;

    //fNotify->Notify();

    return kTRUE;
}

// --------------------------------------------------------------------------
//
//  If you want to read the given tree over several files you must add
//  the files here before PreProcess is called. Be careful: If the tree
//  doesn't have the same contents (branches) it may confuse your
//  program (trees which are are not existing in later files are not read
//  anymore, tree wich are not existing in the first file are never read)
//
//  Name may use the wildcarding notation, eg "xxx*.root" means all files
//  starting with xxx in the current file system directory.
//
//  AddFile returns the number of files added to the chain.
//
Int_t MReadTree::AddFile(const char *fname)
{
    //
    // FIXME! A check is missing whether the file already exists or not.
    //
    //
    // returns the number of file which were added
    //
    return fChain->Add(fname);
}

// --------------------------------------------------------------------------
//
//  This function is called if Branch choosing method should get enabled.
//  Branch choosing means, that only the enabled branches are read into
//  memory. To use an enableing scheme we have to disable all branches first.
//  This is done, if this function is called the first time.
//
void MReadTree::EnableBranchChoosing()
{
    if (fBranchChoosing)
        return;

    *fLog << inf << GetDescriptor() << ": Branch choosing method enabled (only enabled branches are read)." << endl;
    fChain->SetBranchStatus("*", kFALSE);
    fBranchChoosing = kTRUE;
}

// --------------------------------------------------------------------------
//
// The first time this function is called all branches are disabled.
// The given branch is enabled. By enabling only the branches you
// are processing you can speed up your calculation many times (up to
// a factor of 10 or 20)
//
void MReadTree::EnableBranch(const char *name)
{
    EnableBranchChoosing();

    TNamed branch(name, "");
    SetBranchStatus(&branch, kTRUE);
}

// --------------------------------------------------------------------------
//
// Set branch status of branch name
//
void MReadTree::SetBranchStatus(const char *name, Bool_t status)
{
    fChain->SetBranchStatus(name, status);

    *fLog << inf << (status ? "Enabled" : "Disabled");
    *fLog << " subbranch '" << name << "'." << endl;
}

// --------------------------------------------------------------------------
//
// Checks whether a branch with the given name exists in the chain
// and sets the branch status of this branch corresponding to status.
//
void MReadTree::SetBranchStatus(TObject *branch, Bool_t status)
{
    //
    // Get branch name
    //
    const char *name = branch->GetName();

    //
    // Check whether this branch really exists
    //
    if (fChain->GetBranch(name))
        SetBranchStatus(name, status);

    //
    // Remove trailing '.' if one and try to enable the subbranch without
    // the master branch name. This is to be compatible with older mars
    // and camera files.
    //
    const char *dot = strrchr(name, '.');
    if (!dot)
        return;

    if (fChain->GetBranch(dot+1))
        SetBranchStatus(dot+1, status);
}

// --------------------------------------------------------------------------
//
// Set the status of all branches in the list to status.
//
void MReadTree::SetBranchStatus(const TList *list, Bool_t status)
{
    //
    // Loop over all subbranches in this master branch
    //
    TIter Next(list);

    TObject *obj;
    while ((obj=Next()))
        SetBranchStatus(obj, status);
}

// --------------------------------------------------------------------------
//
//  This is the implementation of the Auto Enabling Scheme.
//  For more information see MTask::AddBranchToList.
//  This function loops over all tasks and its filters in the tasklist
//  and enables all branches which are requested by the tasks and its
//  filters.
//
//  To enable 'unknown' branches which are not in the branchlist of
//  the tasks you can call EnableBranch
//
void MReadTree::EnableBranches(MParList *plist)
{
    //
    // check whether branch choosing must be switched on
    //
    EnableBranchChoosing();

    //
    // request the tasklist from the parameter list.
    // FIXME: Tasklist can have a different name
    //
    const MTaskList *tlist = (MTaskList*)plist->FindObject("MTaskList");
    if (!tlist)
    {
        *fLog << warn << GetDescriptor() << "Cannot use auto enabeling scheme for branches. 'MTaskList' not found." << endl;
        return;
    }

    //
    // This loop is not necessary. We could do it like in the commented
    // loop below. But this loop makes sure, that we don't try to enable
    // one branch several times. This would not harm, but we would get
    // an output for each attempt. To have several outputs for one subbranch
    // may confuse the user, this we don't want.
    // This loop creates a new list of subbranches and for each branch
    // which is added we check before whether it already exists or not.
    //
    TList list;

    MTask *task;
    TIter NextTask(tlist->GetList());
    while ((task=(MTask*)NextTask()))
    {
        TObject *obj;

        TIter NextTBranch(task->GetListOfBranches());
        while ((obj=NextTBranch()))
            if (!list.FindObject(obj->GetName()))
                list.Add(obj);

        const MFilter *filter = task->GetFilter();

        if (!filter)
            continue;

        TIter NextFBranch(filter->GetListOfBranches());
        while ((obj=NextFBranch()))
            if (!list.FindObject(obj->GetName()))
                list.Add(obj);
    }

    SetBranchStatus(&list, kTRUE);
/*
    //
    // Loop over all tasks iand its filters n the task list.
    //
    MTask *task;
    TIter NextTask(tlist->GetList());
    while ((task=(MTask*)NextTask()))
    {
        SetBranchStatus(task->GetListOfBranches(), kTRUE);

        const MFilter *filter = task->GetFilter();
        if (!filter)
           continue;

        SetBranchStatus(filter->GetListOfBranches(), kTRUE);

    }
*/
}

// --------------------------------------------------------------------------
//
//  The disables all subbranches of the given master branch.
//
void MReadTree::DisableSubBranches(TBranch *branch)
{
    //
    // This is not necessary, it would work without. But the output
    // may confuse the user...
    //
    if (fAutoEnable || fBranchChoosing)
        return;

    SetBranchStatus(branch->GetListOfBranches(), kFALSE);
}

// --------------------------------------------------------------------------
//
//  The PreProcess loops (till now) over the branches in the given tree.
//  It checks if the corresponding containers (containers with the same
//  name than the branch name) are existing in the Parameter Container List.
//  If not, a container of objec type 'branch-name' is created (everything
//  after the last semicolon in the branch name is stripped). Only
//  branches which don't have a veto (see VetoBranch) are enabled If the
//  object isn't found in the root dictionary (a list of classes known by the
//  root environment) the branch is skipped and an error message is printed
//  out.
//
Bool_t MReadTree::PreProcess(MParList *pList)
{
    //
    // Make sure, that all the following calls doesn't result in
    // Notifications. This may be dangerous, because the notified
    // tasks are not preprocessed.
    //
    fChain->SetNotify(NULL);

    //
    // get number of events in this tree
    //
    fNumEntries = (UInt_t)fChain->GetEntries();

    if (!fNumEntries)
    {
        *fLog << warn << dbginf << "No entries found in file(s)" << endl;
        return kFALSE;
    }

    //
    // output logging information
    //
    *fLog << inf << fNumEntries << " entries found in file(s)." << endl;

    //
    // Get all branches of this tree and
    // create the Iterator to loop over all branches
    //
    TIter Next(fChain->GetListOfBranches());
    TBranch *branch=NULL;

    Int_t num=0;
    //
    // loop over all tasks for processing
    //
    while ( (branch=(TBranch*)Next()) )
    {
        //
        // Get Name of Branch and Object
        //
        const char *bname = branch->GetName();

        TString oname(bname);
        if (oname.EndsWith("."))
            oname.Remove(oname.Length()-1);

        //
        // Check if enabeling the branch is allowed
        //
        if (fVetoList->FindObject(oname))
        {
            *fLog << inf << "Master branch " << bname << " has veto... skipped." << endl;
            DisableSubBranches(branch);
            continue;
        }

        //
        // Create a pointer to the pointer to the object in which the
        // branch data is stored. The pointers are stored in the TChain
        // object and we get the pointers from there to delete it.
        //
        MParContainer **pcont= new MParContainer*;

#if ROOT_VERSION_CODE < ROOT_VERSION(3,02,06)
        const char *classname = oname;
#else
        const char *classname = branch->GetClassName();
#endif

        //
        // check if object is existing in the list
        //
        *pcont=pList->FindCreateObj(classname, oname);

        if (!*pcont)
        {
            //
            // if class is not existing in the (root) environment
            // we cannot proceed reading this branch
            //
            *fLog << warn << dbginf << "Warning: Class '" << classname;
            *fLog << "' for " << oname << " not existing in dictionary. Branch skipped." << endl;
            DisableSubBranches(branch);
            continue;
        }

        //
        // Check whether a Pointer to a pointer already exists, if
        // we created one already delete it.
        //
        TChainElement *element = (TChainElement*)fChain->GetStatus()->FindObject(bname);
        if (element)
            delete (MParContainer**)element->GetBaddress();

        //
        // here pcont is a pointer the to container in which the data from
        // the actual branch should be stored - enable branch.
        //
        fChain->SetBranchAddress(bname, pcont);

        *fLog << inf << "Master branch address " << bname << " [";
        *fLog << classname << "] setup for reading." << endl;

        //*fLog << "Branch " << bname << " autodel: " << (int)branch->IsAutoDelete() << endl;
        //branch->SetAutoDelete();

        num++;
    }

    *fLog << inf << GetDescriptor() << " setup " << num << " master branches addresses." << endl;

    //
    // If auto enabling scheme isn't disabled, do auto enabling
    //
    if (fAutoEnable)
        EnableBranches(pList);

    //
    // If a progress bar is given set its range.
    //
    if (fProgress)
        fProgress->SetRange(0, fNumEntries);

    //
    // Now we can start notifying. Reset tree makes sure, that TChain thinks
    // that the correct file is not yet initialized and reinitilizes it
    // as soon as the first event is read. This is necessary to call
    // the notifiers when the first event is read, but after the
    // PreProcess-function.
    //
    fChain->ResetTree();
    fChain->SetNotify(this);

    return kTRUE;
}

// --------------------------------------------------------------------------
//
//  Set the ready to save flag of all containers which branchaddresses are
//  set for. This is necessary to copy data.
//
void MReadTree::SetReadyToSave(Bool_t flag)
{
    TIter Next(fChain->GetStatus());

    TChainElement *element = NULL;
    while ((element=(TChainElement*)Next()))
    {
        //
        // Check whether the branch is enabled
        //
        if (!element->GetStatus())
            continue;

        //
        // Get the pointer to the pointer of the corresponding container
        //
        MParContainer **pcont = (MParContainer**)element->GetBaddress();

        //
        // Check whether the pointer is not NULL
        //
        if (!pcont || !*pcont)
            continue;

        //
        // Set the ready to save status of the container.
        //
        (*pcont)->SetReadyToSave(flag);
    }

    //
    // Set the ready to save status of this task (used?), too
    //
    MTask::SetReadyToSave(flag);
}

// --------------------------------------------------------------------------
//
//  The Process-function reads one event from the tree (this contains all
//  enabled branches) and increases the position in the file by one event.
//  (Remark: The position can also be set by some member functions
//  If the end of the file is reached the Eventloop is stopped.
//
#if ROOT_VERSION_CODE < ROOT_VERSION(3,02,06)
#include "MRawEvtData.h"
#endif
Bool_t MReadTree::Process()
{
    //
    // This is necessary due to a bug in TChain::LoadTree in root.
    // will be fixed in 3.03
    //
#if ROOT_VERSION_CODE < ROOT_VERSION(3,03,01)
    if (fNumEntry >= fNumEntries)
        return kFALSE;
#endif

#if ROOT_VERSION_CODE < ROOT_VERSION(3,02,06)
    //
    // This fixes 99.9% of a memory leak using a root version prior
    // to 3.02/??
    //
    TChainElement *element=NULL;
    TIter Next(fChain->GetStatus());
    while ((element=(TChainElement*)Next()))
    {
        MParContainer **c = (MParContainer**)element->GetBaddress();
        if (!c) continue;
        if ((*c)->InheritsFrom(MRawEvtData::Class()))
            ((MRawEvtData*)(*c))->DeletePixels(kFALSE);

    }
#endif

    Bool_t rc = fChain->GetEntry(fNumEntry++) != 0;

    if (rc)
        SetReadyToSave();

    return rc;
}

// --------------------------------------------------------------------------
//
// Get the Event with the current EventNumber fNumEntry
//
Bool_t MReadTree::GetEvent()
{
    Bool_t rc = fChain->GetEntry(fNumEntry) != 0;

    if (rc)
        SetReadyToSave();

    return rc;
}

// --------------------------------------------------------------------------
//
// Decrease the number of the event which is read by Process() next
// by one or more
//
Bool_t MReadTree::DecEventNum(UInt_t dec)
{
    if (fNumEntry-dec >= fNumEntries)
    {
        *fLog << warn << GetDescriptor() << ": DecEventNum, WARNING - Event " << fNumEntry << "-";
        *fLog << dec << "=" << (Int_t)fNumEntry-dec << " out of Range." << endl;
        return kFALSE;
    }

    fNumEntry -= dec;
    return kTRUE;
}

// --------------------------------------------------------------------------
//
// Increase the number of the event which is read by Process() next
// by one or more
//
Bool_t MReadTree::IncEventNum(UInt_t inc)
{
    if (fNumEntry+inc >= fNumEntries)
    {
        *fLog << warn << GetDescriptor() << ": IncEventNum, WARNING - Event " << fNumEntry << "+";
        *fLog << inc << "=" << (Int_t)fNumEntry+inc << " out of Range." << endl;
        return kFALSE;
    }

    fNumEntry += inc;
    return kTRUE;
}

// --------------------------------------------------------------------------
//
// This function makes Process() read event number nr next
//
// Remark: You can use this function after instatiating you MReadTree-object
//         to set the event number from which you want to start reading.
//
Bool_t MReadTree::SetEventNum(UInt_t nr)
{
    if (nr >= fNumEntries)
    {
        *fLog << warn << GetDescriptor() << ": SetEventNum, WARNING - " << nr << " out of Range." << endl;
        return kFALSE;
    }

    fNumEntry = nr;
    return kTRUE;
}

// --------------------------------------------------------------------------
//
//  For the branch with the given name:
//   1) no object is automatically created
//   2) the branch address for this branch is not set
//      (because we lack the object, see 1)
//   3) The whole branch (exactly: all its subbranches) are disabled
//      this means are not read in memory by TTree:GetEntry
//
void MReadTree::VetoBranch(const char *name)
{
    fVetoList->Add(new TNamed(name, ""));
}

// --------------------------------------------------------------------------
//
//  Return the name of the file we are actually reading from.
//
TString MReadTree::GetFileName() const
{
    const TFile *file = fChain->GetFile();

    if (!file)
        return TString("<unknown>");

    TString name(file->GetName());
    name.Remove(0, name.Last('/')+1);
    return name;
}

// --------------------------------------------------------------------------
//
//  Return the number of the file in the chain, -1 in case of an error
//
Int_t MReadTree::GetFileIndex() const
{
    return fChain->GetTreeNumber();
    /*
    const TString filename = fChain->GetFile()->GetName();

    int i=0;
    TObject *file = NULL;

    TIter Next(fChain->GetListOfFiles());
    while ((file=Next()))
    {
        if (filename==gSystem->ExpandPathName(file->GetTitle()))
            return i;
        i++;
    }
    return -1;
    */
}

// --------------------------------------------------------------------------
//
//  This schedules a TObject which Notify(9 function is called in case
//  of MReadTree (TChain) switches from one file in the chain to another
//  one.
//
void MReadTree::AddNotify(TObject *obj)
{
    fNotify->Add(obj);
}

void MReadTree::Print(Option_t *o) const
{
    *fLog << all << GetDescriptor() << dec << endl;
    *fLog << setfill('-') << setw(strlen(GetDescriptor())) << "" << endl;
    *fLog << " Files [Tree]:" << endl;

    int i = 0;
    TIter Next(fChain->GetListOfFiles());
    TObject *obj = NULL;
    while ((obj=Next()))
        *fLog << " " << i++ << ") " << obj->GetTitle() << " [" << obj->GetName() << "]" << endl;

    *fLog << " Total Number of Entries: " << fNumEntries << endl;
    *fLog << " Next Entry to read: " << fNumEntry << endl;
}

void MReadTree::SavePrimitive(ofstream &out, Option_t *o="")
{
    TString name = ToLower(fName);

    out << "   " << ClassName() << " " << name << "(";
    out << fChain->GetName() << ", \"" << fName << "\", \"" << fTitle << "\");" << endl;

    TIter Next(fChain->GetListOfFiles());
    TObject *obj = NULL;
    while ((obj=Next()))
        out << "   " << name << ".AddFile(\"" << obj->GetTitle() << "\");" << endl;

    if (!fAutoEnable)
        out << "   " << name << ".DisableAutoScheme();" << endl;

    if (fNumEntry!=0)
       out << "   " << name << ".SetEventNum(" << fNumEntry << ");" << endl;


}