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' ]
strippings = [ 'Stripping21r1', 'Stripping21' ]
eventtypes = eventTypeDict.keys()
# *************************************************************************************** #

import os
import re

knownMagTypes = [ 'MagDown', 'MagUp' ]

known2011StrippingVersions = [ 'Stripping21r1' ]
known2012StrippingVersions = [ 'Stripping21' ]

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'

#baseDir = '/data/lhcb/phsdba/B2KShh/Sim08-MC/Signal/S21-DSTs'
baseDir = '/eos/lhcb/wg/BnoC/B2KShh/S21-SignalMC-DSTs'

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

        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

                # still need to search the bookkeeping to check that this is a valid configuration and to get the DB tags
                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

                s21_ds = LHCbDataset()
                for file in ds.files :
                        s21_file_name = baseDir + '/' + '%s_MC_%s_%s_Pythia%d' % ( mode, datatype, magtype, pythiaversion ) + '/' + os.path.basename( file.name )
                        s21_file = PhysicalFile( s21_file_name )
                        s21_ds.files.append( s21_file )

                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_B2KShh_signalMC_S21.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 = LSF()
                j.inputdata = s21_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 (LocalFile) or stored on the Grid (DiracFile)
                j.outputfiles = [LocalFile(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()