############# Global settings
runFilters = False
MDST = True
data = False
down = False
year = "2015"
stream = "AllStreams"
line = "B2XMuMu_Line"

channel = "B0toKstJpsi"

#subdecay chains
KShortPiPlus = False
KplusPi0Merged = False
KplusPi0Resolved = True
from Gaudi.Configuration import *

from Configurables import DaVinci
from Configurables import FilterDesktop, GaudiSequencer
from Configurables import CombineParticles
from Configurables import TupleToolP2VV, MCTupleToolP2VV


#---------- Debugging tools----------#
#MessageSvc().OutputLevel = VERBOSE
#MessageSvc().OutputLevel = DEBUG
#MessageSvc().Format = "% F%100W%S%7W%R%T %0W%M"

#Run1 or Run2
#DST or MDST?
Run1 = True
if (year == "2011" or year == "2012"):
    Run1 = True
elif (year == "2015" or year == "2016"):
    Run1 = False
else:
    print "[WARNING]\tWrong year input!\n"

#MC Truthmatching fix for run 1 (obsolete)
#if (not data and Run1):
   # from Configurables import PhysConf
   # PhysConf().CaloReProcessing=True

#triggerlines
if (Run1):
    print "[INFO]\tTriggerlines for Run 1 added!\n"
    triggerlines = [
		#l0
		        "L0HadronDecision",
                "L0MuonDecision",
                "L0GlobalDecision",

		#hlt1
                "Hlt1TrackMuonDecision",
                "Hlt1DiMuonHighMassDecision",
                "Hlt1DiMuonLowMassDecision",
                "Hlt1TrackAllL0Decision",
                "Hlt1TrackAllL0TightDecision",
                "Hlt1L0AnyDecision",
                "Hlt1MBNoBiasDecision",
                "Hlt1GlobalDecision",

         #hlt2
                "Hlt2CharmHadD2HHHHDecision",
                "Hlt2CharmHadD2HHHHWideMassDecision",
                "Hlt2IncPhiDecision",
                 "Hlt2ExpressDs2PhiPiDecision",
                 "Hlt2Topo2BodySimpleDecision",
                 "Hlt2Topo3BodySimpleDecision",
                 "Hlt2Topo2BodyBBDTDecision",
                 "Hlt2Topo3BodyBBDTDecision",
                 "Hlt2TopoMu2BodyBBDTDecision",
                 "Hlt2TopoMu3BodyBBDTDecision",
                 "Hlt2TopoE2BodyBBDTDecision",
                 "Hlt2TopoE3BodyBBDTDecision",
                 "Hlt2RadiativeTopoTrackTOSDecision",
                 "Hlt2RadiativeTopoPhotonL0Decision",
                 "Hlt2SingleMuonDecision",
                 "Hlt2DiMuonDetachedDecision",
                 "Hlt2DiMuonDetachedHeavyDecision"]

else:
    print "[INFO]\tTriggerlines for Run 2 added!\n"
    triggerlines = [

            #l0
                "L0HadronDecision",
                "L0MuonDecision",
		"L0DiMuonDecision",
                #"L0GlobalDecision",

            #hlt1
                "Hlt1TrackMuonDecision",
                "Hlt1DiMuonHighMassDecision",
                "Hlt1DiMuonLowMassDecision",
                "Hlt1TrackMVADecision",
		"Hlt1TwoTrackMVADecision",
                "Hlt1L0AnyDecision",
                "Hlt1MBNoBiasDecision",
                "Hlt1GlobalDecision",

             #hlt2
                 "Hlt2Topo2BodyDecision",
                 "Hlt2Topo3BodyDecision",
                 "Hlt2TopoMu2BodyDecision",
                 "Hlt2TopoMu3BodyDecision",
                 "Hlt2TopoMuMu2BodyDecision",
                 "Hlt2TopoMuMu3BodyDecision",
                 "Hlt2SingleMuonDecision",
                 "Hlt2SingleMuonHighPTDecision",
                 "Hlt2SingleMuonLowPTDecision",
                 "Hlt2DiMuonDetachedDecision",
                 "Hlt2DiMuonDetachedHeavyDecision"]

