#!/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 a particular Stripping version extracted from the StrippingArchive or from a local config in StrippingSelections and StrippingSettings
- 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',  'Lb' )
    track1 = params.get( 'track1', 'pi'  )
    track2 = params.get( 'track2', 'pi'  )
    v0type = params.get( 'v0type', 'Lz' )
    whichMC = params.get( 'whichMC', '2011' )
    whichStripping = params.get( 'stripping', 'Stripping20r1p2' )
    isXGen = params.get( 'isXGen', False )
    modesuffix = params.get( 'modesuffix', '' )
    stripFromArchiveOrLocal = params.get( 'stripFromArchiveOrLocal', 'Archive' )
    saveDst = params.get( 'saveDst', False )
    printFreq = params.get( 'printFreq', 1000 )
    dddbtag = params.get( 'dddbtag', '' )
    conddbtag = params.get( 'conddbtag', '' )
    extended_hypos = params.get( 'extended_hypos', False )
    #=================================================#

    modeName = btype+'2'+track1+track2+v0type
    mode = modeName + modesuffix
    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 = []

    knownMCTypes = [ '2011', '2012' ]

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

    knownStrippings = [ 'Stripping20', 'Stripping20r1', 'Stripping20r0p2', 'Stripping20r1p2' ]

    if whichStripping not in knownStrippings :
        e = Exception('Unsupported stripping version')
        raise e

    strippingYear = {}
    strippingYear['Stripping20']     = '2012'
    strippingYear['Stripping20r1']   = '2011'
    strippingYear['Stripping20r0p2'] = '2012'
    strippingYear['Stripping20r1p2'] = '2011'

    if whichMC != strippingYear[whichStripping] :
        e = Exception('Mismatch between stripping version and year')
        raise e

    strippingNames = {}
    strippingNames['Stripping20']     = 'Lb2V0hh'
    strippingNames['Stripping20r1']   = 'Lb2V0hh'
    strippingNames['Stripping20r0p2'] = 'Lb2V0h'
    strippingNames['Stripping20r1p2'] = 'Lb2V0h'

    strippingName = strippingNames[whichStripping]

    inputLocationDD = 'Phys/Lb2V0hhDDLine/Particles'
    inputLocationLL = 'Phys/Lb2V0hhLLLine/Particles'

    reco_daughters = [ 'pi', 'pi', 'Lambda0' ]

    # Configuration of Stripping
    if stripFromArchiveOrLocal == 'Archive' :
        from StrippingArchive import Utils
    elif stripFromArchiveOrLocal == 'Local' :
        from StrippingSelections import Utils
    else :
        e = Exception('Unknown Stripping config location')
        raise e
    lb = Utils.lineBuilder(whichStripping,strippingName)

    from StrippingConf.StrippingStream import StrippingStream
    stream1 = StrippingStream("Bhadron")
    stream1.appendLines( lb.lines() )

    from StrippingConf.Configuration import StrippingConf
    # was trying new approach with the EventNodeKiller - but it didn't seem to work :(
    sc = StrippingConf( HDRLocation = "ReStrip", Streams = [ stream1 ] )
    #sc = StrippingConf( Streams = [ stream1 ] )

    for stream in sc.activeStreams() :
        for line in stream.lines :
            print line.name(),' output=',line.outputLocation()

    from Configurables import StrippingReport
    sr = StrippingReport(Selections = sc.selections())
    sr.ReportFrequency = printFreq

    # Configuration of the EventNodeKiller to remove the old stripping reports
    #from Configurables import EventNodeKiller
    #eventNodeKiller = EventNodeKiller('Stripkiller')
    #eventNodeKiller.Nodes = [ '/Event/AllStreams', '/Event/Strip' ]

    # 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
    #daVinci.appendToMainSequence( [ eventNodeKiller ] )

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

    # Configuration of SelDSTWriter
    try :
        from DSTWriters.__dev__.Configuration import SelDSTWriter
    except ImportError :
        from DSTWriters.Configuration import SelDSTWriter

    suffix = mode+'-MC-'+whichMC+'-'+magtype+'-'+whichStripping

    dstWriter = SelDSTWriter( "MyDSTWriter",
                          OutputFileSuffix = suffix,
                          SelectionSequences = sc.activeStreams()
                          )

    seq = sc.sequence()

    setData( datafiles, catalogues )

    gaudi = appMgr()

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

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

    algLb2V0hhDD = B2KShhReco(
            'Lb2LzhhDD',
            reco_daughters,
            extended_hypos, False, True, True,
            mc_daughters,
            mc_decay_descriptors,
            PP2MCs = [ 'Relations/Rec/ProtoP/Charged'] ,
            Inputs = [ inputLocationDD ]
            )
    algLb2V0hhLL = B2KShhReco(
            'Lb2LzhhLL',
            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 += [ "EventNodeKiller/Stripkiller" ]
    userSeq.Members += [ algGenMCTruth.name() ]
    userSeq.Members += [ "GaudiSequencer/StrippingGlobal" ]
    userSeq.Members += [ sr.name() ]
    userSeq.Members += [ algLb2V0hhDD.name() ]
    userSeq.Members += [ algLb2V0hhLL.name() ]
    if saveDst:
        userSeq.Members += [ 'GaudiSequencer/MyDSTWriterMainSeq' ]

    return SUCCESS

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

if '__main__' == __name__ :

    datafiles = [
            'PFN:/data/lhcb/phsdba/Lb2Lhh/DST/MC2011/Lb2Lambdapipi-MagDown-Pythia6/00028874_00000010_1.allstreams.dst'
    ]

    pars = {}
    pars[ 'btype' ]     = 'Lb'
    pars[ 'track1' ]    = 'pi'
    pars[ 'track2' ]    = 'pi'
    pars[ 'v0type' ]    = 'Lambda0'
    pars[ 'whichMC' ]   = '2011'
    pars[ 'stripping' ] = 'Stripping20r1p2'

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

    run(-1)

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