eventTypeDict = {
    'Bd2pipiKS0-sqDalitz-DecProdCut' : 11104124
  , 'Bd2KpiKS0-sqDalitz-DecProdCut'  : 11304161
  , 'Bd2KKKS0-sqDalitz-DecProdCut'   : 11104114
  , 'Bs2pipiKS0-sqDalitz-DecProdCut' : 13104124
  , 'Bs2KpiKS0-sqDalitz-DecProdCut'  : 13304104
  , 'Bs2KKKS0-sqDalitz-DecProdCut'   : 13104134
}

# *************************************************************************************** #
# pick magnet polarities, stripping versions and an event type from the known types above #
# *************************************************************************************** #
magtypes = [ 'MagDown', 'MagUp' ]
years = [ '2011', '2012' ]
#years = [ '2011' ]
eventtypes = eventTypeDict.keys()
#eventtypes = [ 'Bd2pipiKS0-sqDalitz-DecProdCut' ]
# *************************************************************************************** #

import os
import re

knownMagTypes = [ 'MagDown', 'MagUp' ]

nativeStrippingVersion = {}
nativeStrippingVersion['2011'] = 'Stripping20r1'
nativeStrippingVersion['2012'] = 'Stripping20'

beamEnergy = {}
beamEnergy['2011'] = '3500'
beamEnergy['2012'] = '4000'

nuValue = {}
nuValue['2011'] = '2'
nuValue['2012'] = '2.5'

tck = {}
tck['2011'] = '0x40760037'
tck['2012'] = '0x409f0045'

currentdir = os.getcwd()
daVinciVersion = 'v36r1p2'

for eventtype in eventtypes :

    print 'Creating job(s) for event type: ' + eventtype

    if eventtype not in eventTypeDict.keys() :
        e = Exception('Unknown event type: ' + eventtype )
        raise e

    evtID = eventTypeDict[ eventtype ]

    bkinfo = getBKInfo( evtID )
    bkpaths = bkinfo.keys()

    mode = eventtype.split('-')[0]
    btype = mode.split('2')[0]
    tracks = re.findall( 'K(?!S)|pi|p(?!i)', mode )
    track1 = tracks[0]
    track2 = tracks[1]
    match = re.search( 'KS|KS0|Lambda|Lambda0', mode )
    v0type = match.group(0)
    if not v0type.endswith('0') :
        v0type = v0type+'0'

    for year in years :

        print 'For year: ' + year

        for magtype in magtypes :

            print 'With magnet setting: ' + magtype

            if magtype not in knownMagTypes :
                e = Exception('Unsupported magnet setting: ' + magtype)
                raise e

            for pythiaversion in [ 6, 8 ] :

                print 'With Pythia version: %d' % pythiaversion

                datasets = []
                dddbtags = set()
                condtags = set()
                for simversion in [ 'Sim08a', 'Sim08b' ] :
                    bkpath = '/MC/%s/Beam%sGeV-%s-%s-Nu%s-Pythia%d/%s/Digi13/Trig%s/Reco14a/%sNoPrescalingFlagged/%d/ALLSTREAMS.DST' % ( year, beamEnergy[year], year, magtype, nuValue[year], pythiaversion, simversion, tck[year], nativeStrippingVersion[year], evtID )
                    if bkpath not in bkpaths :
                        continue
                    print 'Trying BK path: ' + bkpath,
                    bkq = BKQuery( type='Path', dqflag='OK', path=bkpath )
                    dstmp = bkq.getDataset()
                    if len(dstmp) != 0 :
                        print ' - found dataset'
                        datasets.append( dstmp )
                        dddbtags.add( bkinfo[bkpath][0] )
                        condtags.add( bkinfo[bkpath][1] )
                    else :
                        print ' - nothing found'

                if len( datasets ) == 0 :
                    print 'Could not find any valid data!!  Skipping this configuration!!'
                    continue
                elif len(datasets) == 1 :
                    ds = datasets[0]
                elif len(datasets) == 2 :
                    ds = datasets[0].union( datasets[1] )
                else :
                    e = Exception('Found more than two datasets!')
                    raise e

                if len( ds.files ) == 0 :
                    e = Exception('Zero files in this dataset!')
                    raise e

                if len( dddbtags ) != 1 or len( condtags ) != 1 :
                    e = Exception('Found multiple DB tags')
                    raise e

                # uncomment this if you want to run a quick test on the CERN batch
                #reduced_ds = LHCbDataset()
                #for file in ds.files :
                    #sites = file.getReplicas().keys()
                    #for site in sites :
                        #if 'CERN' in site :
                            #reduced_ds.files.append( file )
                            #break
                    #if len(reduced_ds.files) > 0 :
                        #break

                if year == '2011' :
                    optsFiles = [ currentdir+'/DV-Stripping21r1-Stripping-MC.py',
                                  currentdir+'/DV-RedoCaloPID-Stripping21.py' ]
                elif year == '2012' :
                    optsFiles = [ currentdir+'/DV-Stripping21-Stripping-MC.py',
                                  currentdir+'/DV-RedoCaloPID-Stripping21.py' ]
                else :
                    e = Exception('Unrecognised year, what is going on?')
                    raise e

                dv = DaVinci(version=daVinciVersion)
                dv.optsfile = optsFiles
                dv.extraopts = """
DaVinci().DDDBtag = '%s'
DaVinci().CondDBtag = '%s'
                """ % ( dddbtags.pop(), condtags.pop() )

                j=Job()
                j.name = mode+'_MC_'+year+'_'+magtype+'_Pythia'+str(pythiaversion)
                j.application = dv
                j.backend = Dirac()
                j.inputdata = ds

                # uncomment this if you want to run a quick test on the CERN batch
                #j.backend = LSF( queue = '8nh' ) 
                #j.inputdata = reduced_ds

                # can pick if you want the ntuple returned to you immediately (SandboxFile) or stored on the Grid (DiracFile)
                j.outputfiles = [LocalFile('000000.AllStreams21.dst')]
                #j.outputfiles = [DiracFile('000000.AllStreams21.dst')]

                # can tweak the Dirac options if you like
                #j.backend.settings['CPUTime'] = 10000 
                #j.backend.settings['Destination'] = 'LCG.CERN.ch' 

                # can change here the number of files you want to run over per job
                j.splitter = SplitByFiles( filesPerJob = 1 ) 

                #j.submit()