###############   Pre Filter, does not really do much except choose only candidates passing the Stripping line, maybe beneficial to performance
#from Configurables import LoKi__HDRFilter as StripFilter
#stripFilter = StripFilter( "stripPassFilter",\
#                          Code = "HLT_PASS('StrippingB2XMuMu_LineDecision')",\
#                          Location= "/Event/Strip/Phys/DecReports")

############# DecayTreeTuple
from DecayTreeTuple.Configuration import *
from Configurables import TupleToolTISTOS
from PhysSelPython.Wrappers import AutomaticData, Selection, SelectionSequence
from Configurables import PrintDecayTree, PrintDecayTreeTool

############# MCDecayTreeTuple
from Configurables import MCDecayTreeTuple
from Configurables import MCTupleToolKinematic, TupleToolMCBackgroundInfo, TupleToolMCTruth
    #Fix Truthmatching
from Configurables import BackgroundCategory,Calo2MCTool,DaVinciSmartAssociator

##subpid stuff
#from Configurables import SubPIDMMFilter
from Configurables import SubstitutePID ,BTaggingTool
from Configurables import TupleToolDecayTreeFitter, TupleToolTrackIsolation, TupleToolConeIsolation
from Configurables import LoKi__Hybrid__TupleTool


##########################################
###B+ -> (K*(892)+ -> KS0 pi+) mu+ mu- ### (> 1 <)
##########################################
    ### AND/OR ###
##########################################################
###B+ -> (K*(892)+ -> K+ (pi0 -> gamma gamma)) mu+ mu- ### (> 2 <)
##########################################################

#This is probably a very stupid idea but python should be okay with it
b2KstmumuTupleList= []

b2KstmumuTupleList.append(DecayTreeTuple("b2KstKs0pimumu_Tuple"))
b2KstmumuTupleList.append(DecayTreeTuple("b2KstKpi0mumuResolvedTuple"))

for b2KstmumuTuple in b2KstmumuTupleList:
    if (MDST):
        b2KstmumuTuple.RootInTES =  '/Event/{0}'.format(stream)
        b2KstmumuTuple.Inputs =  ['Phys/{0}/Particles'.format(line)]
    else:
        b2KstmumuTuple.Inputs =  ['{0}/Phys/{1}/Particles'.format(stream, line)]

#decay descriptor
b2KstmumuTupleList[0].Decay = "[B+ -> ^(K*(892)+ -> ^(KS0 -> ^pi+ ^pi-) ^pi+) ^(J/psi(1S) -> ^mu- ^mu+)]CC"
b2KstmumuTupleList[0].addBranches({
    "B_plus" : 	    "^([B+ -> (K*(892)+ -> (KS0 -> pi+ pi-) pi+) (J/psi(1S) -> mu- mu+)]CC)" ,
    "K_star_plus" : "[B+ -> ^(K*(892)+ -> (KS0 -> pi+ pi-) pi+) (J/psi(1S) -> mu- mu+)]CC",
    "K_short" :     "[B+ -> (K*(892)+ -> ^(KS0 -> pi+ pi-) pi+) (J/psi(1S) -> mu- mu+)]CC",
    "Ks_pi_plus" :  "[B+ -> (K*(892)+ -> (KS0 -> ^pi+ pi-) pi+) (J/psi(1S) -> mu- mu+)]CC",
    "Ks_pi_minus" : "[B+ -> (K*(892)+ -> (KS0 -> pi+ ^pi-) pi+) (J/psi(1S) -> mu- mu+)]CC",
    "pi_plus" :     "[B+ -> (K*(892)+ -> (KS0 -> pi+ pi-) ^pi+) (J/psi(1S) -> mu- mu+)]CC",
    "J_psi" : 	    "[B+ -> (K*(892)+ -> (KS0 -> pi+ pi-) pi+) ^(J/psi(1S) -> mu- mu+)]CC",
    "mu_minus" :    "[B+ -> (K*(892)+ -> (KS0 -> pi+ pi-) pi+) (J/psi(1S) -> ^mu- mu+)]CC",
    "mu_plus" :     "[B+ -> (K*(892)+ -> (KS0 -> pi+ pi-) pi+) (J/psi(1S) -> mu- ^mu+)]CC"
    })

