from __future__ import division
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/'

end2011 = 106014  # fill 2354
endjune = 119892  # fill 2772

sel_strings = {
	'L0'           : '(L0HadronDecision_TOS || L0Global_TIS)',
	'HLT1'         : 'Hlt1TrackAllL0Decision_TOS',
	'2011'         : '(runNumber <= %s)' % end2011,
	'HLT2-2011'    : '((Hlt2Topo2BodyBBDTDecision_TOS || Hlt2Topo3BodyBBDTDecision_TOS || Hlt2Topo4BodyBBDTDecision_TOS) || (Hlt2Topo2BodySimpleDecision_TOS || Hlt2Topo3BodySimpleDecision_TOS || Hlt2Topo4BodySimpleDecision_TOS))',
	'2012pre-june' : '(runNumber > %s && runNumber <= %s)' % (end2011, endjune),
	'2012post-june': '(runNumber > %s)' % endjune,
	'HLT2-2012'    : '(Hlt2Topo2BodyBBDTDecision_TOS || Hlt2Topo3BodyBBDTDecision_TOS || Hlt2Topo4BodyBBDTDecision_TOS)',
}

sel_strings['Trigger'] = ' && '.join([sel_strings['L0'], sel_strings['HLT1'], sel_strings['HLT2-2011']])
sel_strings['Trigger'] += ' && '
if 'MC_2011' in source:
	sel_strings['Trigger'] += sel_strings['HLT2-2011']
elif 'MC_2012' in source:
	sel_strings['Trigger'] += sel_strings['HLT2-2012']
elif not 'MC' in source:
	pass
else:
	raise Exception('DANGER: problem with file name parsing and trigger conditions!')


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)
	t.cptree = t.tree.CopyTree(' && '.join([sel_strings[sel_name] for sel_name in sel_names]))
	n_new, n_old = t.cptree.GetEntriesFast(), t.tree.GetEntriesFast()
	print '\t Copying {0}/{1} tree: {2:.1%} out of {3} events selected...'.format(t.location, t.tree.GetName(), n_new / n_old, n_old)
	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