# -*- coding: utf-8 -*-
# @Author: Marcin Chrzaszcz
# @Date:   2016-08-27 
# @Last Modified by:   Anna Lupato
# @Last Modified time: 2016-09-28 
########################################################################
from os import environ
import math


# use CommonParticlesArchive
from CommonParticlesArchive import CommonParticlesArchiveConf
CommonParticlesArchiveConf().redirect('stripping28')


from Gaudi.Configuration import *
#from GaudiPython import AppMgr
from Configurables import DecayTreeTuple, BTaggingTool, SubstitutePID, TrackScaleState
from Configurables import CombineParticles, FilterDesktop, TupleToolDecayTreeFitter, TupleToolDecay
from Configurables import TupleToolTrigger, BackgroundCategory, TupleToolTISTOS, TupleToolRecoStats, TupleToolTagging
from Configurables import DaVinci, HltSelReportsDecoder, HltVertexReportsDecoder, HltDecReportsDecoder, LoKi__Hybrid__TupleTool, TupleToolJets
from DecayTreeTuple.Configuration import *
from StrippingConf.Configuration import StrippingConf, StrippingStream
from StrippingSettings.Utils import strippingConfiguration
from StrippingArchive.Utils import buildStreams, cloneLinesFromStream
from Configurables import CombineParticles
from StrippingArchive import strippingArchive
from Configurables import FilterInTrees
from PhysSelPython.Wrappers import Selection, DataOnDemand, AutomaticData, SelectionSequence
msgSvc = getConfigurable("MessageSvc")
msgSvc.Format = "% F%60W%S%7W%R%T %0W%M"

#try and match it with what is in main_opt.py

trigger_list = [
    'L0MuonDecision',
    'L0ElectronDecision',
    'L0HadronDecision',
    'Hlt1TrackMVADecision',
    'Hlt1TrackMVALooseDecision',
    'Hlt2RareCharmD2PiMuMuOSDecision',
    'Hlt2RareCharmD2PiEEOSDecision',
]

#StdPions = DataOnDemand(Location = "Phys/StdAllLoosePions/Particles")

particle_list = []

#no need for a WS decay so get rid of it


decay_OSee = "[D_s+ -> ^(phi(1020) -> ^e+ ^e-) ^pi+]CC"

#decay_OSmumu = "[D+ -> ^pi+ ^mu+ ^mu-]CC"

#for mc /Event/
# find my own decay_tree_location using main_opt and make_decay_tree_tuples.py
decay_tree_location_Dplus_electrons = 'Phys/D2XMuMuSS_Pi_EE_OSLine/Particles'
#decay_tree_location_Dplus_muons = 'Phys/D2XMuMuSS_PiOSLine/Particles'

#ACCESSING TO THE DECAY DAUGHTERS
StrippingSel = AutomaticData(decay_tree_location_Dplus_electrons)


electron_from_Ds = FilterInTrees('electron_from_Ds', Code="('e-' == ABSID)")
electron_from_Ds_sel = Selection("electron_from_Ds_sel",
                              Algorithm=electron_from_Ds,
                              RequiredSelections=[StrippingSel])

electron_Filter = FilterDesktop('electron_Filter', Code="PPINFO(PROBNNe)>0.05")#
electron_Filter_sel = Selection("Filtered_electron_sel",
                              Algorithm=electron_Filter,
                              RequiredSelections=[electron_from_Ds_sel])


pion_from_Ds = FilterInTrees('pion_from_Ds', Code="('pi+' == ABSID)")
pion_from_Ds_sel = Selection("pion_from_Ds_sel",
                              Algorithm=pion_from_Ds,
                              RequiredSelections=[StrippingSel])

pion_Filter = FilterDesktop('pionFilter', Code="PPINFO(PROBNNpi)>0.2")
pion_Filter_sel = Selection("Filtered_pion_sel",
                              Algorithm=pion_Filter,
                              RequiredSelections=[pion_from_Ds_sel])


#############################################################################
#RECOMBINE THE LEPTONS TO A PHI


phi = CombineParticles("phi")

phi.DecayDescriptors = ["phi(1020) -> e+ e-"] #no cc here
#phi.MotherCut = "ALL" #"(VFASPF(VCHI2/VDOF) < 6.0) " #leaving the chisquared cut 
phi.MotherCut = "BPVDIRA>0.999" #in order to cut uninteresting events from the data BPV = Best Primary Vertex


phi_sel = Selection("phi_sel",
                              Algorithm=phi,
                              RequiredSelections=[electron_Filter_sel])

############
#Remake the Ds