b2KstmumuTupleList[1].Decay = "[B+ -> ^(K*(892)+ -> ^K+ ^(pi0 -> ^gamma ^gamma)) ^(J/psi(1S) -> ^mu- ^mu+)]CC"
b2KstmumuTupleList[1].addBranches({
    "B_plus" :           "^([B+ -> (K*(892)+ -> K+ (pi0 -> gamma gamma)) (J/psi(1S) -> mu- mu+)]CC)",
    "K_star_plus" :      "[B+ -> ^(K*(892)+ -> K+ (pi0 -> gamma gamma)) (J/psi(1S) -> mu- mu+)]CC",
    "K_plus" :           "[B+ -> (K*(892)+ -> ^K+ (pi0 -> gamma gamma)) (J/psi(1S) -> mu- mu+)]CC",
    "pi_zero_resolved" : "[B+ -> (K*(892)+ -> K+ ^(pi0 -> gamma gamma)) (J/psi(1S) -> mu- mu+)]CC",
    "gamma1" : 		 "[B+ -> (K*(892)+ -> K+ (pi0 -> ^gamma gamma)) (J/psi(1S) -> mu- mu+)]CC",
    "gamma2" : 		 "[B+ -> (K*(892)+ -> K+ (pi0 -> gamma ^gamma)) (J/psi(1S) -> mu- mu+)]CC",
    "J_psi" : 		 "[B+ -> (K*(892)+ -> K+ (pi0 -> gamma gamma)) ^(J/psi(1S) -> mu- mu+)]CC",
    "mu_minus" : 	 "[B+ -> (K*(892)+ -> K+ (pi0 -> gamma gamma)) (J/psi(1S) -> ^mu- mu+)]CC",
    "mu_plus" : 	 "[B+ -> (K*(892)+ -> K+ (pi0 -> gamma gamma)) (J/psi(1S) -> mu- ^mu+)]CC"
    })


#add tools
for b2KstmumuTuple in b2KstmumuTupleList:
    b2KstmumuTuple.ToolList =  ["TupleToolKinematic",
                                 "TupleToolPrimaries",
                                 "TupleToolEventInfo",
                                 "TupleToolTrackInfo",
                                 "TupleToolRecoStats",
                                 "TupleToolAngles",
                                 "TupleToolPid",
                                 "TupleToolPhotonInfo"]

    #add MCtools
    if not (data):
        b2KstmumuMCTuple = b2KstmumuTuple.addTupleTool("TupleToolMCTruth/b2KstmumuMCTuple")
        b2KstmumuMCTuple.addTupleTool("MCTupleToolHierarchy")
        b2KstmumuTuple.ToolList += ["MCTupleToolKinematic",
					"TupleToolMCBackgroundInfo"]
	#Adding Angle Calculator
	b2KstmumuTuple.B_plus.addTool(TupleToolMCTruth, name="B_plusMCTruth")
	b2KstmumuTuple.B_plus.ToolList = [ "TupleToolMCTruth/B_plusMCTruth"]
	b2KstmumuTuple.B_plus.B_plusMCTruth.ToolList = [ "MCTupleToolP2VV/B2KsMuMu"]
	b2KstmumuTuple.B_plus.B_plusMCTruth.addTool( MCTupleToolP2VV, name="B2KsMuMu" )
	b2KstmumuTuple.B_plus.B_plusMCTruth.B2KsMuMu.Calculator = 'MCBu2KstarMuMuAngleCalculator'

    if ( (channel == "BtoXJpsi" and not Run1) or (channel == "B0toKstMuMu" and  not Run1) ):
        from Configurables import Calo2MCTool
        for b2KstmumuTuple in b2KstmumuTupleList:
            b2KstmumuTuple.addTool(Calo2MCTool,name="Calo2MCTool")
            b2KstmumuTuple.Calo2MCTool.Hypo2Cluster=True

