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

Description:
    Add a branch called branchname and filled with branchvalue to a tree identified by (filename,treename) and save the output in a new file.

How to run it:
    python AddNewBranch.py --filename [filename] --treename [treename] --branchname [branchname] --branchvalue [branchvalue]
"""

import ROOT
import argparse
from array import array

def AddNewBranch(filename,treename,branchname,branchvalue) :
    print 'Adding a new branch...'
    print 'Input file: ', filename
    print 'Branch name: ', branchname
    print 'Branch value: ', branchvalue

    fIn = ROOT.TFile(filename,'READ')

    # Check that the tree exists.
    if not fIn.GetListOfKeys().Contains(treename) and not fIn.GetListOfKeys().Contains(treename.rpartition('/')[2]) :
        print 'ERROR! Object not found in the file.'
        return

    tIn = fIn.Get(treename)
    
    fOut = ROOT.TFile(filename.replace('.root','')+'_With'+branchname+'Label'+'.root','RECREATE')
    
    fOut.cd()
    
    tOut = tIn.CloneTree()

    new = array('i',[0])
    # new = int(branchvalue)
    new[0] = int(branchvalue)
    print new
    bNew = tOut.Branch(branchname,new,branchname+'/I')

    for iEvent in range(tOut.GetEntries()) :
        tOut.GetEntry(iEvent)
        bNew.Fill()

    # Write output.
    fOut.Write()

    # Close files.
    fIn.Close();
    fOut.Close()

    return

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

if __name__ == "__main__" :
    parser = argparse.ArgumentParser(description='Add a branch to a tree.')

    parser.add_argument('--filename',required=True,help='filename')
    parser.add_argument('--treename',required=True,help='treename')
    parser.add_argument('--branchname',required=True,help='branchname')
    parser.add_argument('--branchvalue',required=True,help='branchvalue')

    args = parser.parse_args()

    # Parameters and configuration.
    filename = args.filename
    treename = args.treename
    branchname = args.branchname
    branchvalue = args.branchvalue

    AddNewBranch(filename,treename,branchname,branchvalue)