from Gaudi.Configuration import *
from Configurables import DaVinci, MicroDSTWriter
from PhysConf.MicroDST import uDstConf
from PhysSelPython.Wrappers import DataOnDemand
from PhysSelPython.Wrappers import Selection, SelectionSequence
# make a FilterDesktop
from Configurables import FilterDesktop

from DecayTreeTuple import *
from DecayTreeTuple.Configuration import *
from Configurables import LoKi__Hybrid__EvtTupleTool as LokiEventTool
from Configurables import LoKi__Hybrid__TupleTool as LokiTool
from Configurables import TupleToolTwoParticleMatching as MatcherTool

from Configurables import StoreExplorerAlg
from Production.Configuration import ProbNNRecalibrator as ProbNNcalib
from Configurables import ApplySWeights
from Configurables import BinnedPrescaler
from PhysConf.MicroDST import uDstConf
## from Configurables import DaVinci
from copy import copy


##################################################
###      C O N F I G U R A T I O N   1     #######
###. . . . . . . . . . . . . . . . . . . . #######
###        Variables by branch type        #######
##################################################
###  List here variables for each TYPE of branch.
### Branches are: "HEAD", the head of the decay chain
###               "INTERMEDIATE", produced and decayed
###               "TRACK", charged basic particle
###               "NEUTRAL", photons or pi0
##################################################

LokiVarsByType = {
  "HEAD" : {
    "ENDVERTEX_CHI2"   : "VFASPF(VCHI2)"
    , "ENDVERTEX_NDOF" : "VFASPF(VDOF)"
    , "IPCHI2" : "BPVIPCHI2()"
    , "IP"     : "BPVIP()"
    ,"BPVLTCHI2" : "BPVLTCHI2()"
    , "BPVDLS"   : "BPVDLS"
    ,"Mass" : "M"
  },
  
  "INTERMEDIATE" : {
    "ENDVERTEX_CHI2"   : "VFASPF(VCHI2)"
    , "ENDVERTEX_NDOF" : "VFASPF(VDOF)"
   , "IPCHI2" : "BPVIPCHI2()"
    , "IP"     : "BPVIP()"
    , "BPVDLS"   : "BPVDLS"
    ,"BPVLTCHI2" : "BPVLTCHI2()"
    ,"Mass" : "M"
  }, 

  "TRACK" : {
    "LoKi_PT"      : "PT"     ## These variables are also saved from the PIDCalib
    , "LoKi_P"     : "P"      ## tuple tool, so they are a waste of disk space
    , "LoKi_ETA"   : "ETA"    ## however the correct thing would be to keep them
    , "Loki_PIDK"  : "PIDK"   ## here and removing them from PIDCalib.
    , "Loki_PIDmu" : "PIDmu"  ##  ... to be done ...
    , "Loki_PIDp"  : "PIDp"   ## 
    , "sWeight"    : "WEIGHT"
    , "Loki_MINIPCHI2": "BPVIPCHI2()"
    , "RichDLLe"      : "PPFUN ( PP_RichDLLe )"
    , "RichDLLpi"     : "PPFUN ( PP_RichDLLpi )"
    , "RichDLLmu"     : "PPFUN ( PP_RichDLLmu )"
    , "RichDLLk"      : "PPFUN ( PP_RichDLLk )"
    , "RichDLLp"      : "PPFUN ( PP_RichDLLp )"
    , "RichDLLbt"     : "PPFUN ( PP_RichDLLbt )"
    , "MuonMuLL"      : "PPFUN ( PP_MuonMuLL )"
    , "MuonBgLL"      : "PPFUN ( PP_MuonBkgLL )"
    , "Charge"        : "switch ( ID > 0, +1, -1 )"
    , "MuonUnbiased"  : "switch ("\
                        "(TIS('L0.*Decision', 'L0TriggerTisTos')) & "\
                        "(TIS('Hlt1(?!ODIN)(?!L0)(?!Lumi)(?!Tell1)(?!MB)(?!NZS)(?!Velo)(?!BeamGas)(?!Incident).*Decision', 'Hlt1TriggerTisTos')) &"\
                        "(TIS('Hlt2(?!Forward)(?!DebugEvent)(?!Express)(?!Lumi)(?!Transparent)(?!PassThrough).*Decision', 'Hlt2TriggerTisTos')) , " \
                        "1,0)"
    , "ElectronUnbiased" : "switch ( "\
                        "(TIS('L0ElectronDecision', 'L0TriggerTisTos')) & "\
                        "(TIS('Hlt1PhysDecision', 'Hlt1TriggerTisTos')) &"\
                        "(TIS('Hlt2PhysDecision', 'Hlt2TriggerTisTos')) , " \
                        "1,0)"
    , "TRCHI2NDOF" :  "TRCHI2DOF"
    , "TRACK_GHOSTPROB" : "TRGHP"
  },
  
  "NEUTRAL" : {
  }
}