#add tools to branches
if (KplusPi0Resolved):
    b2KstmumuTupleList[1].pi_zero_resolved.addTupleTool("TupleToolPi0Info")
    b2KstmumuTupleList[1].pi_zero_resolved.addTupleTool("TupleToolCaloHypo")
    b2KstmumuTupleList[1].pi_zero_resolved.addTupleTool("TupleToolGeometry/Geo")
    b2KstmumuTupleList[1].pi_zero_resolved.Geo.Verbose = True

    b2KstmumuTupleList[1].gamma1.addTupleTool("TupleToolProtoPData")
    b2KstmumuTupleList[1].gamma1.addTupleTool("TupleToolPhotonInfo")
    b2KstmumuTupleList[1].gamma1.addTupleTool("TupleToolCaloHypo")
    b2KstmumuTupleList[1].gamma2.addTupleTool("TupleToolProtoPData")
    b2KstmumuTupleList[1].gamma2.addTupleTool("TupleToolPhotonInfo")
    b2KstmumuTupleList[1].gamma2.addTupleTool("TupleToolCaloHypo")


for b2KstmumuTuple in b2KstmumuTupleList:

    b2KstmumuTuple.B_plus.addTupleTool("TupleToolGeometry/Geo")
    b2KstmumuTuple.B_plus.Geo.Verbose = True
    b2KstmumuTuple.B_plus.addTupleTool("TupleToolPropertime")

	#Adding Angle Calculator
    TupleToolP2VV_Bu = TupleToolP2VV('TupleToolP2VV_Bu')
    TupleToolP2VV_Bu.Calculator = 'Bu2KstarMuMuAngleCalculator'
    b2KstmumuTuple.B_plus.addTool(TupleToolP2VV_Bu)
    b2KstmumuTuple.B_plus.ToolList +=  ["TupleToolP2VV/TupleToolP2VV_Bu"]

    b2KstmumuTuple.J_psi.addTupleTool("TupleToolGeometry/Geo")
    b2KstmumuTuple.J_psi.Geo.Verbose = True
    b2KstmumuTuple.mu_minus.addTupleTool("TupleToolGeometry/Geo")
    b2KstmumuTuple.mu_minus.Geo.Verbose = True
    b2KstmumuTuple.mu_plus.addTupleTool("TupleToolGeometry/Geo")
    b2KstmumuTuple.mu_plus.Geo.Verbose = True

if (KShortPiPlus):
    b2KstmumuTupleList[0].K_short.addTupleTool("TupleToolGeometry/Geo")
    b2KstmumuTupleList[0].K_short.Geo.Verbose = True
    b2KstmumuTupleList[0].pi_plus.addTupleTool("TupleToolGeometry/Geo")
    b2KstmumuTupleList[0].pi_plus.Geo.Verbose = True

if (KplusPi0Resolved):
    b2KstmumuTupleList[1].K_plus.addTupleTool("TupleToolGeometry/Geo")
    b2KstmumuTupleList[1].K_plus.Geo.Verbose = True

### DST isolation info tools ###
for b2KstmumuTuple in b2KstmumuTupleList:
    if not( MDST ):
        b2KstmumuTuple.addTupleTool("TupleToolTrackIsolation/TrackIso")
        b2KstmumuTuple.TrackIso.Verbose = True
        b2KstmumuTuple.addTupleTool("TupleToolConeIsolation/ConeIso")

    b2KstmumuTuple.addTool(TupleToolDecay, name="TupleToolDecayBplus")

    #decay tree fitter (DTF)
    b2KstmumuTuple.B_plus.addTupleTool('TupleToolDecayTreeFitter/DTF')
    b2KstmumuTuple.B_plus.DTF.constrainToOriginVertex = True
    b2KstmumuTuple.B_plus.DTF.UpdateDaughters = True
    b2KstmumuTuple.B_plus.DTF.Verbose = True

#Add mass constraints
if (KShortPiPlus):
    b2KstmumuTupleList[0].B_plus.DTF.daughtersToConstrain = ['KS0']
if (KplusPi0Resolved):
    b2KstmumuTupleList[1].B_plus.DTF.daughtersToConstrain = ['pi0']

#LoKi tools
ConeIsoInfo = 'ConeIsoInfo'
VtxIsoInfo = 'VtxIsoInfo'
if (not Run1 and not MDST):
    if ( (channel == "BtoXJpsi") or (channel == "B0toKstMuMu")):
        ConeIsoInfo = 'ConeIsoInfoBDT'
        ConeIsoInfo = 'VtxIsoInfoBDT'


