source: trunk/MagicSoft/Mars/mfileio/MWriteRootFile.cc@ 5245

/* ======================================================================== *\
!
! *
! * 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, 6/2001 <mailto:tbretz@astro.uni-wuerzburg.de>
!
!   Copyright: MAGIC Software Development, 2000-2004
!
!
\* ======================================================================== */

/////////////////////////////////////////////////////////////////////////////
//
// MWriteRootFile
//
// This is a writer to store several containers to a root file.
// The containers are added with AddContainer.
// To understand how it works, see base class MWriteFile
//
// Warning: Look at the Warning in MTaskList.
//
// There is a special mode of operation which opens a new file for each new
// file read by the reading task (opening the new file is initiated by
// ReInit()) For more details se the corresponding constructor.
//
/////////////////////////////////////////////////////////////////////////////
#include "MWriteRootFile.h"

#include <fstream>

#include <TFile.h>
#include <TTree.h>
#include <TRegexp.h>

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

#include "MRead.h"
#include "MParList.h"

ClassImp(MRootFileBranch);
ClassImp(MWriteRootFile);

using namespace std;

const TString MWriteRootFile::gsDefName  = "MWriteRootFile";
const TString MWriteRootFile::gsDefTitle = "Task which writes a root-output file";

void MWriteRootFile::Init(const char *name, const char *title)
{
    fName  = name  ? name  : gsDefName.Data();
    fTitle = title ? title : gsDefTitle.Data();

    //
    // Set the Arrays the owner of its entries. This means, that the
    // destructor of the arrays will delete all its entries.
    //
    fBranches.SetOwner();

    //
    // Believing the root user guide, TTree instanced are owned by the
    // directory (file) in which they are. This means we don't have to
    // care about their destruction.
    //
    //fTrees.SetOwner();
}

// --------------------------------------------------------------------------
//
// Default constructor. It is there to support some root stuff.
// Don't use it.
//
MWriteRootFile::MWriteRootFile() : fOut(NULL)
{
    Init();

    //
    // Set the Arrays the owner of its entries. This means, that the
    // destructor of the arrays will delete all its entries.
    //
    fBranches.SetOwner();
}

// --------------------------------------------------------------------------
//
// Use this constructor to run in a special mode.
//
// In this mode for each input file a new output file is written. This
// happens in ReInit.
//
// comp:        Compression Level (see TFile, TBranch)
// rule:        Rule to create output file name (see GetNewFileName())
// overwrite:   Allow newly created file to overwrite old files ("RECREATE")
// ftitle:      File title stored in the file (see TFile)
// name, title: Name and title of this object
//
MWriteRootFile::MWriteRootFile(const Int_t comp,
                               const char *rule,
                               const Bool_t overwrite,
                               const char *ftitle,
                               const char *name,
                               const char *title) : fSplitRule(rule)
{
    Init(name, title);

    //
    // Open a TFile in dummy mode! This is necessary to be able to create
    // the trees and branches, which are then (in ReInit) moved to
    // a valid file. (Stupid workaround - but does a good job)
    //
    fOut = new TFile("/dev/null", overwrite?"RECREATE":"NEW", ftitle, comp);
}

// --------------------------------------------------------------------------
//
// Specify the name of the root file. You can also give an option ("UPDATE"
// and "RECREATE" would make sense only) as well as the file title and
// compression factor. To a more detaild description of the options see
// TFile.
//
MWriteRootFile::MWriteRootFile(const char *fname,
                               const Option_t *opt,
                               const char *ftitle,
                               const Int_t comp,
                               const char *name,
                               const char *title)
{
    Init(name, title);

    //
    // If no name is given we open the TFile in some kind of dummy mode...
    //
    if (!fname)
    {
        fOut = new TFile("/dev/null", "READ", ftitle, comp);
        return;
    }

    TString str(fname);
    if (!str.EndsWith(".root", TString::kIgnoreCase))
        str += ".root";

    //
    // Open the rootfile
    //
    fOut = new TFile(str, opt, ftitle, comp);
}

