#!/usr/bin/env python

"""
Bender module to run the following sequence over B2KShh signal MC samples:
- run an algorithm to store the MC truth DP position (and other info) for all generated events
- 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 =========#
    btype  = params.get( 'btype',  'Bd' )
    track1 = params.get( 'track1', 'K'  )
    track2 = params.get( 'track2', 'K'  )
    v0type = params.get( 'v0type', 'KS0' )
    whichMC = params.get( 'whichMC', '2012' )
    whichStripping = params.get( 'stripping', 'Stripping21' )
    isXGen = params.get( 'isXGen', False )
    printFreq = params.get( 'printFreq', 1000 )
    dddbtag = params.get( 'dddbtag', '' )
    conddbtag = params.get( 'conddbtag', '' )
    extended_hypos = params.get( 'extended_hypos', False )
    #=================================================#

    mode = btype+'2'+track1+track2+v0type
    mc_daughters = [ track1, track2, v0type ]

    # decay descriptors can be provided if running over resonant signal MC, e.g. B0 -> K*+ pi-; K*+ -> KS0 pi+
    mc_decay_descriptors = []

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

    knownMCTypes = [ '2011', '2012' ]

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

    nativeStrippingVersion = {}
    nativeStrippingVersion['2011'] = 'Stripping21r1'
    nativeStrippingVersion['2012'] = 'Stripping21'

    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

    algNamesAndLocations = {}
    if whichMC == '2011' :
        algNamesAndLocations['B2KShh_DD_2011'] = '/Event/AllStreams21/Phys/B2KShh_MVAFilter_DD_2011_OS/Particles'
        algNamesAndLocations['B2KShh_LL_2011'] = '/Event/AllStreams21/Phys/B2KShh_MVAFilter_LL_2011_OS/Particles'
    else :
        algNamesAndLocations['B2KShh_DD_2012a'] = '/Event/AllStreams21/Phys/B2KShh_MVAFilter_DD_2012a_OS/Particles'
        algNamesAndLocations['B2KShh_LL_2012a'] = '/Event/AllStreams21/Phys/B2KShh_MVAFilter_LL_2012a_OS/Particles'
        algNamesAndLocations['B2KShh_DD_2012b'] = '/Event/AllStreams21/Phys/B2KShh_MVAFilter_DD_2012b_OS/Particles'
        algNamesAndLocations['B2KShh_LL_2012b'] = '/Event/AllStreams21/Phys/B2KShh_MVAFilter_LL_2012b_OS/Particles'


    # Configuration of DaVinci
    from Configurables import DaVinci

    daVinci = DaVinci()
    daVinci.DataType        = whichMC
    daVinci.Simulation      = True
    daVinci.Lumi            = False
    daVinci.InputType       = "DST"
    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+'-withMCtruth.root'

    setData( datafiles, catalogues )

    gaudi = appMgr()

    from B2KShh.MCTruthAlgo import B2KShhMCTruth
    from B2KShh.RecoAlgo import B2KShhReco

    algGenMCTruth = B2KShhMCTruth( mode, btype, track1, track2, v0type, isXGen )

    reco_algs = []

    for algName,inputLoc in algNamesAndLocations.items() :
        alg = B2KShhReco(   algName,
                            reco_daughters,
                            extended_hypos, False, True, True,
                            mc_daughters,
                            mc_decay_descriptors,
                            PP2MCs = [ 'Relations/Rec/ProtoP/Charged'] ,
                            Inputs = [ inputLoc ]
                        )
        reco_algs.append( alg )

    userSeq = gaudi.algorithm('GaudiSequencer/DaVinciUserSequence' , True )

    userSeq.Members += [ algGenMCTruth.name() ]
    for alg in reco_algs :
        userSeq.Members += [ alg.name() ]

    return SUCCESS

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

if '__main__' == __name__ :

    datafiles = [ '/data/lhcb/phsdba/gangadir-KShh-Stripping21/workspace/phsdba/LocalXML/DSTs/Bd2KKKS0_MC_2012_MagDown_Pythia8/00029527_%08d_1.allstreams.dst' % i for i in range(1,117) ]

    for file in datafiles :
        print file

    pars = {}
    pars[ 'btype' ]     = 'Bd'
    pars[ 'track1' ]    = 'K'
    pars[ 'track2' ]    = 'K'
    pars[ 'v0type' ]    = 'KS0'
    pars[ 'whichMC' ]   = '2012'
    pars[ 'stripping' ] = 'Stripping21'
    pars[ 'dddbtag' ]   = 'Sim08-20130503-1'
    pars[ 'conddbtag' ] = 'Sim08-20130503-1-vc-md100'

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

    run(-1)

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