#
#  Get LFN files from BKKquery
#  Generate xml files from BKKquery
#
#
max_LFNs = 30 #define how many lfns are included in the xml file

channel_list = ['Kplus',
		'Kshort'
		]

decay_list = [	'BtoKstMuMu',
 		#'BtoKstJpsi',
		#'PHSP'
		#'BtoXJpsi'
		#'B0toKstMuMu',
		#'B0toKstJpsi',
		#'BtoK1MuMu',
		#'BtoK1Jpsi'
			]

polarity_list = ['MagDown',
		 'MagUp']

decay_no_dict = {
		'BtoKstMuMu' : {'Kplus':  12113100,
				'Kshort': 12115102},
		'BtoKstJpsi' : {'Kplus':  12143401,
				'Kshort': 12145102},
		'PHSP'	     : {'Kplus':  12113445,
				'Kshort': 12115178},
		'BtoXJpsi'   : {'Kplus':  12442001,
				'Kshort': 12442001},
		'B0toKstMuMu': {'Kplus':  11114001,
				'Kshort': 11114001},
		'B0toKstJpsi': {'Kplus':  11144001,
				'Kshort': 11144001},
		#'BtoK1MuMu'  : {'Kplus':  12435001, # it's B+ to JpsiKst everything
		#		'Kshort': 12435001},
		'BtoK1Jpsi'  : {'Kplus':  12245000,
	                        'Kshort': 12245000}
		}


decay_year_dict = { 	'BtoKstMuMu' : ['2017'],#['2011','2012','2015','2016','2017'],
			'BtoKstJpsi' : ['2011','2012'],
			'PHSP'	     : ['2012'],
#			'PHSP'	     : ['2011','2012','2015','2016'],
			'BtoXJpsi'   : ['2011','2012','2016'],
			'B0toKstMuMu': ['2011','2012','2015','2016'],
			'B0toKstJpsi': ['2011','2012','2015','2016'],
			'BtoK1MuMu'  : ['2011','2012'],
			'BtoK1Jpsi'  : ['2012']
		}

year_E_dict = { '2011': '3500GeV',
		'2012': '4000GeV',
		'2015': '6500GeV',
		'2016': '6500GeV',
		'2017': '6500GeV'
		}

year_nu_dict = { '2011': '2',
		 '2012': '2.5',
		 '2015': '1.6-25ns',
		 '2016': '1.6-25ns',
		 '2017': '1.6-25ns'
		}

year_trigg_dict = { 	'2011': 'Trig0x40760037',
			#'2012': 'Trig0x409f0045',	#TODO! two different trigger settings! for PHSP
			'2012': 'Trig0x4097003d',  #Seccond trigger settings for PHSP
			'2015': 'Trig0x411400a2',
			'2016': 'Trig0x6138160F',
			'2017': 'Trig0x62661709'
		}

decay_simVer_dict = { 	'BtoKstMuMu' : {'2011': 'Sim09b',
					'2012': 'Sim09a',
					'2015': 'Sim09b',
					'2016': 'Sim09c',
					'2017': 'Sim09e'
					},
			'BtoKstJpsi' : {'2011': 'Sim09a',
					'2012': 'Sim09a'
					},
			'PHSP'	     : {'2011': 'Sim09b',
					#'2012': 'Sim09b',
					'2012': 'Sim09c',
					'2015': 'Sim09b',
					'2016': 'Sim09c'
					},
	                'BtoXJpsi'   : {'2011': 'Sim08c/Digi13',
					'2012': 'Sim08a/Digi13',
					'2016': 'Sim09b'
					},
	                'B0toKstMuMu': {'2011': 'Sim08e/Digi13',
					'2012': 'Sim08b/Digi13',
					'2015': 'Sim09c',
					'2016': 'Sim09b'
					},
	                'B0toKstJpsi': {'2011': 'Sim08f/Digi13',
					'2012': 'Sim08f/Digi13',
					'2015': 'Sim09c',
					'2016': 'Sim09c'
					},
	                'BtoK1MuMu'  : {'2011': 'Sim08e/Digi13',
					'2012': 'Sim08e/Digi13',
					},
	                'BtoK1Jpsi'  : {'2012': 'Sim08a/Digi13'
					}
		}

