import sys
import os
import ROOT as r

source = str(sys.argv[1])
a = str(sys.argv[2])
sel_names = [c for c in a.split(",")]

source_path = '/disk/data3/lhcb/elena/B2KShh/ntuples/jobs/'
target_path = '/disk/data3/lhcb/elena/B2KShh/ntuples/sel/'
sel_strings = {
	'L0'   : '(L0HadronDecision_TOS || L0Global_TIS)',
	'HLT1' : 'Hlt1TrackAllL0Decision_TOS',
}


def getall(d, basepath="/"):
    "Generator function to recurse into a ROOT file/dir and yield (path, obj) pairs"
    keylist = d.GetListOfKeys()
    for key in keylist:
        kname = key.GetName()
        if key.IsFolder() and d.Get(kname).ClassName() == 'TDirectoryFile' :
            for i in getall(d.Get(kname), basepath+kname+"/"):
                yield i
        else:
            yield basepath+kname, d.Get(kname)


infile = r.TFile(source_path + source, 'read')
trees = []

for k, o in getall(infile):
	if 'TTree' in o.ClassName():
		trees.append(o)

recotrees = [t for t in trees if 'MCTruth' not in t.GetName()]
truthtrees = [t for t in trees if 'MCTruth' in t.GetName()]
if len(truthtrees) == 1:
	truth_tree = truthtrees[0]
else:
	truth_tree = None

os.system('mkdir -p %s' % target_path)
if not os.path.isfile(target_path + source):
	print "\tSource file not in target directory, I'll bring it along..."
	os.system('cp %s %s' % (source_path + source, target_path + source))

outfilename = source.split('.')[0] + '_' + '_'.join(sel_names) + '.root'
print '\tCreating file %s ...' % outfilename
outfile = r.TFile(target_path + outfilename, 'recreate')
clones = []
if truth_tree:
	print '\tCopying MCTruth tree...'
	#truth_tree.AutoSave()
	clones.append(truth_tree.CopyTree(''))
for t in recotrees:
	print '\tCopying %s tree...' % t.GetName()
	clones.append(t.CopyTree(' && '.join([sel_strings[sel_name] for sel_name in sel_names])))
for c in clones:
	c.AutoSave()
outfile.Close()
infile.Close()
print '\tAll done.'
print