// --------------------------------------------------------------------------
//
// Prints some statistics about the file to the screen. And closes the file
// properly.
//
void MWriteRootFile::Close()
{
    //
    // Print some statistics to the looging out.
    //
    Print();

    //
    // If the file is still open (no error) write the keys. This is necessary
    // for appearance of the all trees and branches.
    //
    if (IsFileOpen())
        fOut->Write();

    //
    // Delete the file. This'll also close the file (if open)
    //
    delete fOut;
    fOut = 0;

    //
    // Remark:
    // - Trees are automatically deleted by the the file
    //   (unless file.SetDirectory(0) was called)
    // - Branches are automatically deleted by the tree destructor
    //

    *fLog << inf << "Output File closed and object deleted." << endl;
}

// --------------------------------------------------------------------------
//
// call Close()
//
MWriteRootFile::~MWriteRootFile()
{
    Close();
}

// --------------------------------------------------------------------------
//
// Prints all trees with the actually number of written entries to log-out.
//
void MWriteRootFile::Print(Option_t *) const
{
    *fLog << all << underline << "File: " << GetFileName() << endl;

    if (fTrees.GetEntries()==0)
    {
        *fLog << " No contents." << endl;
        return;
    }

    TObject *obj;
    TIter NextBranch(&fBranches);
    while ((obj=NextBranch()))
    {
        MRootFileBranch *b = (MRootFileBranch*)obj;

        if (!b->GetTree() || b->GetTree()->TestBit(kIsNewTree))
            continue;

        TBranch *branch = b->GetBranch();

        TString name = b->GetTree()->GetName();
        name += '.';
        name += branch->GetName();

        *fLog << " " << name.Strip(TString::kTrailing, '.') << ": \t" << (ULong_t)branch->GetEntries() << " entries." << endl;
    }

    TTree *t = NULL;
    TIter NextTree(&fTrees);
    while ((t=(TTree*)NextTree()))
        if (t->TestBit(kIsNewTree))
            *fLog << " " << t->GetName() << ": \t" << (ULong_t)t->GetEntries() << " entries." << endl;
    *fLog << endl;
}

// --------------------------------------------------------------------------
//
// Add a new Container to list of containers which should be written to the
// file. Give the name of the container which will identify the container
// in the parameterlist. tname is the name of the tree to which the
// container should be written (Remark: one tree can hold more than one
// container). The default is the same name as the container name.
// You can slso specify a title for the tree. This is only
// used the first time this tree in 'mentioned'. As default the title
// is the name of the tree.
//
void MWriteRootFile::AddContainer(const char *cname, const char *tname, Bool_t must)
{
    TIter Next(&fBranches);
    TObject *o=0;
    while ((o=Next()))
        if (TString(o->GetName())==TString(tname) && TString(o->GetTitle())==TString(cname))
        {
            *fLog << "WARNING - Container '" << cname << "' in Tree '" << tname << "' already scheduled... ignored." << endl;
            return;
        }

    //
    // create a new entry in the list of branches to write and
    // add the entry to the list.
    //
    MRootFileBranch *entry = new MRootFileBranch(AddSerialNumber(cname), tname, must);
    fBranches.AddLast(entry);

    if (tname && tname[0])
        AddToBranchList(Form("%s.%s", (const char*)AddSerialNumber(cname), tname));
}

// --------------------------------------------------------------------------
//
// Add a new Container to list of containers which should be written to the
// file. Give the pointer to the container. tname is the name of the tree to
// which the container should be written (Remark: one tree can hold more than
// one container). The default is the same name as the container name.
// You can slso specify a title for the tree. This is only
// used the first time this tree in 'mentioned'. As default the title
// is the name of the tree.
//
void MWriteRootFile::AddContainer(MParContainer *cont, const char *tname,
                                  Bool_t must)
{
    TIter Next(&fBranches);
    TObject *o=0;
    while ((o=Next()))
        if (TString(o->GetName())==TString(tname) &&
            static_cast<MRootFileBranch*>(o)->GetContainer()==cont)
        {
            *fLog << "WARNING - Container " << cont << " in Tree '" << tname << "' already scheduled... ignored." << endl;
            return;
        }

    //
    // create a new entry in the list of branches to write and
    // add the entry to the list.
    //
    MRootFileBranch *entry = new MRootFileBranch(cont, tname, must);
    fBranches.AddLast(entry);
}