EventInfo = {
  "nPVs"              : "RECSUMMARY( LHCb.RecSummary.nPVs              , -9999)"
#  , "nLongTracks"       : "RECSUMMARY( LHCb.RecSummary.nLongTracks       , -9999)"
#  , "nDownstreamTracks" : "RECSUMMARY( LHCb.RecSummary.nDownstreamTracks , -9999)"
#  , "nUpstreamTracks"   : "RECSUMMARY( LHCb.RecSummary.nUpstreamTracks   , -9999)"
#  , "nVeloTracks"       : "RECSUMMARY( LHCb.RecSummary.nVeloTracks       , -9999)"
#  , "nTTracks"          : "RECSUMMARY( LHCb.RecSummary.nTTracks          , -9999)"
#  , "nBackTracks"       : "RECSUMMARY( LHCb.RecSummary.nBackTracks       , -9999)"
  , "nTracks"           : "RECSUMMARY( LHCb.RecSummary.nTracks           , -9999)"
  , "nRich1Hits"        : "RECSUMMARY( LHCb.RecSummary.nRich1Hits        , -9999)"
  , "nRich2Hits"        : "RECSUMMARY( LHCb.RecSummary.nRich2Hits        , -9999)"
#  , "nVeloClusters"     : "RECSUMMARY( LHCb.RecSummary.nVeloClusters     , -9999)"
#  , "nITClusters"       : "RECSUMMARY( LHCb.RecSummary.nITClusters       , -9999)"
#  , "nTTClusters"       : "RECSUMMARY( LHCb.RecSummary.nTTClusters       , -9999)"
#  , "nUTClusters"       : "RECSUMMARY( LHCb.RecSummary.nUTClusters       , -9999)"
#  , "nOTClusters"       : "RECSUMMARY( LHCb.RecSummary.nOTClusters       , -9999)"
#  , "nFTClusters"       : "RECSUMMARY( LHCb.RecSummary.nFTClusters       , -9999)"
  , "nSPDhits"          : "RECSUMMARY( LHCb.RecSummary.nSPDhits          , -9999)"
#  , "nMuonCoordsS0"     : "RECSUMMARY( LHCb.RecSummary.nMuonCoordsS0     , -9999)"
#  , "nMuonCoordsS1"     : "RECSUMMARY( LHCb.RecSummary.nMuonCoordsS1     , -9999)"
#  , "nMuonCoordsS2"     : "RECSUMMARY( LHCb.RecSummary.nMuonCoordsS2     , -9999)"
#  , "nMuonCoordsS3"     : "RECSUMMARY( LHCb.RecSummary.nMuonCoordsS3     , -9999)"
#  , "nMuonCoordsS4"     : "RECSUMMARY( LHCb.RecSummary.nMuonCoordsS4     , -9999)"
  , "nMuonTracks"       : "RECSUMMARY( LHCb.RecSummary.nMuonTracks       , -9999)"
}   


##################################################
###      C O N F I G U R A T I O N   2     #######
###. . . . . . . . . . . . . . . . . . . . #######
###        Variables by branch name        #######
##################################################
###  List here variables for each Branch name
###   The purpose of this configuration is to allow 
###   to set a variable for all "Lambda_c+" as long
###   as all Lambda_c+ will be named "Lc"
##################################################

