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 = []


class treeWithDirectory(object):
	def __init__(self, key, obj):
		self.tree = obj
		self.location = key.split('/')[1]


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

recotrees = [t for t in trees if 'MCTruth' not in t.tree.GetName()]
#truthtrees = [t for t in trees if 'MCTruth' in t.tree.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.tree.CopyTree(''))
#	truth_tree.cptree = truth_tree.tree.CopyTree('')
for t in recotrees:
	if not outfile.GetDirectory(t.location):
		print '\tMaking directory %s ...' % t.location
		outfile.mkdir(t.location, t.location)
	outfile.cd(t.location)
	print '\tCopying %s/%s tree: ...' % (t.location, t.tree.GetName())
	t.cptree = t.tree.CopyTree(' && '.join([sel_strings[sel_name] for sel_name in sel_names]))
	t.cptree.SetDirectory(outfile.GetDirectory(t.location))
	t.cptree.AutoSave()
	#clones.append(t.tree.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