#!/usr/bin/env python

"""
Bender module to run the following sequence over B2KShh background MC samples:
- run an algorithm to store the reco and MC truth-matched info for all Stripped candidates
"""

from Bender.Main import *

def configure( datafiles, catalogues = [], params = {}, castor = False ) :

    #======= B->KShh Configuration and Setup =========#
    mode = params.get( 'mode', 'Bd2etapKS' )
    whichMC = params.get( 'whichMC', '2011' )
    whichStripping = params.get( 'stripping', 'Stripping20r1' )
    printFreq = params.get( 'printFreq', 1000 )
    dddbtag = params.get( 'dddbtag', '' )
    conddbtag = params.get( 'conddbtag', '' )
    extended_hypos = params.get( 'extended_hypos', True )
    #=================================================#

    reco_daughters = [ 'pi', 'pi', 'KS0' ]

    knownMCTypes = [ '2011', '2012' ]

    if whichMC not in knownMCTypes :
        e = Exception('Unsupported MC version')
        raise e

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

    if whichStripping != nativeStrippingVersion[ whichMC ] :
        e = Exception('Requested stripping version %s is not the native version for this MC %s, you need to use the version of the script that will first re-strip the MC.' % (whichStripping, whichMC) )
        raise e

    inputLocationDD = '/Event/AllStreams/Phys/B2KShhDDLine/Particles'
    inputLocationLL = '/Event/AllStreams/Phys/B2KShhLLLine/Particles'

    from PhysConf.Filters import LoKi_Filters

    filterVoid  = "EXISTS('/Event/Strip/Phys/DecReports')"
    filterStrip = "HLT_PASS_RE('StrippingB2KShh[DL]{2}LineDecision')"
    filterAlg = LoKi_Filters( VOID_Code = filterVoid, STRIP_Code = filterStrip  )
    filters = filterAlg.filters('Filters')
    filters.reverse()

    from Configurables import DaVinci

    daVinci = DaVinci()
    daVinci.DataType        = whichMC
    daVinci.Simulation      = True
    daVinci.Lumi            = False
    daVinci.InputType       = "DST"
    daVinci.EventPreFilters = filters
    daVinci.EvtMax          = -1
    daVinci.PrintFreq       = printFreq

    ## try to get the tags from Rec/Header
    if dddbtag != '' and conddbtag != '' :
        daVinci.DDDBtag = dddbtag
        daVinci.CondDBtag = conddbtag
    else :
        from BenderTools.GetDBtags import getDBTags
        tags = getDBTags ( datafiles[0] , castor  )

        logger.info ( 'Extract tags from DATA : %s' % tags )
        if tags.has_key ( 'DDDB' ) and tags ['DDDB'] :
            daVinci.DDDBtag   = tags['DDDB'  ]
            logger.info ( 'Set DDDB    %s ' % daVinci.DDDBtag   )
        if tags.has_key ( 'CONDDB' ) and tags ['CONDDB'] :
            daVinci.CondDBtag = tags['CONDDB']
            logger.info ( 'Set CONDDB  %s ' % daVinci.CondDBtag )
        if tags.has_key ( 'SIMCOND' ) and tags ['SIMCOND'] :
            daVinci.CondDBtag = tags['SIMCOND']
            logger.info ( 'Set SIMCOND %s ' % daVinci.CondDBtag )

    magtype = "MagUp"
    if "md" in daVinci.CondDBtag :
        magtype = "MagDown"

    daVinci.TupleFile = mode+'-MC-'+whichMC+'-'+magtype+'-'+whichStripping+'.root'

    setData( datafiles, catalogues )

    gaudi = appMgr()

    from B2KShh.RecoAlgo import B2KShhReco

    algB2KShhDD = B2KShhReco(
            'B2KShhDD',
            reco_daughters,
            extended_hypos, False, True, False,
            PP2MCs = [ 'Relations/Rec/ProtoP/Charged'] ,
            Inputs = [ inputLocationDD ]
            )
    algB2KShhLL = B2KShhReco(
            'B2KShhLL',
            reco_daughters,
            extended_hypos, False, True, False,
            PP2MCs = [ 'Relations/Rec/ProtoP/Charged'] ,
            Inputs = [ inputLocationLL ]
            )

    userSeq = gaudi.algorithm('GaudiSequencer/DaVinciUserSequence' , True )
    userSeq.Members += [ algB2KShhDD.name() ]
    userSeq.Members += [ algB2KShhLL.name() ]

    return SUCCESS

#############

if '__main__' == __name__ :

    datafiles = [
            'PFN:/data/lhcb/phsdba/B2KShh/DST/MC/2011/Bd2etapKS-MagUp-Pythia8/00030589_00000001_1.allstreams.dst'
    ]

    pars = {}
    pars[ 'mode' ]      = 'Bd2etapKS'
    pars[ 'whichMC' ]   = '2011'
    pars[ 'stripping' ] = 'Stripping20r1'

    configure( datafiles, params = pars, castor=False )

    run(-1)

#############