LokiVarsByName = {
  "mu" : {
     "MuonTOS"  :      "switch((TOS('L0MuonDecision', 'L0TriggerTisTos')) & "\
                        "(TOS('Hlt1(TrackAllL0|TrackMuon|SingleMuonHighPT)Decision', 'Hlt1TriggerTisTos')), "\
                        "1,0)"
 ,"MuonProbe" :     "switch((P>3000) & (PT>800) & (BPVIPCHI2()>10) & (TRCHI2DOF<3), 1,0)"
     ,"MuonTag" :     "switch((P>6000) & (PT>1500) & (BPVIPCHI2()>25) & (TRCHI2DOF<3) & (ISMUON), 1,0)"


  },
  
  "e" : {
    "ElectronTOS" : "switch ( "\
                        "(TOS('L0ElectronDecision', 'L0TriggerTisTos')) & "\
                        "(TOS('Hlt1PhysDecision', 'TriggerTisTos')) &"\
                        "(TOS('Hlt2PhysDecision', 'TriggerTisTos')) , " \
                        "1,0)"
    ,"ElectronProbe" :   "switch((P>3000) & (PT>500) & (BPVIPCHI2()>9) , 1,0)"
    ,"ElectronTag" :     "switch((P>6000) & (PT>1500) & (BPVIPCHI2()>9) & (PIDe>5), 1,0)"    

}

}

################################################################################
################################################################################
####                                                                        ####
####                                                                        ####
####              S T O P   H E R E                                         ####
####                                                                        ####
####  99% of use cases won't require you to modify what is below            ####
####                                                                        ####
################################################################################
################################################################################
                     
