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

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

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

import ROOT
import argparse

def MergeTrees(filename1,treename1,filename2,treename2,filename,treename) :
    fIn1 = ROOT.TFile(filename1,'READ')
    fIn2 = ROOT.TFile(filename2,'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
    
    tIn1 = fIn1.Get(treename1)
    tIn2 = fIn2.Get(treename2)
    tList = ROOT.TList()
    tList.Add(tIn1)
    tList.Add(tIn2)

    print 'Merging the two trees...'
    print 'First file: ', filename1
    print 'First tree: ', treename1
    print 'Second file: ', filename2
    print 'Second tree: ', treename2
    
    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();
    fOut.Close()

    return

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

if __name__ == "__main__" :
    parser = argparse.ArgumentParser(description='Merge two 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('--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
    filename = args.filename
    treename = args.treename

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