#!/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', 'Stripping20' )
    isXGen = params.get( 'isXGen', False )
    printFreq = params.get( 'printFreq', 1000 )
    dddbtag = params.get( 'dddbtag', '' )
    conddbtag = params.get( 'conddbtag', '' )
    extended_hypos = params.get( 'extended_hypos', True )
    #=================================================#

    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'] = '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'


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

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

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

    return SUCCESS

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

if '__main__' == __name__ :

    datafiles = [
            #'/data/lhcb/phsdba/B2KShh/DST/MC/2012/Bd2KSKK-SqDalitz-MagDown-Pythia8/00029527_00000114_1.allstreams.dst'
            '/disk/users/elena/B2KShh/ntuples/tests/00031273_00000013_1.allstreams.dst'
    ]

    pars = {}
    pars[ 'btype' ]     = 'Xib0'  # 'Bd'
    pars[ 'track1' ]    = 'p'  # 'K'
    pars[ 'track2' ]    = 'pi'  # 'K'
    pars[ 'v0type' ]    = 'KS0'
    pars[ 'whichMC' ]   = '2011'  # '2012'
    pars[ 'stripping' ] = 'Stripping20r1'  # 'Stripping20'
    pars[ 'dddbtag' ]   = 'dddb-20130929'  # 'Sim08-20130503-1'
    pars[ 'conddbtag' ] = 'sim-20130522-vc-md100'  # 'Sim08-20130503-1-vc-md100'

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

    run(300)

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