Ds = CombineParticles("D_s+")
Ds.DecayDescriptors = ["[D_s+ -> phi(1020) pi+]cc"]
Ds.MotherCut = "(M<6500) & (VFASPF(VCHI2/VDOF) < 6.0) & (BPVDIRA>0.999)"
Ds_sel = Selection("Ds_sel",
                              Algorithm=Ds,
                              RequiredSelections=[phi_sel, pion_Filter_sel]) #the order matters!! It makes sense to put the rarer decay first in order to 
                                                                            #save time -> put the most distinctive decay first

Ds_phi_pi_sel_seq = SelectionSequence('Ds_phi_pi_selection_seq',
                                           TopSelection=Ds_sel
                                           )


#############################################################################
# Now WS:
# Lb_SL_WS = CombineParticles("Lb_SL_WS")
# Lb_SL_WS.DecayDescriptors = ["[Lambda_b0 -> Lambda_c(2595)+ mu+]cc"]
# Lb_SL_WS.MotherCut = "(M<6500) & (VFASPF(VCHI2/VDOF) < 6.0) & (BPVDIRA>0.999)"



# Creating the DecayTreeTuple objects
decay_trees = []

#replace thedecays you're intereste in here
Ds_OfflineTree = DecayTreeTuple("Ds_OfflineTree")
Ds_OfflineTree.Inputs = Ds_phi_pi_sel_seq.outputLocations()
Ds_OfflineTree.Decay = decay_OSee
Ds_OfflineTree.addBranches({
    "Ds":"[D_s+ -> phi(1020) pi+]CC",

    "phi":"[D_s+ -> ^phi(1020) pi+]CC",
    "pi":"[D_s+ -> phi(1020) ^pi+]CC",

    "e_plus":"[D+ -> (phi(1020) -> ^e+ e-) pi+]CC",
    "e_minus":"[D+ -> (phi(1020) -> e+ ^e-) pi+]CC",


})

Ds_OfflineTree.Ds.addTupleTool('TupleToolDecayTreeFitter/ConsD')

#Ds_OfflineTree.Ds.ConsD.constrainToOriginVertex = True
Ds_OfflineTree.Ds.ConsD.Verbose = True
Ds_OfflineTree.Ds.ConsD.daughtersToConstrain = ['phi(1020)']
Ds_OfflineTree.Ds.ConsD.UpdateDaughters = True

decay_trees.append(Ds_OfflineTree)

# Lb_SL_OfflineTree_WS = DecayTreeTuple("Lb_SL_OfflineTree_WS")
# Lb_SL_OfflineTree_WS.Inputs = ["Phys/Lb_SL_WS/Particles"]
# Lb_SL_OfflineTree_WS.Decay = decay0_WS
# Lb_SL_OfflineTree_WS.addBranches({
#     "Lambda_b0": "[Lambda_b0 -> (Lambda_c(2595)+ -> (Lambda_c+ -> K- p+ pi+ ) pi+ pi-) mu+ ]CC",
#     "Lambda_c_STAR": "[Lambda_b0 -> ^(Lambda_c(2595)+ -> (Lambda_c+ -> K- p+ pi+ ) pi+ pi-) mu+ ]CC",
#     "mu": "[Lambda_b0 -> (Lambda_c(2595)+ -> (Lambda_c+ -> K- p+ pi+ ) pi+ pi-) ^mu+ ]CC",
#     "Lambda_c": "[Lambda_b0 -> (Lambda_c(2595)+ -> ^(Lambda_c+ -> K- p+ pi+ ) pi+ pi-) mu+ ]CC",
#     "pi_plus": "[Lambda_b0 -> (Lambda_c(2595)+ -> (Lambda_c+ -> K- p+ pi+ ) ^pi+ pi-) mu+ ]CC",
#     "pi_minus": "[Lambda_b0 -> (Lambda_c(2595)+ -> (Lambda_c+ -> K- p+ pi+ ) pi+ ^pi-) mu+ ]CC",
#     "p": "[Lambda_b0 -> (Lambda_c(2595)+ -> (Lambda_c+ -> K- ^p+ pi+ ) pi+ pi-) mu+ ]CC",
#     "pi_Lc": "[Lambda_b0 -> (Lambda_c(2595)+ -> (Lambda_c+ -> K- p+ ^pi+ ) pi+ pi-) mu+ ]CC",
#     "K": "[Lambda_b0 -> (Lambda_c(2595)+ -> (Lambda_c+ -> ^K- p+ pi+ ) pi+ pi-) mu+ ]CC"
# })
# decay_trees.append(Lb_SL_OfflineTree_WS)
# Lambda_b from wrong sign Lambda_c* tree

# Lambda_b from wrong sign Lambda_c* tree

# Double charm tree




