"""
Author: Federica Lionett
Date: March 2nd, 2016

Description:
    Merge three trees identified by (filename1,treename1), (filename2,treename2), and (filename3,treename3) and save the output in (filename,treename).

How to run it:
    python MergeTrees.py --filename1 [filename1] --treename1 [treename1] --filename2 [filename2] --treename2 [treename2] --filename3 [filename3] --treename3 [treename3] --filename [filename] --treename [treename]
"""

import ROOT
import argparse

def Merge3Trees(filename1,treename1,filename2,treename2,filename3,treename3,filename,treename) :
    fIn1 = ROOT.TFile(filename1,'READ')
    fIn2 = ROOT.TFile(filename2,'READ')
    fIn3 = ROOT.TFile(filename3,'READ')
    
    # Check that the tree exists.
    if not fIn1.GetListOfKeys().Contains(treename1) :
        print 'ERROR! Object not found in the first file.'
        return
    if not fIn2.GetListOfKeys().Contains(treename2) :
        print 'ERROR! Object not found in the second file.'
        return
    if not fIn3.GetListOfKeys().Contains(treename3) :
        print 'ERROR! Object not found in the third file.'
        return
    
    tIn1 = fIn1.Get(treename1)
    tIn2 = fIn2.Get(treename2)
    tIn3 = fIn3.Get(treename3)
    tList = ROOT.TList()
    tList.Add(tIn1)
    tList.Add(tIn2)
    tList.Add(tIn3)

    print 'Merging the three trees...'
    print 'First file: ', filename1
    print 'First tree: ', treename1
    print 'Second file: ', filename2
    print 'Second tree: ', treename2
    print 'Third file: ', filename3
    print 'Third tree: ', treename3
    
    fOut = ROOT.TFile(filename,'RECREATE')
    
    fOut.cd()
    
    tOut = ROOT.TTree.MergeTrees(tList)
    tOut.SetName(treename)

    # Write output.
    fOut.Write()

    # Close files.
    fIn1.Close();
    fIn2.Close();
    fIn3.Close();
    fOut.Close()

    return

###############
# 
# Main function
#

if __name__ == "__main__" :
    parser = argparse.ArgumentParser(description='Merge three trees.')

    parser.add_argument('--filename1',required=True,help='filename1')
    parser.add_argument('--treename1',required=True,help='treename1')
    parser.add_argument('--filename2',required=True,help='filename2')
    parser.add_argument('--treename2',required=True,help='treename2')
    parser.add_argument('--filename3',required=True,help='filename3')
    parser.add_argument('--treename3',required=True,help='treename3')
    parser.add_argument('--filename',required=True,help='filename')
    parser.add_argument('--treename',required=True,help='treename')

    args = parser.parse_args()

    # Parameters and configuration.
    filename1 = args.filename1
    treename1 = args.treename1
    filename2 = args.filename2
    treename2 = args.treename2
    filename3 = args.filename3
    treename3 = args.treename3
    filename = args.filename
    treename = args.treename

    MergeTrees(filename1,treename1,filename2,treename2,filename3,treename3,filename,treename)