#
# B->DX selection for stripping (consider as an example). 
# This is an example of a simple wrapper over user's own algorithm. 
# No StrippingMember functionality is used here. 
#
# B->DX,
# where X=pi,K,K*
# and D=D+,D0,Ds.
# The allowed decay modes of the intermediates are
# K* -> Kpi, pipi (to include the pho)
# D+ -> hhh (h=pi,K all combinations except KKK),
# Ds -> hhh (h=pi,K all combinations except KKK), 
# D0 -> hh (h=pi,K all combinations).
# D0 -> hhhh (K3pi FAV and DCS, and KKpipi are considered), 
# D0 -> Kpipi0
# D0 -> Kshh (h=pi,K all combinations)
# 
# This is a loosed version of the B ->DX stripping to be used for early data
#

# The selection is tuned on the L0xHLT1 sample of minbias, achieving 
# a reduction factor of ~3/1000 for MC09 mbias data. 
# The signal efficiency for the various channels ranges between 75 and 90%.
#
# Authors: J. Nardulli
# 18-03-2010
#

from os import environ
from Gaudi.Configuration import *
from Configurables import DecayTreeTuple
from Configurables import CombineParticles, FilterDesktop
from Configurables import TupleToolMCBackgroundInfo, TupleToolMCTruth, BackgroundCategory
from StrippingConf.StrippingLine import StrippingLine

########################################################################
#Define the shared cuts
Dcut = "((M > 1768.*MeV) & (M < 2068.*MeV) & (VFASPF(VCHI2/VDOF)<12) & (BPVVDCHI2 > 4)) "
#Dcut = "((M > 1768.*MeV) & (M < 2068.*MeV) & (VFASPF(VCHI2/VDOF)<6) & (BPVVDCHI2 > 100) & (BPVVD>1*mm) & (BPVVDR>0.2*mm)) "

Dptcut = "PT>1*GeV"
Dcombcut = "(A" + Dptcut + ")"   
KStarMothercut = "((VFASPF(VCHI2/VDOF) < 15) & (MIPCHI2DV(PRIMARY) > 4.))"

Kstarcut = "(M > 0)"

Bcut = "((VFASPF(VCHI2/VDOF)<15)  & (BPVIPCHI2() < 49) & (BPVVDCHI2 > 4) & (BPVDIRA > 0.9995))"
#Bcut = "((VFASPF(VCHI2/VDOF)<12)  & (BPVIPCHI2() < 49) & (BPVVDCHI2 > 4) & (BPVIP() < 0.07*mm) & (BPVVD > 1*mm) & (BPVDIRA > 0.9995))"
Bcombcut = "((AM > 4800.*MeV) & (AM < 5900.*MeV))"

Daughtercut = "((TRCHI2DOF<10) & (PT >250*MeV) & (P > 2*GeV) & (MIPCHI2DV(PRIMARY) > 4.))"
Bachelorcut = "((TRCHI2DOF<10) & (PT >350*MeV) & (P > 2*GeV) & (MIPCHI2DV(PRIMARY) > 4.))"

########################################################################
#Now the various selections etc.
########################################################################
# the K*
StrippingSelKstarLoose = CombineParticles("StrippingSelKstarLoose")
StrippingSelKstarLoose.InputLocations = ["StdNoPIDsPions","StdNoPIDsKaons"]
StrippingSelKstarLoose.DecayDescriptors = ["[K*(892)0 -> K+ pi-]cc","[K*(892)0 -> pi+ pi-]cc"]
StrippingSelKstarLoose.DaughtersCuts =  { "K+"        : Daughtercut,
                                          "pi+"        : Daughtercut}
StrippingSelKstarLoose.CombinationCut = Dcombcut
StrippingSelKstarLoose.MotherCut = KStarMothercut

########################################################################
#The bachelor
StrippingBachelorLoose = FilterDesktop("StrippingBachelorLoose")
StrippingBachelorLoose.InputLocations  = [ "StdNoPIDsPions","StdNoPIDsKaons" ]
StrippingBachelorLoose.Code = Bachelorcut
########################################################################
SubXLoose = GaudiSequencer("SubXLoose")
SubXLoose.ModeOR = 1
SubXLoose.IgnoreFilterPassed = True
SubXLoose.Members += [StrippingSelKstarLoose]
SubXLoose.Members += [StrippingBachelorLoose]