from Configurables import TupleToolNeutrinoReco
from Configurables import TupleToolKinematic
from Configurables import TupleToolGeometry
from Configurables import TupleToolTrackIsolation
from Configurables import TupleToolANNPID
# from Configurables import Rich__Rec__DetailedTrSegMakerFromRecoTracks
# from Configurables import TupleToolRICHPid
#from Configurables import TupleToolApplyIsolation

#from Configurables import TupleToolSLTools

for DTT in decay_trees:
    DTT.ToolList = [
        "TupleToolKinematic",
        "TupleToolGeometry",
        "TupleToolTrackInfo",
        "TupleToolEventInfo",
        "TupleToolRecoStats",
        "TupleToolANNPID",
        "TupleToolPid",
        # "TupleToolMCTruth",
        # 'TupleToolTrigger',
        # 'TupleToolTISTOS',
        # 'TupleToolRICHPid',
    ]
    # DTT.pi_plus.addTool(TupleToolANNPID,name='TupleToolANNPID')
    # DTT.pi_plus.ToolList+=['TupleToolANNPID']
    # DTT.pi_plus.TupleToolANNPID.ANNPIDTunes = ['MC15TuneV1','MC15TuneDNNV1','MC15TuneFLAT4dV1','MC15TuneCatBoostV1']
    # print 'PID tunes',DTT.pi_plus.TupleToolANNPID.ANNPIDTunes
    DTT.Ds.addTool(TupleToolTrigger, name='TupleToolTrigger')
    DTT.Ds.TupleToolTrigger.Verbose = True
    DTT.Ds.TupleToolTrigger.TriggerList = trigger_list
    DTT.Ds.addTool(TupleToolTISTOS, name='TupleToolTISTOS')
    DTT.Ds.TupleToolTISTOS.Verbose = True
    DTT.Ds.TupleToolTISTOS.TriggerList = trigger_list
    DTT.Ds.ToolList += ["TupleToolTISTOS"]
    #DTT.Lambda_b0.addTool(TupleToolSLTools, name='TupleToolSLTools')
    #DTT.Lambda_b0.ToolList += ["TupleToolSLTools"]

# Isolation from the R(D*) SL analysis
#for DTT in decay_trees:

    #toolIso = DTT.Lambda_b0.addTupleTool(
    #    TupleToolApplyIsolation, name="TupleToolApplyIsolation")
    #toolIso.WeightsFile = "weights.xml"
    #DTT.Lambda_b0.ToolList += ["TupleToolApplyIsolation/TupleToolApplyIsolation"]

# Isolation from the R(D*) SL analysis (adding VELO tracks)
# from Configurables import ChargedProtoParticleMaker
# veloprotos = ChargedProtoParticleMaker("myProtoPMaker")
# veloprotos.Inputs = ["Rec/Track/Best"]
# veloprotos.Output = "Rec/ProtoP/myProtoPMaker/ProtoParticles"
# DaVinci().appendToMainSequence([veloprotos])
from Configurables import ProtoParticleCALOFilter, CombinedParticleMaker, NoPIDsParticleMaker
from CommonParticles.Utils import *
# algorithm = NoPIDsParticleMaker('StdNoPIDsVeloPions', Particle='pion',)
# algorithm.Input = "Rec/ProtoP/myProtoPMaker/ProtoParticles"
# selector = trackSelector(algorithm, trackTypes=['Velo'])
# locations = updateDoD(algorithm)
# DaVinci().appendToMainSequence([algorithm])
from PhysConf.Filters import LoKi_Filters
# fltrs = LoKi_Filters(
#   STRIP_Code = "HLT_PASS('Strippingb2LcMuXB2DMuForTauMuLineDecision')")
##################################################
from Configurables import DaVinci
# DaVinci().EventPreFilters = fltrs.filters('Filters')
DaVinci().EvtMax = -1  # 200                          # Number of events
DaVinci().SkipEvents = 0                       # Events to skip
DaVinci().PrintFreq = 10000
DaVinci().Simulation = False
DaVinci().Lumi = True
DaVinci().RootInTES = '/Event/Charm'
DaVinci().InputType = 'MDST'
#DaVinci().appendToMainSequence([muons_from_Lc_sel_seq.sequence(),
#                                muons_from_Lc_sel_WC_seq.sequence()])
DaVinci().appendToMainSequence([Ds_phi_pi_sel_seq.sequence()])
DaVinci().appendToMainSequence(decay_trees)

#DaVinci().UserAlgorithms.extend(decay_trees)

DaVinci().TupleFile = "Ds_eepi.root"
DaVinci().Input = ['PFN:/disk/gangadir/davide/workspace/davide/local_data/00069601_00000009_1.charm.mdst']


DaVinci().DataType  = "2016"