// --------------------------------------------------------------------------
//
// Add a new Container to list of containers which should be written to the
// file. Give the pointer to the container. tname is the name of the tree to
// which the container should be written (Remark: one tree can hold more than
// one container). The default is the same name as the container name.
// You can slso specify a title for the tree. This is only
// used the first time this tree in 'mentioned'. As default the title
// is the name of the tree.
//
Bool_t MWriteRootFile::GetContainer(MParList *pList)
{
    //
    // loop over all branches which are 'marked' as branches to get written.
    //
    MRootFileBranch *entry;

    TIter Next(&fBranches);
    while ((entry=(MRootFileBranch*)Next()))
    {
        //
        // Get the pointer to the container. If the pointer is NULL it seems,
        // that the user identified the container by name.
        //
        MParContainer *cont = entry->GetContainer();
        if (!cont)
        {
            //
            // Get the name and try to find a container with this name
            // in the parameter list.
            //
            const char *cname = entry->GetContName();
            cont = (MParContainer*)pList->FindObject(cname);
            if (!cont)
            {
                //
                // No corresponding container is found
                //
                if (entry->MustHave())
                {
                    *fLog << err << "Cannot find parameter container '" << cname << "'." << endl;
                    return kFALSE;
                }

                *fLog << inf << "Unnecessary parameter container '" << cname << "' not found..." << endl;
                delete fBranches.Remove(entry);
                continue;
            }

            //
            // The container is found. Put the pointer into the entry.
            //
            entry->SetContainer(cont);
        }

        //
        // Get container name, tree name and tree title of this entry.
        //
        const char *cname = cont->GetName();
        const char *tname = entry->GetName();
        const TString ttitle(Form("Tree containing %s", cont->GetDescriptor()));

        //
        // if the tree name is NULL this idetifies it to use the default:
        // the container name.
        //
        if (tname[0] == '\0')
            tname = cname;

        //
        // Check if the tree is already existing (part of the file)
        //
        TTree *tree = (TTree*)fOut->Get(tname);
        if (!tree)
        {
            //
            // if the tree doesn't exist create a new tree. Use the tree
            // name as title if title is NULL.
            // And add the tree to the list of trees
            //
            TDirectory *save = gDirectory;
            fOut->cd();

            tree = new TTree(tname, ttitle);
            fTrees.AddLast(tree);

            //
            // If the tree does not already exist in the file mark this
            // tree as a branch created by MWriteRootFile
            //
            tree->SetBit(kIsNewTree);

            gDirectory = save;

            *fLog << inf << "Tree " << tname << " created." << endl;
        }

        //
        // In case the file is opened as 'UPDATE' the tree may still not
        // be in the list. Because it neither was created previously,
        // nor this time, so the corresponding entries is marked as a
        // single branch to be filled. --> Add it to the list of trees.
        //
        if (!fTrees.FindObject(tree))
            fTrees.AddLast(tree);

        //
        // Now we have a valid tree. Search the list of trees for this tree
        // Add a pointer to the entry in the tree list to this branch-entry
        //
        entry->SetTree(tree);

        TString branchname(cname);
        branchname.Append(".");

        //
        // Try to get the branch from the file. 
        // If the branch already exists the user specified one branch twice.
        //
        TBranch *branch = tree->GetBranch(branchname);
        if (branch)
        {
            *fLog << inf << "Branch '" << cname << "' already existing... updating." << endl;
            branch->SetAddress(entry->GetAddress());

            if (!fSplitRule.IsNull())
            {
                *fLog << warn << endl;
                *fLog << "WARNING: You are updating an existing branch. For this case" << endl;
                *fLog << "         file-splitting mode  is  not  allowed...  disabled!" << endl;
                *fLog << endl;
                fSplitRule = "";
            }
        }
        else
        {
            //
            // Create a new branch in the actual tree. The branch has the name
            // container name. The type of the container is given by the
            // ClassName entry in the container. The Address is the address of a
            // pointer to the container (gotten from the branch entry). As
            // Basket size we specify a (more or less) common default value.
            // The containers should be written in Splitlevel=1
            //
            *fLog << inf << "Creating Branch " << cname << " of " << cont->ClassName();
            *fLog << " in tree " << tree->GetName() << "... " << flush;

            branch = tree->Branch(branchname, cont->ClassName(), entry->GetAddress());

            //
            // If the branch couldn't be created we have a problem.
            //
            if (!branch)
            {
                *fLog << endl;
                *fLog << err << "Unable to create branch '" << cname << "'." << endl;
                return kFALSE;
            }

            *fLog << "done." << endl;

            if (!tree->TestBit(kIsNewTree) && !fSplitRule.IsNull())
            {
                *fLog << warn << endl;
                *fLog << "WARNING:   You  have  created  a new branch  in  an existing tree." << endl;
                *fLog << "     For this case file-splitting mode is not allowed... disabled!" << endl;
                *fLog << endl;
                fSplitRule= "";
            }
        }

        //
        // Tell the entry also which branch belongs to it (this is necessary
        // for branches belonging to already existing tree, UPDATE-mode)
        //
        entry->SetBranch(branch);
    }

    return kTRUE;
}

