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

/* ======================================================================== *\
!
! *
! * 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-2003
!
!
\* ======================================================================== */

/////////////////////////////////////////////////////////////////////////////
//                                                                         //
// 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: Checkout the Warning in MTaskList.                             //
//                                                                         //
/////////////////////////////////////////////////////////////////////////////
#include "MWriteRootFile.h"

#include <fstream>

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

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

#include "MParList.h"

ClassImp(MRootFileBranch);
ClassImp(MWriteRootFile);

using namespace std;

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

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

    fBranches.SetOwner();
}

// --------------------------------------------------------------------------
//
// 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)
{
    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();

    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.
//
MWriteRootFile::~MWriteRootFile()
{
    //
    // 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;

    //
    // 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;
}

// --------------------------------------------------------------------------
//
// 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" << branch->GetEntries() << " entries." << endl;
    }

    TTree *t = NULL;
    TIter NextTree(&fTrees);
    while ((t=(TTree*)NextTree()))
        if (t->TestBit(kIsNewTree))
            *fLog << " " << t->GetName() << ": \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)
{
    //
    // create a new entry in the list of branches to write and
    // add the entry to the list.
    //
    MRootFileBranch *entry = new MRootFileBranch(cname, tname, must);
    fBranches.AddLast(entry);

    if (tname && tname[0])
        AddToBranchList(Form("%s.%s", 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)
{
    //
    // 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());
        }
        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;
        }

        //
        // 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() const
{
    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
    //
    TIter NextTree(&fTrees);
    while ((obj=NextTree()))
    {
        TTree *t = (TTree*)obj;

        //
        // 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;
        }
    }
    return kTRUE;
}

// --------------------------------------------------------------------------
//
// return open state of the root file.
//
Bool_t MWriteRootFile::IsFileOpen() const
{
    return fOut->IsOpen();
}

// --------------------------------------------------------------------------
//
// return name of the root-file
//
const char *MWriteRootFile::GetFileName() const
{
    return fOut->GetName();
}

// --------------------------------------------------------------------------
//
// 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;
    }
}