decay_strip_dict = { 	'BtoKstMuMu' : {'2011': 'Reco14c/Stripping21r1NoPrescalingFlagged',
					'2012': 'Reco14c/Stripping21NoPrescalingFlagged',
					'2015': 'Reco15a/Turbo02/Stripping24NoPrescalingFlagged',
					'2016': 'Reco16/Turbo03/Stripping28NoPrescalingFlagged',
					'2017': 'Reco17/Turbo04a-WithTurcal/Stripping29r2NoPrescalingFlagged',
					},
			'BtoKstJpsi' : {'2011': 'Reco14c/Stripping21r1NoPrescalingFlagged',
					'2012': 'Reco14c/Stripping21NoPrescalingFlagged',
					},
			'PHSP'	     : {'2011': 'Reco14c/Stripping21r1Filtered',
					#'2012': 'Reco14c/Stripping21Filtered',
					'2012': 'Reco14c/Stripping21NoPrescalingFlagged', #Seccond trigger settings for PHSP
					'2015': 'Reco15a/Turbo02/Stripping24Filtered',
					'2016': 'Reco16/Turbo03/Stripping28Filtered'
					},
	                'BtoXJpsi'   : {'2011': 'Reco14a/Stripping20r1NoPrescalingFlagged',
					'2012': 'Reco14a/Stripping20NoPrescalingFlagged',
					'2016': 'Reco16/Turbo03/Stripping26NoPrescalingFlagged'
					},
	                'B0toKstMuMu': {'2011': 'Reco14a/Stripping20r1NoPrescalingFlagged',
					'2012': 'Reco14a/Stripping20NoPrescalingFlagged',
					'2015': 'Reco15a/Turbo02/Stripping24r1NoPrescalingFlagged',
					'2016': 'Reco16/Turbo03/Stripping26NoPrescalingFlagged'
					},
	                'B0toKstJpsi': {'2011': 'Reco14a/Stripping20r1NoPrescalingFlagged',
					'2012': 'Reco14a/Stripping20NoPrescalingFlagged',
					'2015': 'Reco15a/Turbo02/Stripping24r1NoPrescalingFlagged', #todo
					'2016': 'Reco16/Turbo03/Stripping28r1NoPrescalingFlagged'
					},
	                'BtoK1MuMu'  : {'2011': 'Reco14a/Stripping20r1NoPrescalingFlagged',
					'2012': 'Reco14a/Stripping20NoPrescalingFlagged',
					},
	                'BtoK1Jpsi'  : {'2012': 'Reco14a/Stripping20NoPrescalingFlagged'
					}
		}

decay_file_dict = { 	'BtoKstMuMu' : {'2011': 'ALLSTREAMS.DST',
					'2012': 'ALLSTREAMS.DST',
					'2015': 'ALLSTREAMS.MDST',
					'2016': 'ALLSTREAMS.MDST',
					'2017': 'ALLSTREAMS.MDST'
					},
			'BtoKstJpsi' : {'2011': 'ALLSTREAMS.DST',
					'2012': 'ALLSTREAMS.DST'
					},
			'PHSP'	     : {'2011': 'B2XMUMU.STRIP.DST',
					'2012': 'B2XMUMU.STRIP.DST',
					'2015': 'B2XMUMU.STRIP.MDST',
					'2016': 'B2XMUMU.STRIP.MDST'
					},
	                'BtoXJpsi'   : {'2011': 'ALLSTREAMS.DST',
					'2012': 'ALLSTREAMS.DST',
					'2016': 'ALLSTREAMS.DST'
					},
	                'B0toKstMuMu': {'2011': 'ALLSTREAMS.DST',
					'2012': 'ALLSTREAMS.DST',
					'2015': 'ALLSTREAMS.DST',
					'2016': 'ALLSTREAMS.DST'
					},
	                'B0toKstJpsi': {'2011': 'ALLSTREAMS.DST',
					'2012': 'ALLSTREAMS.DST',
					'2015': 'ALLSTREAMS.MDST',
					'2016': 'ALLSTREAMS.MDST'
					},
	                'BtoK1MuMu'  : {'2011': 'ALLSTREAMS.DST',
					'2012': 'ALLSTREAMS.DST',
					},
	                'BtoK1Jpsi'  : {'2012': 'ALLSTREAMS.DST'
					}
		}

