############# Global settings
runFilters = False
reference = False
phsp = False
BkgStudy = False
data = True
down = True
year = "2015"
stream = "Leptonic"
line = "B2XMuMu_Line"

#subdecay chains
KShortPiPlus = True
KplusPi0Merged = False
KplusPi0Resolved = True


from Gaudi.Configuration import *

from Configurables import DaVinci

from Configurables import FilterDesktop

from Configurables import CombineParticles

from Configurables import GaudiSequencer

from Configurables import TupleToolP2VV, MCTupleToolP2VV

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

#DST or MDST? #This is a horrible piece of code but oh well
if (BkgStudy): 
    MDST = False
else: 
    if( not data and Run1 ):
       MDST = False
if (not KShortPiPlus): 
    if (year == "2015" and (not data) and (not phsp) and (not reference)): 
        MDST = False
    if (year == "2016" and (not data) and (not phsp) and (not reference)): 
        MDST = False

if (year == "2018" and (not data) and (not phsp) and (not reference)): 
    MDST = False

if (reference and Run1 == False and not KShortPiPlus):
    MDST = False


if (MDST): 
    print "\n[INFO]\tUsing MDST as input.\n"
else:
    print "\n[INFO]\tUsing DST as 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'


#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")
    if (year != "2018" and year !="2015"): b2KstmumuTupleList[1].gamma1.addTupleTool("TupleToolCaloHypo")
    b2KstmumuTupleList[1].gamma2.addTupleTool("TupleToolProtoPData")
    b2KstmumuTupleList[1].gamma2.addTupleTool("TupleToolPhotonInfo")
    if (year != "2018" and year !="2015"): 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.K_star_plus.addTupleTool("TupleToolGeometry/Geo")
    b2KstmumuTuple.K_star_plus.Geo.Verbose = True
    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
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
    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.EvtMax = -1 #10000 #100000 #100000
dv.PrintFreq = 1000
dv.TupleFile = "b2Kstmumu.root"

#This should have been a dictionary, big time, but oh well :)
#tags from http://lhcb-release-area.web.cern.ch/LHCb-release-area/DOC/dbase/conddb/release_notes.html

DBtags_data = {
    "2011": ["cond-20150409-1", "dddb-20150928"],
    "2012": ["cond-20150409-1", "dddb-20150928"],
    "2015": ["cond-20150828",   "dddb-20150724"],
    "2016": ["cond-20161004",   "dddb-20150724"],
    "2017": ["cond-20170724",   "dddb-20170721-3"],
    "2018": ["cond-20180202",   "dddb-20171030-3"]
}

DBtags_Kplus_sigMC = {
    "2011": ["20160614-1", "20160318-1"],
    "2012": ["20160321-2", "20150928"   ],
    "2015": ["20161124",   "20170721-3" ],
    "2016": ["20170721-2", "20170721-3" ],
    "2017": ["20180411",   "20170721-3" ],
    "2018": ["20190430",   "20170721-3" ]
}
   
DBtags_Kplus_refMC = {
    "2011": ["20160614-1", "20160318-1"],
    "2012": ["20160321-2", "20150928"  ],
    "2015": ["20161124",   "20170721-3"],
    "2016": ["20170721-2", "20170721-3"],
    "2017": ["", ""],
    "2018": ["", ""]
}
   
DBtags_Kplus_PHSP = {
    "2011": ["20160614-1", "20170721-1"],
    "2012": ["20160321-2", "20170721-2"],
    "2015": ["20161124",   "20170721-3"],
    "2016": ["20170721-2", "20170721-3"],
    "2017": ["20180411",   "20170721-3"],
    "2018": ["20190128",   "20170721-3"]
}

DBtags_Kshort_sigMC = {
    "2011": ["20160614-1", "20160318-1"],
    "2012": ["20160321-2", "20150928"  ],
    "2015": ["20161124",   "20150724"  ],
    "2016": ["20170721-2", "20170721-3"],
    "2017": ["20180411",   "20170721-3"],
    "2018": ["", ""]
}
   
DBtags_Kshort_refMC = {
    "2011": ["20160614-1", "20160318-1"],
    "2012": ["20160321-2", "20150928"  ],
    "2015": ["", ""],
    "2016": ["20170721-2", "20170721-3"],
    "2017": ["", ""],
    "2018": ["", ""]
}
   
DBtags_Kshort_PHSP = {
    "2011": ["20160614-1", "20170721-1"],
    "2012": ["20160321-2", "20170721-2"],
    "2015": ["20161124",   "20170721-3"],
    "2016": ["20170721-2", "20170721-3"],
    "2017": ["20180411",   "20170721-3"],
    "2018": ["20190128",   "20170721-3"]
}

polarity_tag = "-vc-md100" if(down) else "-vc-mu100"

used_dict = {}
if (data):
    dv.CondDBtag = DBtags_data[year][0]
    dv.DDDBtag = DBtags_data[year][1]
else:    
    if (KplusPi0Resolved):
        if (reference): used_dict = DBtags_Kplus_refMC
        elif (phsp): used_dict = DBtags_Kplus_PHSP
        else: used_dict = DBtags_Kplus_sigMC
    if (KShortPiPlus):
        if (reference): used_dict = DBtags_Kshort_refMC
        elif (phsp): used_dict = DBtags_Kshort_PHSP
        else: used_dict = DBtags_Kshort_sigMC

    dv.CondDBtag = "sim-"+used_dict[year][0]+polarity_tag
    dv.DDDBtag   = "dddb-"+used_dict[year][1]