from math import exp
def parseConfiguration_Run2():
  prescalingSlope = .4
  def cut ( particle, prescale = 1. ):
    ret = "(mcMatch ( '{}' ) ) & (".format ( particle )
    boundaries = [0.0, 0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0]
    for i in range(0, len(boundaries)-1):
      middle = 0.5 * (boundaries [ i ]  + boundaries [ i + 1 ])
      ret += "(scale(in_range ( {}*GeV, PT, {}*GeV), SCALE({}))) | ".format (
          boundaries [ i ], boundaries [ i + 1 ],
          (min(1.,1e-5*exp(middle/prescalingSlope)))*prescale )

    ret += "(scale( PT > 5*GeV, SCALE({}))))".format ( prescale )

    return ret
  


  preambulo = ["from LoKiPhysMC.decorators import *",
               "from LoKiPhysMC.functions import mcMatch",
               "from LoKiNumbers.decorators import *" ,
  ]
  posPionFilter     = FilterDesktop("PosPionFilter",     Code = cut('pi+',  6* 0.025 ) , Preambulo = preambulo)
  negPionFilter     = FilterDesktop("NegPionFilter",     Code = cut('pi-',  6* 0.025 ) , Preambulo = preambulo)
  posKaonFilter     = FilterDesktop("PosKaonFilter",     Code = cut('K+' ,  6* 0.15  ) , Preambulo = preambulo) 
  negKaonFilter     = FilterDesktop("NegKaonFilter",     Code = cut('K-' ,  6* 0.15  ) , Preambulo = preambulo)
  posMuonFilter     = FilterDesktop("PosMuonFilter",     Code = cut('mu+',     1.0   ) , Preambulo = preambulo)
  negMuonFilter     = FilterDesktop("NegMuonFilter",     Code = cut('mu-',     1.0   ) , Preambulo = preambulo)
  posElectronFilter = FilterDesktop("PosElectronFilter", Code = cut('e+' ,  2* 0.5   ) , Preambulo = preambulo)
  negElectronFilter = FilterDesktop("NegElectronFilter", Code = cut('e-' ,  2* 0.5   ) , Preambulo = preambulo)
  posProtonFilter   = FilterDesktop("PosProtonFilter",   Code = cut('p+' ,  2* 0.3   ) , Preambulo = preambulo)
  negProtonFilter   = FilterDesktop("NegProtonFilter",   Code = cut('p~-',  2* 0.35  ) , Preambulo = preambulo)
  
  
  posPionSel     = Selection("PosPionSel",     RequiredSelections = [DataOnDemand("Phys/StdAllNoPIDsPions/Particles")], Algorithm = posPionFilter     )
  negPionSel     = Selection("NegPionSel",     RequiredSelections = [DataOnDemand("Phys/StdAllNoPIDsPions/Particles")], Algorithm = negPionFilter     )
  posKaonSel     = Selection("PosKaonSel",     RequiredSelections = [DataOnDemand("Phys/StdAllNoPIDsKaons/Particles")], Algorithm = posKaonFilter     ) 
  negKaonSel     = Selection("NegKaonSel",     RequiredSelections = [DataOnDemand("Phys/StdAllNoPIDsKaons/Particles")], Algorithm = negKaonFilter     )
  posMuonSel     = Selection("PosMuonSel",     RequiredSelections = [DataOnDemand("Phys/StdAllNoPIDsMuons/Particles")], Algorithm = posMuonFilter     )
  negMuonSel     = Selection("NegMuonSel",     RequiredSelections = [DataOnDemand("Phys/StdAllNoPIDsMuons/Particles")], Algorithm = negMuonFilter     )
  posElectronSel = Selection("PosElectronSel", RequiredSelections = [DataOnDemand("Phys/StdAllNoPIDsElectrons/Particles")], Algorithm = posElectronFilter )
  negElectronSel = Selection("NegElectronSel", RequiredSelections = [DataOnDemand("Phys/StdAllNoPIDsElectrons/Particles")], Algorithm = negElectronFilter )
  posProtonSel   = Selection("PosProtonSel",   RequiredSelections = [DataOnDemand("Phys/StdAllNoPIDsProtons/Particles")], Algorithm = posProtonFilter   )
  negProtonSel   = Selection("NegProtonSel",   RequiredSelections = [DataOnDemand("Phys/StdAllNoPIDsProtons/Particles")], Algorithm = negProtonFilter   )
  
  
  posPionSeq     = SelectionSequence("PosPionSeq",     TopSelection = posPionSel    )
  negPionSeq     = SelectionSequence("NegPionSeq",     TopSelection = negPionSel    )
  posKaonSeq     = SelectionSequence("PosKaonSeq",     TopSelection = posKaonSel    )
  negKaonSeq     = SelectionSequence("NegKaonSeq",     TopSelection = negKaonSel    )
  posMuonSeq     = SelectionSequence("PosMuonSeq",     TopSelection = posMuonSel    )
  negMuonSeq     = SelectionSequence("NegMuonSeq",     TopSelection = negMuonSel    )
  posElectronSeq = SelectionSequence("PosElectronSeq", TopSelection = posElectronSel)
  negElectronSeq = SelectionSequence("NegElectronSeq", TopSelection = negElectronSel)
  posProtonSeq   = SelectionSequence("PosProtonSeq",   TopSelection = posProtonSel  )
  negProtonSeq   = SelectionSequence("NegProtonSeq",   TopSelection = negProtonSel  )
  
  filterSequences = [
    posPionSeq    ,
    negPionSeq    ,
    posKaonSeq    ,
    negKaonSeq    ,
    posMuonSeq    ,
    negMuonSeq    ,
    posElectronSeq,
    negElectronSeq,
    posProtonSeq  ,
    negProtonSeq  ,
  ]
  
  
  posPionTuple     = DecayTreeTuple("PosPionTuple",     Decay = "pi+" , Inputs = [posPionSeq    . outputLocation() ], ToolList = [])
  negPionTuple     = DecayTreeTuple("NegPionTuple",     Decay = "pi-" , Inputs = [negPionSeq    . outputLocation() ], ToolList = [])
  posKaonTuple     = DecayTreeTuple("PosKaonTuple",     Decay = "K+"  , Inputs = [posKaonSeq    . outputLocation() ], ToolList = [])
  negKaonTuple     = DecayTreeTuple("NegKaonTuple",     Decay = "K-"  , Inputs = [negKaonSeq    . outputLocation() ], ToolList = [])
  posMuonTuple     = DecayTreeTuple("PosMuonTuple",     Decay = "mu+" , Inputs = [posMuonSeq    . outputLocation() ], ToolList = [])
  negMuonTuple     = DecayTreeTuple("NegMuonTuple",     Decay = "mu-" , Inputs = [negMuonSeq    . outputLocation() ], ToolList = [])
  posElectronTuple = DecayTreeTuple("PosElectronTuple", Decay = "e+"  , Inputs = [posElectronSeq. outputLocation() ], ToolList = [])
  negElectronTuple = DecayTreeTuple("NegElectronTuple", Decay = "e-"  , Inputs = [negElectronSeq. outputLocation() ], ToolList = [])
  posProtonTuple   = DecayTreeTuple("PosProtonTuple",   Decay = "p+"  , Inputs = [posProtonSeq  . outputLocation() ], ToolList = [])
  negProtonTuple   = DecayTreeTuple("NegProtonTuple",   Decay = "p~-" , Inputs = [negProtonSeq  . outputLocation() ], ToolList = [])
  
  tuples = [
      posPionTuple    
      , negPionTuple    
      , posKaonTuple    
      , negKaonTuple                         
      , posMuonTuple    
      , negMuonTuple    
      , posElectronTuple
      , negElectronTuple
      , posProtonTuple  
      , negProtonTuple  
  ]
  
  for tuple in tuples:
      b = tuple.addBranches ({'probe' : tuple.Decay})
      b = b['probe']
      eventTool = tuple.addTupleTool("LoKi::Hybrid::EvtTupleTool/LoKiEvent")
      eventTool.VOID_Variables = EventInfo
      eventTool.Preambulo = [
        "from LoKiTracks.decorators import *",
        "from LoKiCore.functions import *"
      ]
  
      lokitool = tuple.addTupleTool("LoKi::Hybrid::TupleTool/LoKiPart")
      lokitool.Variables = LokiVarsByType [ 'TRACK' ]
      lokitool.Preambulo = [
        "from LoKiProtoParticles.functions import *"
      ]
      pidcalibtool = b.addTupleTool ( "TupleToolPIDCalib/PIDCalibTool")


  return filterSequences + tuples
  
  
                      