path_list_dict = {}
path_list = []

for channel in channel_list:
	for decay in decay_list:
	    for polarity in  polarity_list:
        	for year in decay_year_dict[decay]:
			script_opts = { "DATA_TYPE":"MC",
                                "CHAN": channel,
				"YEAR":year,
				"POL":polarity,
				"ENERGY":year_E_dict[year],
				"NU":year_nu_dict[year],
				"SIM": decay_simVer_dict[decay][year],
				"TRIG":year_trigg_dict[year],
				"STRIP":decay_strip_dict[decay][year],
				"NUM":decay_no_dict[decay][channel],
				"FILE":decay_file_dict[decay][year],
				"DEC":decay
					}
			PATH = "/%(DATA_TYPE)s/%(YEAR)s/Beam%(ENERGY)s-%(YEAR)s-%(POL)s-Nu%(NU)s-Pythia8/%(SIM)s/%(TRIG)s/%(STRIP)s/%(NUM)s/%(FILE)s" % script_opts
			path_list += [PATH]
			outFile = "%(DEC)s_%(YEAR)s%(POL)s_%(CHAN)s" % script_opts
			path_list_dict.update({PATH: outFile})

get_xml_files = False #God knows why the second popen.subprocess doesn't work

if ('BtoXJpsi' in decay_list):
	PATH = "/MC/Dev/Beam6500GeV-Jun2015-MagDown-Nu1.6-25ns-Pythia8/Sim09Dev03/Trig0x410700a1/Reco15/Turbo01a/Stripping23r1NoPrescalingFlagged/12442001/ALLSTREAMS.MDST"
	path_list += [PATH]
	path_list_dict.update({PATH: 'BtoXJpsi_2015MagDown.py'})

	PATH = "/MC/Dev/Beam6500GeV-Jun2015-MagUp-Nu1.6-25ns-Pythia8/Sim09Dev03/Trig0x410700a1/Reco15/Turbo01a/Stripping23r1NoPrescalingFlagged/12442001/ALLSTREAMS.MDST"
	path_list += [PATH]
	path_list_dict.update({PATH: 'BtoXJpsi_2015MagUp.py'})

import subprocess
from subprocess import Popen, PIPE

if __name__ == '__main__':
    from subprocess import call
    from sys import argv
    '''  	
    if (not get_xml_files):
    	for f in path_list:
		address_file = open(path_list_dict[f]+'.py','w')
		address_file.write("""from Gaudi.Configuration import *\n""")
		address_file.write("""from GaudiConf import IOHelper\n""")
		address_file.write("""IOHelper('ROOT').inputFiles([\n""")
		print('Getting file {0}.'.format(f))
		dirac_command = subprocess.Popen('lb-run LHCbDIRAC dirac-bookkeeping-get-files --BKQuery {0}'.format(f), shell=True, stdout=subprocess.PIPE)
		stdout = dirac_command.communicate()[0]
		stdout = stdout.split('\n')
		#Loop only first 10 files (or less, if there are less LFNs)
		#First 3 entries are junk from LHCbDIRAC dirac-bookkeeping-get-files
		max_file = min(max_LFNs+3,len(stdout)-2)
		for i in range(3, max_file-1):
			stdout_tmp = [stdout[i].split('  ')][0]
			address_file.write("""'LFN:""" + stdout_tmp[0] + """',\n""")
		stdout_tmp = [stdout[max_file].split('  ')][0]
		address_file.write("""'LFN:""" + stdout_tmp[0] + """'\n""")
		address_file.write("""], clear=True)\n""")
    else:
	print "You have to manually uncomment... "
    '''
    for f in path_list:
	command = 'lb-run LHCbDIRAC dirac-bookkeeping-genXMLCatalog --Options=' + path_list_dict[f]+'.py --Catalog=../xml_files/catalog_'+path_list_dict[f]+'.xml'
	print command
	xml_command  = subprocess.check_call(command, shell=True)
   
