eventTypeDict = {
    'Lb2pipiLambda0-sqDalitz-DecProdCut' : 15104125
  , 'Lb2KpiLambda0-sqDalitz-DecProdCut'  : 15304111
  , 'Lb2KKLambda0-sqDalitz-DecProdCut'   : 15104132
  , 'Lb2ppLambda0-sqDalitz-DecProdCut'   : 15104142
}

# *************************************************************************************** #
# pick magnet polarities, stripping versions and an event type from the known types above #
# *************************************************************************************** #
#magtypes = [ 'MagDown', 'MagUp' ]
#strippings = [ 'Stripping20r1', 'Stripping20', 'Stripping20r1p2', 'Stripping20r0p2' ]
#eventtypes = eventTypeDict.keys()
magtypes = [ 'MagUp' ]
strippings = [ 'Stripping20r1p2' ]
eventtypes = [ 'Lb2pipiLambda0-sqDalitz-DecProdCut' ]
# *************************************************************************************** #

import os
import re

knownMagTypes = [ 'MagDown', 'MagUp' ]

known2011StrippingVersions = [ 'Stripping20r1', 'Stripping20r1p2' ]
known2012StrippingVersions = [ 'Stripping20',   'Stripping20r0p2' ]

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'

path = os.getcwd()
pathList = path.split('/')

match1 = re.search( '/Bender_(v\d+r\d+[p(?=\d)]?[(?<=p)\d]*)/', path )
match2 = re.search( 'cmtuser', path )

if not path.endswith('job') or 'B2KShh' not in pathList or not match1 or not match2 :
    e = Exception('You do not appear to be in the \'job\' directory of the B2KShh package within a Bender project')
    raise e

benderVersion = match1.groups()[0]

modulePath = path.replace('job','options')

userReleaseArea = path[:match2.end()]


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 stripping in strippings :

        print 'With stripping: ' + stripping

        datatype = ''
        if stripping in known2011StrippingVersions:
            datatype = '2011'
        elif stripping in known2012StrippingVersions:
            datatype = '2012'
        else :
            e = Exception('Unsupported Stripping version: ' + stripping)
            raise e

        restrip = False
        if stripping != nativeStrippingVersion[datatype] :
            restrip = True

        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' % ( datatype, beamEnergy[datatype], datatype, magtype, nuValue[datatype], pythiaversion, simversion, tck[datatype], nativeStrippingVersion[datatype], 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

                params = {}
                params = {}
                params['btype']     = btype
                params['track1']    = track1
                params['track2']    = track2
                params['v0type']    = v0type
                params['whichMC']   = datatype
                params['stripping'] = stripping
                params['dddbtag']   = dddbtags.pop()
                params['conddbtag'] = condtags.pop()

                moduleFile = modulePath+'/bender_Lb2V0hh_signalMC.py'
                if restrip :
                    moduleFile = modulePath+'/bender_Lb2V0hh_stripping_signalMC.py'

                b = Bender(version=benderVersion)
                b.user_release_area = userReleaseArea
                b.module = File(moduleFile)
                b.params = params

                j=Job()
                j.name = mode+'_MC_'+datatype+'_'+magtype+'_'+stripping+'_Pythia'+str(pythiaversion)
                j.application = b
                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

                # NB remember to change the tuple name in the Bender script to match this!
                tupleFile = mode+'-MC-'+datatype+'-'+magtype+'-'+stripping+'-withMCtruth.root'

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

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

                #j.submit()