if ( MDST ) :
	for b2KstmumuTuple in b2KstmumuTupleList:
    		LoKiToolb2KstmumuTuple = b2KstmumuTuple.addTupleTool("LoKi::Hybrid::TupleTool/LoKiToolb2KstmumuTuple")
	        LoKiToolb2KstmumuTuple.Variables = {"ETA" : "ETA" ,
                          "PHI" : "PHI" ,
                          "DOCA1" : "DOCA(1,2)" ,
                          "CONEANGLE" : "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + ConeIsoInfo + "', 'CONEANGLE', -1.)",
                          "CONEMULT" :  "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + ConeIsoInfo + "', 'CONEMULT', -1.)",
                          'CONEPTASYM' : "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + ConeIsoInfo + "', 'CONEPTASYM', -1.)",
                          'CONEPT': "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + ConeIsoInfo + "', 'CONEPT', -1.)",
                           'CONEP': "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + ConeIsoInfo + "', 'CONEP', -1.)",
                           'CONEPASYM' : "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + ConeIsoInfo + "', 'CONEPASYM', -1.)",
                           'CONEDELTAETA' : "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + ConeIsoInfo + "', 'CONEDELTAETA', -1.)",
                           'CONEDELTAPHI' : "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + ConeIsoInfo + "', 'CONEDELTAPHI', -1.)",
                           'VTXISONUMVTX' : "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + VtxIsoInfo + "', 'VTXISONUMVTX', -1.)",
                           'VTXISODCHI2ONETRACK': "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + VtxIsoInfo + "', 'VTXISODCHI2ONETRACK', -1.)",
                           'VTXISODCHI2MASSONETRACK' : "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + VtxIsoInfo + "', 'VTXISODCHI2MASSONETRACK', -1.)",
                           'VTXISODCHI2TWOTRACK' : "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + VtxIsoInfo + "', 'VTXISODCHI2TWOTRACK', -1.)",
                          'VTXISODCHI2MASSTWOTRACK': "RELINFO('/Event/" + stream + "/Phys/B2XMuMu_Line/" + VtxIsoInfo + "', 'VTXISODCHI2MASSTWOTRACK', -1.)"
                         };

#trigger config
for b2KstmumuTuple in b2KstmumuTupleList:
    b2KstmumuTT = b2KstmumuTuple.addTupleTool(TupleToolTISTOS)
    b2KstmumuTT.TriggerList = triggerlines
    b2KstmumuTT.FillL0 = True
    b2KstmumuTT.FillHlt1 = True
    b2KstmumuTT.FillHlt2 = True
    b2KstmumuTT.Verbose = True
    b2KstmumuTT.VerboseL0 = True
    b2KstmumuTT.VerboseHlt1 = True
    b2KstmumuTT.VerboseHlt2 = True

###########################
##Main sequence
###########################

b2KstmumuSeq = GaudiSequencer("b2KstmumuSeq")
if (KShortPiPlus):
    b2KstmumuSeq.Members += [b2KstmumuTupleList[0]]
if (KplusPi0Resolved):
    b2KstmumuSeq.Members += [b2KstmumuTupleList[1]]
b2KstmumuSeq.ModeOR = True
b2KstmumuSeq.ShortCircuit = False

###########################
##DaVinci configurations
###########################

dv = DaVinci()

#dv.EventPreFilters = [stripFilter]
dv.UserAlgorithms += [b2KstmumuSeq]

if(MDST):
	dv.InputType = "MDST"
else:
	dv.InputType = "DST"


dv.RootInTES = "/Event/{0}".format(stream)
dv.DataType = year
dv.Simulation = not data
dv.Lumi = True
#dv.VerboseMessages = True
#dv.SkipEvents = 0
dv.PrintFreq = 1000
dv.EvtMax = -1
dv.TupleFile = "B0toKstJpsi_2015MagUp.root"
dv.CondDBtag = "sim-20161124-vc-mu100"
dv.DDDBtag = "dddb-20170721-3"