// --------------------------------------------------------------------------
//
// Checks all given containers (branch entries) for the write flag.
// If the write flag is set the corresponding Tree is marked to get filled.
// All Trees which are marked to be filled are filled with all their
// branches.
// In case of a file opened in 'UPDATE' mode, single branches can be
// filled, too. WARNING - for the moment there is no check whether
// you filled the correct number of events into the branch, so that
// each of the other branches in the tree has the correct corresponding
// number of new entries in the new branch!
// Be carefull: If only one container (corresponding to a branch) of a tree
// has the write flag, all containers in this tree are filled!
//
Bool_t MWriteRootFile::CheckAndWrite()
{
    TObject *obj;

    //
    // Loop over all branch entries
    //
    TIter NextBranch(&fBranches);
    while ((obj=NextBranch()))
    {
        MRootFileBranch *b = (MRootFileBranch*)obj;

        //
        // Check for the Write flag
        //
        if (!b->GetContainer()->IsReadyToSave())
            continue;

        //
        // If the write flag of the branch entry is set, set the write flag of
        // the corresponding tree entry.
        //
        if (b->GetTree()->TestBit(kIsNewTree))
            b->GetTree()->SetBit(kFillTree);
        else
        {
            if (!b->GetBranch()->Fill())
            {
                *fLog << err << "ERROR - Zero bytes written to branch '" << b->GetBranch()->GetName() << "'... abort." << endl;
                return kFALSE;
            }
        }
    }

    //
    // Loop over all tree entries
    //
    const Int_t n = fTrees.GetEntriesFast();

    for (int idx=0; idx<n; idx++)
    {
        TTree *t = (TTree*)fTrees[idx];

        //
        // Check the write flag of the tree
        //
        if (!t->TestBit(kFillTree))
            continue;

        //
        // If the write flag is set, fill the tree (with the corresponding
        // branches/containers), delete the write flag and increase the number
        // of written/filled entries.
        //
        t->ResetBit(kFillTree);

        if (!t->Fill())
        {
            *fLog << err << "ERROR - Zero bytes written to tree '" << t->GetName() << "'... abort." << endl;
            return kFALSE;
        }
    }

    //
    // For more information see TTree:ChangeFile()
    //
    TTree *t0 = (TTree*)fTrees[0];
    if (!t0 || fOut==t0->GetCurrentFile())
        return kTRUE;

    *fLog << warn << endl;
    *fLog << "WARNING - MWriteRootFile:   Root's  TTree/TFile   has  opened   a  new  file" << endl;
    *fLog << "  automatically.  You can change this behaviour using TTree::SetMaxTreeSize." << endl;
    *fLog << "  You won't be able to read splitted  files  correctly with MReadMarsFile if" << endl;
    *fLog << "  they have more than one entry in 'RunHeaders' or you try to read more than" << endl;
    *fLog << "  one of such sequences at once." << endl;
    *fLog << endl;

    return kTRUE;
}

