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

truth_match = {
	'pi'       : 211,
	'pi0'      : 111,
	'K'        : 321,
	'K0'       : 311,
	'KS'       : 310,
	'p'        : 2212,
	'Lc'       : 4122,
	'Sigma_c0' : 4112,
	'Xib0'     : 5232,
	'Lb'       : 5122,
}

hypos = {
	'B2KpKS'   : {'h1': 'K',  'h2': 'p',  'KS': 'KS'},
	'B2pKKS'   : {'h1': 'p',  'h2': 'K',  'KS': 'KS'},
	'B2KpiKS'  : {'h1': 'K',  'h2': 'pi', 'KS': 'KS'},
	'B2piKKS'  : {'h1': 'pi', 'h2': 'K',  'KS': 'KS'},
	'B2ppiKS'  : {'h1': 'p',  'h2': 'pi', 'KS': 'KS'},
	'B2pipKS'  : {'h1': 'pi', 'h2': 'p',  'KS': 'KS'},
	'B2pipiKS' : {'h1': 'pi', 'h2': 'pi', 'KS': 'KS'},
	'B2ppKS'   : {'h1': 'p',  'h2': 'p',  'KS': 'KS'},
	'B2KKKS'   : {'h1': 'K',  'h2': 'K',  'KS': 'KS'},
}


def truth_match_string(source, hypo):
	string  = '('
	mother = 'MOTHER'
	b_id = truth_match['Lb']
	if 'Lc' in source:
		mother = 'GDMOTHER'
	if 'Xib0' in source:
		b_id = truth_match['Xib0']
	for h in ['h1', 'h2', 'KS']:
		string += ' && %s_TRUEID == %s' % (h, truth_match[hypos[hypo][h]])
		string += ' && %s_%s_TRUEID == %s' % (h, mother, b_id)
	string += ')'
	return string


def make_sel_trings(source, recotree):
	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!')
	
	sel_strings['TruthMatch'] = truth_match_string(source, recotree)
	return sel_strings


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)


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


#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))


def find_input(target_path, source, sel_names):
	"""
	If some selections were already performed,
	use the selected sample as input!
	"""
	infile_name = target_path + source
	available_files = os.listdir(target_path)
	for i in range(len(sel_names)):
		sel_subset = sel_names[:i+1]
		for f in available_files:
			if f == source.split('.')[0] + '_' + '_'.join(sel_subset) + '.root':
				infile_name = target_path + f
	return infile_name


infile = r.TFile(find_input(target_path, source, sel_names), 'read')
trees = []

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()]

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)
	sel_strings = make_sel_strings(source, t.GetName())
	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