#######################################################################
# All possible D decays starting from the D->hh
StrippingSelD0Loose = CombineParticles("StrippingSelD0Loose")
StrippingSelD0Loose.InputLocations = ["StdNoPIDsKaons", "StdNoPIDsPions"]
StrippingSelD0Loose.DecayDescriptors = ["[D0 -> K- pi+]cc","[D0 -> pi- pi+]cc","[D0 -> K- K+]cc","[D0 -> K+ pi-]cc"]
StrippingSelD0Loose.DaughtersCuts =  { "K+"        : Daughtercut,"pi+"        : Daughtercut }
StrippingSelD0Loose.CombinationCut = Dcombcut
StrippingSelD0Loose.MotherCut = Dcut

########################################################################
B2DXStrippingSelLoose = CombineParticles("B2DXStrippingSelLoose")
B2DXStrippingSelLoose.InputLocations  = ["StrippingSelD0Loose",
                                         "StrippingSelKstarLoose",
                                         "StrippingBachelorLoose"]

B2DXStrippingSelLoose.DecayDescriptors  = ["[B+ -> D~0 K+]cc",
                                           "[B+ -> D~0 pi+]cc",
                                           "[B+ -> D0 K+]cc",
                                           "[B+ -> D0 pi+]cc",
                                           "[B0 -> D~0 K*(892)0]cc",
                                           "[B0 -> D0 K*(892)0]cc"]
                                         

B2DXStrippingSelLoose.DaughtersCuts = { "D+" : Dcut , "K*(892)0" : Kstarcut ,  "D0" : Dcut , "pi+" : Bachelorcut, "K+" : Bachelorcut}
B2DXStrippingSelLoose.CombinationCut = Bcombcut
B2DXStrippingSelLoose.MotherCut = Bcut


########################################################################
B2DXStrippingSequenceLoose = GaudiSequencer("B2DXStrippingSequenceLoose")
B2DXStrippingSequenceLoose.Members += [StrippingSelD0Loose]
B2DXStrippingSequenceLoose.Members += [SubXLoose]
B2DXStrippingSequenceLoose.Members += [B2DXStrippingSelLoose]
########################################################################




#Plots
AnalysisDecayTreeTuple = DecayTreeTuple("AnalysisDecayTreeTuple")
BachelorDecayTreeTuple = DecayTreeTuple("BachelorDecayTreeTuple")
KstarDecayTreeTuple = DecayTreeTuple("KstarDecayTreeTuple")
BachelorOrKstarDecayTreeTuple = GaudiSequencer("BachelorOrKstarDecayTreeTuple")
BachelorOrKstarDecayTreeTuple.ModeOR = 1
BachelorOrKstarDecayTreeTuple.Members = [BachelorDecayTreeTuple,KstarDecayTreeTuple]
B2DXStrippingSequenceLoose.Members += [AnalysisDecayTreeTuple]
B2DXStrippingSequenceLoose.Members += [BachelorOrKstarDecayTreeTuple]

AnalysisDecayTreeTuple = DecayTreeTuple("AnalysisDecayTreeTuple")
AnalysisDecayTreeTuple.InputLocations = ["Phys/B2DXStrippingSelLoose"]
AnalysisDecayTreeTuple.Decay = "[[B0]cc ->  (^D~0 => {^K+ ^pi-, ^K- ^pi+,^K+ ^K-,^pi+ ^pi-}) ? ]cc"
AnalysisDecayTreeTuple.Branches = {
    "D" : "[[B0]cc ->  (^D~0 => {K+ pi-, K- pi+, K+ K-, pi+ pi-}) {K+, pi+}]cc"
    ,"Positive_D_daughter" : "[[B0]cc ->  (D~0 => {^K+ pi-, K- ^pi+, ^K+ K-, ^pi+ pi-, ^K+ ^pi+,K- pi-}) ? ]cc"
    ,"Negative_D_daughter" : "[[B0]cc ->  (D~0 => {K+ ^pi-, ^K- pi+, K+ ^K-, pi+ ^pi-, K+ pi+,^K- ^pi-}) ? ]cc"
    ,"B" : "[B0]cc : [[B0]cc ->  (D~0 => {K+ pi-, K- pi+,K+ K-,pi+ pi-, K+ pi+, K- pi-}) ? ]cc"
    }