// --------------------------------------------------------------------------
//
// Open a new file with the ame fname. Move all trees and branches from the
// old file to the new file.
//
Bool_t MWriteRootFile::ChangeFile(const char *fname)
{
    //
    // The following code is more or less a copy of TTree::ChangeFile
    //
    const Int_t   compr = fOut->GetCompressionLevel();
    const TString title = fOut->GetTitle();

    *fLog << inf << "MWriteRootFile - Open new file " << fname << " (Title=" << title << ", Compression=" << compr << ")" << endl;

    // Open new file with old setup
    TFile *newfile = TFile::Open(fname, "RECREATE", title, compr);
    if (!newfile)
    {
        *fLog << err << "ERROR - Cannot open new file " << fname << endl;
        return kFALSE;
    }

    // Print statistics of old file
    const TString n = GetFileName();
    if (!n.IsNull() && n!=TString("/dev/null"))
        Print();

    if (fOut->IsOpen())
        fOut->Write();

    // Move all trees from the old file to the new file
    TObject *obj=0;
    while ((obj = fOut->GetList()->First()))
    {
        // Remove obj from old file (otherwise deleting
        // the old file will delete the objs)
        fOut->GetList()->Remove(obj);

        // If this is not a tree do nothing.
        if (!obj->InheritsFrom("TTree"))
            continue;

        // process all trees in the old file
        TTree *t = (TTree*)obj;

        // reset and move to new file (this is done implicitly for all branches)
        t->Reset();
        t->SetDirectory(newfile);
    }

    // Close/delete the old file (keys already written above)
    delete fOut;

    // Replace current with new file
    fOut = newfile;

    // Change current directory to new file
    gFile = fOut;

    return kTRUE;
}

// --------------------------------------------------------------------------
//
// A rule looks like:
//   "outputpath{s/source/destination}"
//
// outpath:                the output path into which the files are written
// {s/source/destination}  a substitution rule for the filename
//   while source can be a regular expression everything which matches this
//   regular expression (see TRegexp) will be replaced by destination.
//   Warning: The algorithm is recursive you may create endless loops!
//
// Example:
//   inputfile: /data/MAGIC/Period016/rootdata/20040621_23210_D_Mkn421_E.root
//   rule:      /outpath/{s/_D_/_Y_}
//   outfile:   /outpath/20040621_23210_Y_Mkn421_E.root
//
// If you need more difficult rules please send me an eMail...
//
TString MWriteRootFile::GetNewFileName(const char *inname) const
{
    // Remove the path from the filename
    TString fname(inname);
    if (fname.Last('/')>=0)
        fname.Remove(0, fname.Last('/')+1);

    // Make a copy of the rule
    TString rule(fSplitRule);

    // [characte class], ^ do not, [^{}] do not match { and }, [^{}]+ match at least not one { or }
    const TRegexp subst0("{s/[^{}/]+/[^{}/]+}");

    TString path;
    Bool_t first = kTRUE;

    Ssiz_t idx=0;
    while (1)
    {
        // Find a substitution exprsssion
        Ssiz_t len = 0;
        idx = subst0.Index(rule, &len);
        if (idx<0)
            break;

        // If the first substitution expression is found in the rule
        // determin the path from everything before
        if (first)
        {
            path=rule(0, idx);
            first = kFALSE;
        }

        // Extract a substitution expression
        TString expr = rule(idx, len);
        rule.Remove(idx, len);

        expr.Remove(0,3);
        expr.Remove(expr.Length()-1);

        // In all cases this is well defined (see Regexp)
        const Ssiz_t pos = expr.First('/');

        // Split substitution rule into source and destination
        const TString src  = expr(0, pos);
        const TString dest = expr(pos+1, expr.Length());

        // Replace source by destination
        const TRegexp regexp(src);
        while (1)
        {
            TString sub = fname(idx+dest.Length());
            idx = regexp.Index(fname, &len);
            if (idx<0)
                break;

            fname.Replace(idx, len, dest);
        }
    }

    // Check if we have a trailing '/'
    if (!path.IsNull() && path[path.Length()-1]!='/')
        path.Append("/");

    // Create full qualified pathname
    path += fname;
    return path;
}

