"""
Author: Federica Lionett
Date: February 5th, 2016

Description:
    Create a reduced tree containing only a subset of the branches of the original tree.

How to run it:
    python CreateReducedTree.py --filename [filename] --treename [treename] --branches [list of (branches,type) pairs]
"""

import ROOT
import argparse

def CreateReducedTree(filename,treename,branches) :
    fIn = ROOT.TFile(filename,'READ')
    
    # Check that the tree exists.
    if not fIn.GetListOfKeys().Contains(treename) :
        print 'ERROR! Object not found.'
        return
    
    tIn = fIn.Get(treename)
    
    # Check that the branches exist.
    for branch in branches :
        if not tIn.GetListOfBranches().Contains(branch) :
            print 'ERROR! Branch not found.'
            return

    print 'Creating a reduced tree...'
    print 'File: ', filename
    print 'Tree: ', treename
    print 'Branches: ', branches
    
    tIn.SetBranchStatus('*',0)
    for branch in branches :
        tIn.SetBranchStatus(branch,1)

    fOut = ROOT.TFile(filename.replace('.root','_ReducedTree.root'),'RECREATE')
    tOut = tIn.CloneTree()

    # Write output.
    fOut.cd()
    fOut.Write()

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

    return

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

if __name__ == "__main__" :
    parser = argparse.ArgumentParser(description='Create a reduced tree.')

    parser.add_argument('--filename',required=True,help='filename')
    parser.add_argument('--treename',required=True,help='treename')
    parser.add_argument('--branches',required=True,nargs='+',help='branches') # Each branch should be separated from the others by a space.

    args = parser.parse_args()

    # Parameters and configuration.
    filename = args.filename
    treename = args.treename
    branches = args.branches

    CreateReducedTree(filename,treename,branches)