BachelorDecayTreeTuple = DecayTreeTuple("BachelorDecayTreeTuple")
BachelorDecayTreeTuple.InputLocations = ["Phys/B2DXStrippingSelLoose"]
BachelorDecayTreeTuple.Decay = "[[B+]cc ->  (^D~0 => {^K+ ^pi-, ^K- ^pi+,^K+ ^K-,^pi+ ^pi-,^K+ ^pi+,^K- ^pi-}) {^K+, ^pi+}]cc"
BachelorDecayTreeTuple.Branches = {
    "D" : "[[B+]cc ->  (^D~0 => {K+ pi-, K- pi+, K+ K-, pi+ pi-,K+ pi+,K- pi-}) {K+, pi+}]cc"
    ,"Positive_D_Daughter" : "[[B+]cc ->  (D~0 => {^K+ pi-, K- ^pi+, ^K+ K-, ^pi+ pi-,^K+ ^pi+,K- pi-}) {K+, pi+}]cc"
    ,"Negative_D_Daughter" : "[[B+]cc ->  (D~0 => {K+ ^pi-, ^K- pi+, K+ ^K-, pi+ ^pi-,K+ pi+,^K- ^pi-}) {K+, pi+}]cc"
    ,"Bachelor" : "[[B+]cc ->  (D~0 => {K+ pi-, K- pi+, K+ K-, pi+ pi-,K+ pi+,K- pi-}) {^K+, ^pi+}]cc"
    ,"B" : "[B+]cc : [[B+]cc ->  (D~0 => {K+ pi-, K- pi+, K+ K-, pi+ pi-,K+ pi+,K- pi-}) {K+, pi+}]cc"
    }

KstarDecayTreeTuple = DecayTreeTuple("KstarDecayTreeTuple")
KstarDecayTreeTuple.InputLocations = ["Phys/B2DXStrippingSelLoose"]
KstarDecayTreeTuple.Decay = "[[B0]cc -> (^D~0 => {^K+ ^pi-, ^K- ^pi+,^K+ ^K-,^pi+ ^pi-,^K+ ^pi+,^K- ^pi-}) (^K*(892)0 => ^K+ ^pi-) ]cc"
KstarDecayTreeTuple.Branches = {
    "KStar" : "[[B0]cc ->  (D~0 => {K+ pi-, K- pi+, K+ K-, pi+ pi-,K+ pi+,K- pi-}) (^K*(892)0 => K+ pi-) ]cc"
    ,"D" : "[[B0]cc ->  (^D~0 => {K+ pi-, K- pi+, K+ K-, pi+ pi-,K+ pi+,K- pi-}) (K*(892)0 => K+ pi-) ]cc"
    ,"Positive_D_Daughter" : "[[B0]cc ->  (D~0 => {^K+ pi-, K- ^pi+, ^K+ K-, ^pi+ pi-,^K+ ^pi+,K- pi-}) (K*(892)0 => K+ pi-) ]cc"
    ,"Negative_D_Daughter" : "[[B0]cc ->  (D~0 => {K+ ^pi-, ^K- pi+, K+ ^K-, pi+ ^pi-,K+ pi+,^K- ^pi-}) (K*(892)0 => K+ pi-) ]cc"
    ,"K_from_kstar" : "[[B0]cc ->  (D~0 => {K+ pi-, K- pi+, K+ K-, pi+ pi-,K+ pi+,K- pi-}) (K*(892)0 => ^K+ pi-) ]cc"
    ,"Pi_from_kstar" : "[[B0]cc ->  (D~0 => {K+ pi-, K- pi+, K+ K-, pi+ pi-,K+ pi+,K- pi-}) (K*(892)0 => K+ ^pi-) ]cc"
    ,"B" : "[B0]cc : [[B0]cc ->  (D~0 => {K+ pi-, K- pi+,K+ K-,pi+ pi-}) (K*(892)0 => K+ pi-) ]cc"
    }


KstarDecayTreeTuple.ToolList +=  [
    "TupleToolGeometry"
    , "TupleToolAngles"
    , "TupleToolKinematic"
    , "TupleToolPrimaries"
    , "TupleToolPid"
    , "TupleToolEventInfo"
    , "TupleToolTrackInfo"
     ]


BachelorDecayTreeTuple.ToolList +=  [
    "TupleToolGeometry"
    , "TupleToolAngles"
    , "TupleToolKinematic"
    , "TupleToolPrimaries"
    , "TupleToolPid"
    , "TupleToolEventInfo"
    , "TupleToolTrackInfo"
    ]

AnalysisDecayTreeTuple.ToolList +=  [
    "TupleToolGeometry"
    , "TupleToolAngles"
    , "TupleToolKinematic"
    , "TupleToolPrimaries"
    , "TupleToolPid"
    , "TupleToolEventInfo"
    , "TupleToolTrackInfo"
    , "LoKi::Hybrid::TupleTool/LoKi_All" 
    ]