// --------------------------------------------------------------------------
//
// ReInit. If file splitting is not allowed call MWriteFile::ReInit.
//
// In other cases get MRead from the TaskList (splitting is switched of if
// this is impossible).
//
// Convert the input- into a new output file-name.
//
// Open a new file, change all trees to the new file (calling ChangeFile()),
// and close the old one.
//
// Call MWriteFile::ReInit()
//
Bool_t MWriteRootFile::ReInit(MParList *pList)
{
    if (fSplitRule.IsNull())
        return MWriteFile::ReInit(pList);

    MRead *read = (MRead*)pList->FindTask("MRead");
    if (!read)
    {
        *fLog << warn;
        *fLog << "WARNING: No Task 'MRead' found in the tasklist.  This task is" << endl;
        *fLog << "  necessary  to  get  the filename.  Without  a filename file" << endl;
        *fLog << "  file splitting is not allowed... disabled!" << endl;
        *fLog << endl;
        fSplitRule = "";
        return kTRUE;
    }

    const TString fname = GetNewFileName(read->GetFileName());
    if (!ChangeFile(fname))
        return kFALSE;

    return MWriteFile::ReInit(pList);
}

// --------------------------------------------------------------------------
//
// return open state of the root file.
//
Bool_t MWriteRootFile::IsFileOpen() const
{
    const char *n = fOut->GetName();
    return n==0 || *n==0 ? kTRUE : fOut->IsOpen();
}

// --------------------------------------------------------------------------
//
// return name of the root-file
//
const char *MWriteRootFile::GetFileName() const
{
    const char *n = fOut->GetName();
    return n==0 || *n==0 ? "<dummy>" : n;
}

// --------------------------------------------------------------------------
//
// cd into file. See TFile::cd()
//
Bool_t MWriteRootFile::cd(const char *path)
{
    return fOut->cd(path);
}

// --------------------------------------------------------------------------
//
// Implementation of SavePrimitive. Used to write the call to a constructor
// to a macro. In the original root implementation it is used to write
// gui elements to a macro-file.
//
void MWriteRootFile::StreamPrimitive(ofstream &out) const
{
    out << "   MWriteRootFile " << GetUniqueName() << "(\"";
    out << fOut->GetName() << "\", \"";
    out << fOut->GetOption() << "\", \"";
    out << fOut->GetTitle() << "\", ";
    out << fOut->GetCompressionLevel();

    if (fName!=gsDefName || fTitle!=gsDefTitle)
    {
        out << ", \"" << fName << "\"";
        if (fTitle!=gsDefTitle)
            out << ", \"" << fTitle << "\"";
    }
    out << ");" << endl;


    MRootFileBranch *entry;
    TIter Next(&fBranches);
    while ((entry=(MRootFileBranch*)Next()))
    {
        out << "   " << GetUniqueName() << ".AddContainer(";

        if  (entry->GetContainer())
        {
            entry->GetContainer()->SavePrimitive(out);
            out << "&" << entry->GetContainer()->GetUniqueName();
        }
        else
            out << "\"" << entry->GetContName() << "\"";

        out << ", \"" << entry->GetName() << "\"";
        if (!entry->MustHave())
            out << ", kFALSE";

        out <<");" << endl;
    }
}