from Configurables import (
    DaVinci,
    EventSelector,
    PrintMCTree,
    MCDecayTreeTuple, 
    TupleToolMCTruth
)
from DecayTreeTuple.Configuration import *

"""Configure the variables below with:
decay: Decay you want to inspect, using 'newer' LoKi decay descriptor syntax,
decay_heads: Particles you'd like to see the decay tree of,
datafile: Where the file created by the Gauss generation phase is, and
year: What year the MC is simulating.
"""

# Select correct MC sample
RefChan = False
PHSP = True

# https://twiki.cern.ch/twiki/bin/view/LHCb/FAQ/LoKiNewDecayFinders
decay = "[B+ => ^(K*(892)+ => ^K+ ^(pi0=> ^gamma ^gamma)) ^mu- ^mu+]CC"
if (RefChan): decay = "[B+ => ^(K*(892)+ => ^K+ ^(pi0=> ^gamma ^gamma)) ^(J/psi(1S) => ^mu- ^mu+)]CC"
decay_heads = ["B+", "B-"]

#Path to datafile
datafile = "/afs/cern.ch/work/r/rekopecn/public/B_2_KstarMuMu/EvtGen/GaussDev_v49r12/MCTruth.sim"
if (RefChan): datafile =datafile.replace(".sim","_RefChan.sim",1)
if (PHSP):    datafile =datafile.replace(".sim","_PHSP.sim",1)

year = 2017

# For a quick and dirty check, you don't need to edit anything below here.
##########################################################################

# Create an MC DTT containing any candidates matching the decay descriptor
mctuple = MCDecayTreeTuple("MCDecayTreeTuple")
mctuple.Decay = decay

# Name of the .xgen file produced by Gauss
EventSelector().Input = ["DATAFILE='{0}' TYP='POOL_ROOTTREE' Opt='READ'".format(datafile)]

# Configure DaVinci
DaVinci().TupleFile =datafile.replace(".sim","_orig.root",1)
DaVinci().HistogramFile = "DVntuple_histo.root"
DaVinci().Simulation = True
DaVinci().Lumi = False
DaVinci().DataType = str(year)
DaVinci().UserAlgorithms = [mctuple]

def doIt():
    """
    specific post-config action for (x)GEN-files 
    """
    extension = "xgen"
    ext = extension.upper()

    from Configurables import DataOnDemandSvc
    dod  = DataOnDemandSvc ()
    from copy import deepcopy 
    algs = deepcopy ( dod.AlgMap ) 
    bad  = set() 
    for key in algs :
        if     0 <= key.find ( 'Rec'     )                  : bad.add ( key )
        elif   0 <= key.find ( 'Raw'     )                  : bad.add ( key )
        elif   0 <= key.find ( 'DAQ'     )                  : bad.add ( key )
        elif   0 <= key.find ( 'Trigger' )                  : bad.add ( key )
        elif   0 <= key.find ( 'Phys'    )                  : bad.add ( key )
        elif   0 <= key.find ( 'Prev/'   )                  : bad.add ( key )
        elif   0 <= key.find ( 'Next/'   )                  : bad.add ( key )
        elif   0 <= key.find ( '/MC/'    ) and 'GEN' == ext : bad.add ( key )
        
    for b in bad :
        del algs[b]
            
    dod.AlgMap = algs
    
    from Configurables import EventClockSvc, CondDB 
    EventClockSvc ( EventTimeDecoder = "FakeEventTime" )
    CondDB  ( IgnoreHeartBeat = True )
    
appendPostConfigAction( doIt )