"""
Author: Federica Lionetto
Date: April 19th, 2016

Description:
    Create a new tree, consisting in the original tree, plus a branch from another tree.
    
How to run it:
    python MergeBranchesWithTree.py --filename1 [filename1] --treename1 [treename1] --filename2 [filename2] --treename2 [treename2] --branches [list of (branches,type) pairs]
"""

import ROOT
import argparse
from array import array
import numpy as np

def MergeBranchesWithTree(filename1,treename1,filename2,treename2,branches,branchTypes) :
    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)
   
    # Check that the branches exist.
    for branch in branches :
        if not tIn2.GetListOfBranches().Contains(branch) :
            print 'ERROR! Branch not found.'
            return

    filename = filename1.replace('.root','_Extended.root')
    treename = treename1

    print 'Creating a new tree...'
    print 'File: ', filename
    print 'Tree: ', treename
    print 'Branches: ', branches
    
    
    fOut = ROOT.TFile(filename,'RECREATE')
    tOut = tIn1.CloneTree()

    nEntries1 = tIn1.GetEntries()
    nEntries2 = tIn2.GetEntries()
    if not nEntries1 == nEntries2 :
        print 'ERROR! Different number of entries in the two trees.'
        return

    nEntries = nEntries2
    print 'Number of entries: ', nEntries

    branchInList = []
    branchOutList = []

    for iBranch, branch in enumerate(branches) :
        if not tIn1.GetListOfBranches().Contains(branch) :
            print 'Branch', branch
            branchInList.append(array(branchTypes[iBranch],[0]))
            tIn2.SetBranchAddress(branch,branchInList[-1])

            branchOutList.append(tOut.Branch(branch,branchInList[-1],branch+'/'+branchTypes[iBranch].upper()))

    for iEvent in range(nEntries) :
        tIn2.GetEntry(iEvent)
        for iBranch in range(len(branchInList)) :
                if iEvent < 5 :
                    print branchInList[iBranch][0]
                    print branchInList[iBranch]
                branchOutList[iBranch].Fill()

    # Write output.
    fOut.cd()
    fOut.Write()
            
    # Close files.
    fIn1.Close()
    fIn2.Close()
    fOut.Close()

    return

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

if __name__ == "__main__" :
    parser = argparse.ArgumentParser(description='Create a new tree, consisting in the original tree, plus some branches from another tree.')

    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('--branches',required=True,nargs='+',help='branches') # Each branch should be separated from the others by a space.
    parser.add_argument('--branchTypes',required=True,nargs='+',help='branch types') # Each branch type should be separated from the others by a space.

    args = parser.parse_args()

    # Parameters and configuration.
    filename1 = args.filename1
    treename1 = args.treename1
    filename2 = args.filename2
    treename2 = args.treename2
    branches = args.branches
    branchTypes = args.branchTypes

    MergeBranchesWithTree(filename1,treename1,filename2,treename2,branches,branchTypes)