# ************************************************************************ #
# pick magnet polarities and stripping versions from the known types below #
# ************************************************************************ #
magtypes = [ 'MagDown', 'MagUp' ]
strippings = [ 'Stripping20r1', 'Stripping20' ]
# ************************************************************************ #

import os
import re

knownMagTypes = [ 'MagDown', 'MagUp' ]
known2011StrippingVersions = [ 'Stripping20r1', 'Stripping20r1p1' ]
known2012StrippingVersions = [ 'Stripping20',   'Stripping20r0p1' ]

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

recoVersion = {}
recoVersion['Stripping20'] = '14'

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')
moduleFile = modulePath+'/bender_B2KShh_data.py'

userReleaseArea = path[:match2.end()]


for stripping in strippings :

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

    print 'Creating job(s) for stripping: ' + stripping

    for magtype in magtypes :

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

        print 'With magnet setting: ' + magtype

        bkpath = '/LHCb/Collision'+datatype[2:]+'/Beam'+beamEnergy[datatype]+'GeV-VeloClosed-'+magtype+'/Real Data/Reco'+recoVersion[stripping[:11]]+'/'+stripping+'/90000000/BHADRONCOMPLETEEVENT.DST'
        print 'Using BK path: ' + bkpath
        bkq = BKQuery( type='Path', dqflag='OK', path=bkpath )
        ds = bkq.getDataset()

        # 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['magtype'] = magtype
        params['stripping'] = stripping

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

        j=Job()
        j.name = 'Collision'+datatype[2:]+'_'+magtype+'_'+stripping
        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 = 'B2KShh-Collision'+datatype[2:]+'-'+magtype+'-'+stripping+'.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 = 25, maxFiles = -1 ) 

        #j.submit()
        queues.add(j.submit)