#
##below for local test
#
#TupleFile = 'PIDMC.root'
#DataType = '2012'
#InputType = 'DST'
#Simulation = True
#Lumi = False
#EvtMax = -1
#
#dv = DaVinci()
#dv.DataType   = DataType
#dv.InputType  = InputType
#dv.EvtMax = EvtMax
#dv.TupleFile = TupleFile
#
#
#dv.Simulation = Simulation
#dv.Lumi       = Lumi
#
#dv.UserAlgorithms = filterSequences + tuples
#
#
###############
                    


## Entry point for ganga configuration
def configurePIDCalibTupleProduction(
      DataType
      , TupleFile
      , Simulation 
      , Lumi
      , Stream
      , InputType = 'DST'
      , EvtMax    = -1
      , tesFormat = "/Event/<stream>/<line>/Particles"
      , sTableFile = None
      , protoRecalibrationSequences = None
   ):
#
  from Configurables import MessageSvc
  MessageSvc().Format = "% F%60W%S%7W%R%T %0W%M"

  if protoRecalibrationSequences == None: 
    protoRecalibrationSequences = []

#  MessageSvc().OutputLevel = DEBUG
#
  dv = DaVinci()
  dv.DataType   = DataType
  dv.InputType  = InputType
  dv.EvtMax = EvtMax
  dv.TupleFile = TupleFile
#   
  if InputType == 'MDST':
    rootInTes = "/Event/"+Stream
#    uDstConf ( rootInTes )


  dv.Simulation = Simulation
  dv.Lumi       = Lumi

  tesFormat = tesFormat.replace('<stream>', Stream)
  
  configuredAlgorithms = parseConfiguration_Run2() 

  dv.appendToMainSequence ( configuredAlgorithms )



