## LivZBosonExample_topOptions.py

## This set of job options demonstrates some basic functionality of the code in LiverpoolAnalysis.
## It allows you to:
## - Perform particle preselection (defined in LivZBosonExample_preselection.py)
## - Write/Read a DPD of these preselected particles if desired
## - Makes combinations of electrons to form Z candidates
## - Dumps Z candidates, preselected electrons, all electrons, various trigger decisions and Z->ee truth
## This uses various scripts from LivTools and LivZAnalysis, some of which take configuration flags.
## There are many algorithms used in this example. For full documentation, refer to Doxygen for each specific class.
## Generally, particles are "transmitted" from one algorithm to the next via StoreGate.

## Say you have these lines in the job options:
##
##     from LivAlgorithms.LivAlgorithmsConf import Analysis__Combine2Particles
##     FindZee = Analysis__Combine2Particles(name = 'ZeeCandidates')
## 
## This creates the python equivalent of the Analysis::Combine2Particles class from LivAlgorithms.
## It is generally a good idea to always pass in an explicit name parameter to the constructor. Then,
## when you do:
##
##     jobSequence += FindZee
## 
## you can the always access this algorithm downstream as jobSequence.ZeeCandidates
## Note, 'FindZee' is only a temporary, local name for the algo. The same comments apply to tools.

## The job structure is very flexible - this example shows one possibility.
## For illustration, each "dumper" (except the trigger ones) records its output in a unique tree.
## To change this, simply change the TreeName property of the relevant algorithms.

## -------- Beginning of job options

## Standard job configuration from LivTools
## This sets up an AlgSequence called 'job', and an AthSequence called 'jobSequence' as its first element.
## An AthSequence is like an AlgSequence, but if an algorithm it contains reports that a filter has failed,
## the subsequent algorithms in that AthSequence are not run.
## In this example, the AthSequence is used for the "computational" sections, while the AlgSequence is used 
## for the ntuple dumpers (as well as holding the AthSequence).
## This ensures that the dumpers are run every event, and all the different trees will remain in synch.
## attaching them (at relevant points) to the AthSequence would be more efficient, but this needs careful treatment
## to avoid mistakes.
## Note: In this specific example, none of the algorithms attached to the AthSequence will ever "fail" their filter.
## It's a 1 line addition to enable this behaviour.

# The following block is broken in V15 so we revert back to the job config from the share
# directory while we try to work out hte problem (Carl)
# ## Note the new config. This takes the file from the python directory of the package,
# ## rather than share. It is much more robust and future-proof. The lines here set things
# ## up more or less as they were before - job and the Sequencer can also be accessed
# ## directly from the LivTools_jobConfig module
# #from LivTools import LivTools_jobConfig
# #job = LivTools_jobConfig.job
# #jobSequence = job.Sequencer
# 
# ## This construct allows you to set DetDescrVersion on the command line, using athena -c
# #if 'DetDescrVersion' not in dir():
# #    DetDescrVersion = 'ATLAS-GEO-02-01-00' # For 14.2.10+
# #LivTools_jobConfig.setDetDescrVersion(DetDescrVersion)
# 
# ## Add these lines (or equivalent) to get access to isolation and track isolation tools
# #XIPT   = LivTools_jobConfig.setupXImpactParamTool()
# #TrackIsolationTool = LivTools_jobConfig.setupTrackIsolationTool()
include("LivTools/LivTools_jobConfig.py")

## Helpful utilities from LivTools:
## This line does nothing in itself, but it defines several useful python functions.
from LivTools.LivTools_jobUtils import *

## Trigger setup (comment if not needed)
from LivTools import LivTools_triggerConfig

## Flags to control job flow
ReadAOD = True                           # AOD input
DumpDPD = False                          # AOD input and DPD output (after preselection)
ReadDPD = False                          # DPD input (no nead to rerun preselection)
DumpNtuple = True                        # Ntuple output (may be as well as/instead of DPD)
outputFileName = 'TestNtupleOutput.root' # Ntuple output file name
myOutputLevel = DEBUG                    # Logging level
UseFileStager = False                    # use file stager for running over dpm files much faster

## Input files and number of events.
## Note: To read DPM files, you need a valid grid certificate: see ~gwilliam/DpmTestInstructions.txt
if ReadAOD:
    #filePath = 'rfio:/dpm/ph.liv.ac.uk/home/atlas/atlasliverpooldisk/mc08.106050.PythiaZee_1Lepton.recon.AOD.e347_s462_r541/'
    filePath = '/hepstore/store4/flowerdew/AOD14/mc08.106050.PythiaZee_1Lepton.recon.AOD.e347_s462_r541/'
    #filePath = 'rfio:/dpm/ph.liv.ac.uk/home/atlas/atlasliverpooldisk/mc08.106050.PythiaZee_1Lepton.merge.AOD.e347_s462_r635_t53_tid064376/'
else:
    filePath = '.'
    
