import GaudiPython
from Gaudi.Configuration import *
from GaudiPython.Bindings import gbl, AppMgr, Helper
from Configurables import GaudiSequencer, HltSelReportsMaker, OutputStream, DaVinci
from Configurables import PrintDecayTree
from os import environ
from Configurables import DecayTreeTuple, EventTuple
from Configurables import CombineParticles, FilterDesktop
from Configurables import TupleToolMCBackgroundInfo, TupleToolMCTruth, TupleToolTrigger, TupleToolTISTOS, TupleToolDecay, BackgroundCategory
from StrippingConf.StrippingLine import StrippingLine



DV = DaVinci()
DV.EvtMax = -1
DV.PrintFreq = 100
DV.SkipEvents = 0
DV.DataType = "2010"
DV.Simulation = True

DV.L0 = False #want this to be False for data; True will overwrite decisions. the same applies to the next few settings
DV.Hlt = False
DV.ReplaceL0BanksWithEmulated = False #reruns the L0. need more info
#DV.HltThresholdSettings = 'Physics_25Vis_25L0_2Hlt1_2Hlt2_May10' #only makes sense if re-running trigger

ApplicationMgr().ExtSvc +=  [ "NTupleSvc" ]
NTupleSvc().Output =  [ "FILE1 DATAFILE='/home/hep/uoh35620/stuff/batch-generate/data/up/all/B2Dpi_Up_6_7_2010.direct.root' TYP='ROOT' OPT='NEW'" ]
#NTupleSvc().Output =  [ "FILE1 DATAFILE='castor:/castor/cern.ch/user/j/jnardull/RealDataRootFiles/B2DX_UP_FromHadron_6_7_2010.root' TYP='ROOT' OPT='NEW'" ]

from Configurables import CondDB
CondDB().IgnoreHeartBeat = True

#stripping stream location
location = '/Event/Bhadron/Phys/B2DXStrippingSelChi2Loose'

from PhysSelPython.Wrappers import AutomaticData, Selection, SelectionSequence
# Get the Candidates from the DST.
# Treat particles already on the DST as data-on-demand, but give the full path.
B2DXSel = AutomaticData(Location = location)
# Filter the Candidate. Replace 'ALL' by some cuts.

B2DXFilter = FilterDesktop('B2DXFilter')
B2DXFilter.Code = "ALL" #replace this if wanted

# make a Selection
B2DXFilterSel = Selection(name = 'B2DXFilterSel',
                          Algorithm = B2DXFilter,
                          RequiredSelections = [ B2DXSel])

# build the SelectionSequence
B2DXSeq = SelectionSequence('B2DXSeq', TopSelection = B2DXFilterSel)
DV.appendToMainSequence( [ B2DXSeq.sequence() ] )

########################################################################
# DecayTreeTuples
########################################################################

#
# Initialisation
#

TupleSequence = GaudiSequencer ("TupleSequencer")

AnalysisDecayTreeTuple = DecayTreeTuple("AnalysisDecayTreeTuple")
BachelorDecayTreeTuple = DecayTreeTuple("BachelorDecayTreeTuple")
KstarDecayTreeTuple = DecayTreeTuple("KstarDecayTreeTuple")

BachelorOrKstarDecayTreeTuple = GaudiSequencer("BachelorOrKstarDecayTreeTuple")
BachelorOrKstarDecayTreeTuple.ModeOR = 1
BachelorOrKstarDecayTreeTuple.Members = [BachelorDecayTreeTuple,KstarDecayTreeTuple]

TupleSequence.Members += [AnalysisDecayTreeTuple]
TupleSequence.Members += [BachelorOrKstarDecayTreeTuple]


#Plots
AnalysisDecayTreeTuple.InputLocations = ["B2DXFilterSel"]
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 = ["B2DXFilterSel"]
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 = ["B2DXFilterSel"]
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" 
    ]

########################################################################
# Add sequence to DaVinci
########################################################################

DV.UserAlgorithms += [TupleSequence]