## If myInputCollections is already set, this is taken to be the list of input files (i.e. prepended for farm use)
##myInputCollections=['rfio:/dpm/ph.liv.ac.uk/home/atlas/atlasliverpooldisk/mc08.106050.PythiaZee_1Lepton.merge.AOD.e347_s462_r635_t53_tid064376/AOD.064376._00069.pool.root.1']
## getFileList is defined in LivTools/LivTools_jobUtils.py
try:
    myInputCollections
except NameError:
    myInputCollections = getFileList(filePath, '*AOD*')

if UseFileStager:
    include("LivTools/LivTools_FileStagerConfig.py")
    configureFileStager(jobSequence, myInputCollections)
else:
    svcMgr.EventSelector.InputCollections = myInputCollections

## Number of events to process.
theApp.EvtMax = 10

## Step 1: If we're reading from AOD (ie not DPD), do the preselection.
if ReadAOD:
    include("LivZAnalysis/LivZBosonExample_preselection.py")

## Step 2: If we're writing the DPD, dump the preselected electrons.
if DumpDPD:
    include("LivZAnalysis/LivZBosonExample_DPDWriter.py")

## Step 3: If we're writing an ntuple, do e+e combinations and dump the rest of the information.
if DumpNtuple:

    ## LivAlgorithms/Combine2Particles: used to make the Z candidates
    ## calculated from AOD or DPD contents
    ## This algorithm is booked and configured in 3 stages, called A, B and C in what follows.
    ## The subsequent algorithms follow the same pattern.

    ## Step A: Import and construct the algorithm
    from LivAlgorithms.LivAlgorithmsConf import Analysis__Combine2Particles
    FindZee                             = Analysis__Combine2Particles(name = 'ZeeCandidates')

    ## Step B: Configure it correctly
    ## Combine2Particles is a very general algorithm, we need to tell it exactly what to expect
    ## and how to classify the output (ie as a Z boson)
    FindZee.PdgId                       = 23
    FindZee.MinMass                     = 70*Units.GeV  ## Mass cut: a charge cut can also be applied
    FindZee.MaxMass                     = 110*Units.GeV
    FindZee.OutputLevel                 = myOutputLevel
    FindZee.FirstInputParticleContainer ='PreselectedElectronCollection' ## Use the preselected electrons as input.
    ## Combine2Particles writes a CompositeParticleContainer to StoreGate - we need to provide a name for it.
    FindZee.OutputCombinationContainer  ='ZeeContainer'

    ## Step C: Add this to the AthSequence.
    jobSequence += FindZee

    ## Ntuple algorithms.
    ## Now all the containers we need have been made.
    ## This leaves us free to dump the ones we want in any order we like, but it's probably best to keep it fairly logical!

    ## Electron and trigger dumping.
    ## As an example we'll dump the 'raw' AOD electrons as well as the preselected ones.
    ## This, and the trigger decision, need the full AOD in this setup.

    ## We need DumpEgamma later as well, so let's import outside the if statement.
    from LivDumpers.LivDumpersConf import Analysis__DumpEgamma

    if ReadAOD:
        # Example of an ntupling stage before preselection
        DumpAllE                    = Analysis__DumpEgamma(name = 'AllElectronDump') ## Note the name parameter
        DumpAllE.EgammaContainer    = 'ElectronAODCollection'     ## Input collection: AOD electrons
        DumpAllE.TreeName           = 'RawElectronTree'           ## Output tree
        DumpAllE.OutputBranchPrefix = 'electron'                  ## Branches will be called 'electron_X'
        DumpAllE.OutputLevel        = myOutputLevel               ## Logging level.
        job += DumpAllE                                           ## Add to job (the AlgSequence) so we dump every event.

        ## Trigger decision section, using LivDumpers/DumpTriggerResult.
        ## This dumps the L1, L2, EF decisions for one given EF chain.
        ## This example shows how several instances of this algorithm can be configured.
        ## I've just used a random selection of electron and muon triggers.

        from LivDumpers.LivDumpersConf import Analysis__DumpTriggerResult
        triggers = ['e22i_tight','e20_loose','e15i_medium','e10_medium','mu6','2mu6','3mu6','mu10','mu20','mu20i']

        ## All the trigger dumping is made to one tree in this example.
        triggerTreeName = 'TrigResultTree'

        ## Python loop over the triggers.
        for item in triggers:

	    ## Make a new instance - note how the name parameter is always unique.
            trig                  = Analysis__DumpTriggerResult(name = 'ExampleTrigger_'+item)

            ## Tool configuration.
            trig.TrigDecisionTool = ToolSvc.TrigDecisionTool
            trig.TreeName         = triggerTreeName

            ## Normally, LivDumpers/DumpTriggerResult only needs the EF chain name.
            ## Levels 1 and 2 should only have to be configured manually in exceptional circumstances:
            ## eg, when reading 13.0.30 AOD, each Level 1 item needs to be given too.
            trig.EFTriggerChain   = item

            job += trig

    ## Dump preselected electrons.
    ## Let's use another instance of Analysis__DumpEgamma which was imported above.
    ## Compare the configuration with that of 'AllElectronDump'.
    DumpE                    = Analysis__DumpEgamma(name = 'ExampleElectronDump')
    DumpE.EgammaContainer    = 'PreselectedElectronCollection'
    DumpE.TreeName           = 'PreselectedElectronTree'
    ## Branch prefix: the same as in AllElectronDump, but we have a new tree so there is no conflict.
    DumpE.OutputBranchPrefix = 'electron'
    DumpE.OutputLevel        = myOutputLevel
    ## Let's dump all the isem variables
    DumpE.DoEGTruth            = True
    DumpE.DoIsolCones          = True
    DumpE.DoImpactParameters   = True
    DumpE.DoTrackSummary       = True
    DumpE.DoTrackProbabilities = True
    DumpE.IsolationTool        = TrackIsolationTool
    DumpE.ImpactParameterTool  = ToolSvc.LivIPTool

    job += DumpE

    ## LivDumpers/DumpEventInfo dumps some useful information for real data (lumi block, for example).
    ## It's not really needed for MC, as the event and run numbers are added to each tree by default.
    ## On very recent AOD, it can also dump pdf information, but this is off by default.
    ## This algo is small, it doesn't need much configuration.

    from LivDumpers.LivDumpersConf import Analysis__DumpEventInfo
    DumpEI             = Analysis__DumpEventInfo(name = 'EventInfoDump')
    DumpEI.TreeName    = 'EventInfoTree'
    DumpEI.OutputLevel = myOutputLevel
    job += DumpEI

    ## Now dump the Z candidates using LivDumpers/DumpCombinedParticle.
    ## This takes the CompositeParticleContainer output by the FindZee algo (note the FindZee name is still in scope).
    ## It dumps information on the Z four-momentum, charge and pdg code,
    ## and the same information for each of its constituents.

    from LivDumpers.LivDumpersConf import Analysis__DumpCombinedParticle
    DumpZ = Analysis__DumpCombinedParticle(name = 'ZDump')
    DumpZ.CombinedParticleContainer = FindZee.OutputCombinationContainer
    DumpZ.TreeName                  = 'ZeeTree'
    DumpZ.OutputBranchPrefix        = 'Zee'
    DumpZ.OutputLevel               = myOutputLevel
    job += DumpZ

    ## Finally, let's dump some Z->ee truth information using LivDUmpers/DumpSingleBosonTruth.
    ## The algo can do (almost) any 1->2 signature, it has to be configured for each one.
    ## And we'll use the filtered truth container defined in the preselection.

    ## First, we need the associated tool
    from LivTools.LivToolsConf import Analysis__SingleBosonTruthTool
    TruthTool = Analysis__SingleBosonTruthTool(name = 'TruthTool')
    TruthTool.TruthContainer   = TruthFilter.OutputTruthContainer ## TruthFilter is the truth tool in the preselection
    TruthTool.BosonPDG         = 23                               ## We're looking for Z bosons ..
    TruthTool.DecayProduct1PDG = 11                               ## .. that decay to electrons
    TruthTool.OutputLevel      = myOutputLevel                    ## Logging level

    ## TruthTool is a tool, so we can't add it to the job directly.
    ## Because it will be a "public" tool in the algorithm, we instead add it to ToolSvc,
    ## which manages public tools.
    ## If it was going to be a private tool, we wouldn't do this step, it would belong
    ## to the algorithm, not ToolSvc.
    ToolSvc += TruthTool

    from LivDumpers.LivDumpersConf import Analysis__DumpSingleBosonTruth
    DumpTruth = Analysis__DumpSingleBosonTruth(name = 'ZeeTruthDump')
    DumpTruth.TruthTool          = TruthTool       ## This is the tool we just configured
    DumpTruth.TreeName           = 'ZeeTruth'      ## The output tree name
    DumpTruth.OutputLevel        = myOutputLevel   ## Logging level
    DumpTruth.OutputBranchPrefix = 'Zee'           ## There's no conflict with ZDump here - it's a different tree
    DumpTruth.DoDocuParticles    = False           ## The documentary particles are pre-FSR
    job += DumpTruth

    ## Generic CBNT setup, function defined in LivTools/LivTools_jobUtils.py
    ## Most of this is hidden by the athena framework, but the "Dumper" algorithms don't write directly to file.
    ## They write to a stream - this function defines streams for each unique TreeName in the job.
    ## If some/all of the Dumpers were attached to jobSequence, you'd need to call addStreams(jobSequence, outputFilename)
    ## as well as/instead of this function call.
    addStreams(job, outputFileName)

## This just does some general output formatting.
svcMgr.MessageSvc.OutputLevel = WARNING
svcMgr.MessageSvc.Format = "% F%50W%S%7W%R%T %0W%M"

## If you want to get rid of the event-by-event messages from AthenaEventLoopMgr,
## uncomment the next two lines
#from AthenaServices.AthenaServicesConf import AthenaEventLoopMgr
#AthenaEventLoopMgr.OutputLevel = WARNING

## Uncommenting some or all of these can help debugging.
#print job
#print svcMgr
#print ToolSvc