#!/bin/bash

# load common HLT functions
if [ -f "$CMSSW_BASE/src/HLTrigger/Configuration/common/utils.sh" ]; then
  source "$CMSSW_BASE/src/HLTrigger/Configuration/common/utils.sh"
elif [ -f "$CMSSW_RELEASE_BASE/src/HLTrigger/Configuration/common/utils.sh" ]; then
  source "$CMSSW_RELEASE_BASE/src/HLTrigger/Configuration/common/utils.sh"
else
  exit 1
fi

function log() {
  echo -e "$@"
}

function err() {
  echo -e "$@" 1>&2
}

NAME=$(basename $0)

HELP="Run the integration tests on a given HLT menu.

Usage:
  $NAME -h|--help
  $NAME [-d|--dir WORKDIR] [-s|--setup SETUP] [-i|--input RAW] [-j|--jobs JOBS]
        [--streams STREAMS] [--threads THREADS] [-x|--extra OPTIONS] [--mc]
        [ [-n|--size EVENTS] [-k|--skip EVENTS] | [-e|--events EVENTS] ] MENU

  MENU is the HLT menu to test (format: a local cmsRun cfg file, or the name of a ConfDB configuration).

  -s | --setup SETUP              HLT menu (format: ConfDB configuration) used for Services and EventSetup modules
                                  (useful when testing a ConfDB config that only contains Paths).
                                  Note it is an error to specify the converter/db here,
                                  it uses the same as set by the HLT menu
  -d | --dir         WORKDIR      Create all files and run all tests inside WORKDIR (defauls: ./hltintegration)
  -i | --input       INPUT        Use the specified EDM file as input (data tier: RAW)
  -n | --size        EVENTS       Run on EVENTS events (-1 for all, default is 100)
  -k | --skip        EVENTS       Skip the first EVENTS (default is 0)
  -e | --events      EVENTS       Run on a comma-separated list of EVENTS, a VEventRange
  -j | --jobs        JOBS         Run JOBS single trigger jobs in parallel (default 4)
       --streams     STREAMS      Run with STREAMS parallel streams (i.e. events) (default 0 means as many streams as threads)
       --threads     THREADS      Run with THREADS threads when running the whole HLT (default 4)
  -a | --accelerator ACCELERATOR  Keyword to choose allowed accelerators (examples: \"*\", \"cpu\", \"gpu-nvidia\")
  -p | --paths       PATHS        Comma-separated list of Path names (incl. wildcards)
                                  to select which Paths are tested standalone.
                                  If a Path-name pattern starts with the dash character (-),
                                  the Paths whose name matches that pattern will be ignored.
  -x | --extra       OPTIONS      If the HLT menu is a local cmsRun cfg file, OPTIONS is used as
                                  additional arguments to cmsRun (i.e. \"cmsRun hlt.py [OPTIONS]\")
                                  If the HLT menu is the name of a ConfDB configuration, OPTIONS is used as
                                  additional arguments to hltGetConfiguration (i.e. \"hltGetConfiguration [MENU] [..] [OPTIONS]\")
       --mc                       Run over MC instead of data (the default)
  --dbproxy                       Use a socks proxy to connect to the ConfDB database
  --dbproxyhost      PROXYHOST    Host of the socks proxy (default: \"localhost\")
  --dbproxyport      PROXYPORT    Port of the socks proxy (default: 8080)
  -h | --help                     Print this help message and exit.

  The HLT menu used for the integration tests (MENU) can be specified
  as either (1) a cmsRun cfg file, or (2) the name of a ConfDB configuration.
  MENU is identified as a cmsRun cfg file if it is an existing regular file.
  For ConfDB configurations, the supported formats are
    - /path/to/configuration[/Vn]
    - [[{v1|v2|v3}/]{run3|run2|online|adg}:]/path/to/configuration[/Vn]
    - run:runnumber
  Allowed converters are \"v1\", \"v2\", and \"v3\" (default).
  Allowed databases are
    - \"run3\" (default, used for offline Run-3 development),
    - \"run2\" (previously used for Run-2 development),
    - \"online\" (used to extract online menus within Point 5) and
    - \"adg\" (used to extract the online menus outside Point 5).
  Other converters and databases exist, but they are for expert/special use only.
  If \"run:\" is used instead, the HLT menu used for the given run number is looked up and used.
  If no menu version is specified, the latest one is automatically used.

  It is possible to pass arbitrary command-line options to hltGetConfiguration, using \"-x --option\".
  To pass multiple options, enclose them in quotes, or use \"-x\" more than once.

  Note: '--events' is not supported together with '--size' or '--skip'.


Examples:

  $NAME /dev/CMSSW_X_Y_Z/GRun

      will test the latest version of the GRun menu.


  $NAME /dev/CMSSW_X_Y_Z/GRun -x --l1-emulator

      will test the latest version of the GRun running the L1-Trigger emulator.


  $NAME /users/fwyzard/physics/HLT/V6 -s adg:/cdaq/physics/Run2011/1e33/v1.3/HLT/V6

      will test the paths from /users/fwyzard/physics/HLT/V6
      using the environment from the online menu \"1e33\" v1.3 V6

"

# parse command line argument and options
OPTS=$(getopt -n "$NAME" -o "s:d:i:j:n:k:e:p:a:x:h" \
 -l "setup:,dir:,input:,jobs:,size:,skip:,streams:,threads:,paths:,accelerator:,events:,mc,extra:,help,dbproxy,dbproxyhost:,dbproxyport:" -- "$@")

# invalid options
if [ $? != 0 ]; then
  exit 1
fi

# reload the parsed options into the environment
eval set -- "$OPTS"

# check how many CPUs are available
CPUS=`getconf _NPROCESSORS_ONLN`

MENU=""
SETUP=""
INPUT=""
SIZE=100
SKIP=0
EVENTS=""
JOBS=4
THREADS=4
STREAMS=0
PATHS=""
ACCELERATOR="cpu"
WORKDIR="hltintegration"
EXTRA=""
DATA=""
DBPROXYOPTS=""

SELECTION=""

# parse options
while true; do
  case "$1" in
    "-h" | "--help" )
      echo "$HELP"
      exit 0
      ;;
    "-s" | "--setup" )
      SETUP="$2"
      shift 2
      ;;
    "-d" | "--dir" )
      WORKDIR="$2"
      shift 2
      ;;
    "-i" | "--input" )
      INPUT="$2"
      shift 2
      ;;
    "-n" | "--size" )
      if [ "$SELECTION" == "complex" ]; then
        err "'--events' is not supported together with '--size' or '--skip'"
        exit 1
      fi
      SELECTION="simple"
      SIZE=$2
      if ((SIZE == 0)) && [ "$SIZE" != "0" ]; then
        err "$NAME error: invalid option \"$1 $2\""
        err "Try '$NAME --help' for more information."
        exit 1
      fi
      shift 2
      ;;
    "-k" | "--skip" )
      if [ "$SELECTION" == "complex" ]; then
        err "'--events' is not supported together with '--size' or '--skip'"
        exit 1
      fi
      SELECTION="simple"
      SKIP=$2
      if ((SKIP == 0)) && [ "$SKIP" != "0" ]; then
        err "$NAME error: invalid option \"$1 $2\""
        err "Try '$NAME --help' for more information."
        exit 1
      fi
      shift 2
      ;;
    "-e" | "--events" )
      if [ "$SELECTION" == "simple" ]; then
        err "'--events' is not supported together with '--size' or '--skip'"
        exit 1
      fi
      SELECTION="complex"
      SIZE=-1
      EVENTS="$2"
      shift 2
      ;;
    "-j" | "--jobs" )
      JOBS=$2
      if ((JOBS == 0)); then
        err "$NAME error: invalid option \"$1 $2\""
        err "Try '$NAME --help' for more information."
        exit 1
      fi
      shift 2
      ;;
    "--streams" )
      STREAMS=$2
      shift 2
      ;;
    "--threads" )
      THREADS=$2
      shift 2
      ;;
    "-p" | "--paths" )
      PATHS="$2"
      shift 2
      ;;
    "-a" | "--accelerator" )
      ACCELERATOR="$2"
      shift 2
      ;;
    "-x" | "--extra" )
      EXTRA="$EXTRA $2"
      shift 2
      ;;
    "--mc" )
      DATA="--mc"
      shift 1
      ;;
    "--dbproxy" )
      DBPROXYOPTS="${DBPROXYOPTS} --dbproxy"
      shift 1
      ;;
    "--dbproxyhost" )
      DBPROXYOPTS="${DBPROXYOPTS} --dbproxyhost $2"
      shift 2
      ;;
    "--dbproxyport" )
      DBPROXYOPTS="${DBPROXYOPTS} --dbproxyport $2"
      shift 2
      ;;
    "--" )
      # inserted by getopt to singal the end of options
      shift
      break
      ;;
  esac
done

# remove spurious whitespaces and tabs from EXTRA and DBPROXYOPTS
EXTRA=$(echo "${EXTRA}" | xargs)
DBPROXYOPTS=$(echo "${DBPROXYOPTS}" | xargs)

# parse required argument
if (( $# == 0 )); then
  err "$NAME error: missing argument."
  err "Try '$NAME --help' for more information."
  exit 1
elif (( $# > 1 )); then
  err "$NAME error: too many arguments."
  err "Try '$NAME --help' for more information."
  exit 1
else
  MENU="$1"
fi

log "----------------------------"
log "Starting hltIntegrationTests"
log "----------------------------"

# create empty output directory
rm -rf "${WORKDIR}"
mkdir -p "${WORKDIR}"

# if MENU = local cfg file, copy it to output directory
# (that copy will be customised and used for the integration tests)
if [ -f "${MENU}" ]; then
  cp "${MENU}" "${WORKDIR}"/hlt.py
fi

# move to, and run tests from, the output directory
cd "${WORKDIR}"

if [ -f hlt.py ]; then

  # customise cfg file
  log "Creating customised version of input cfg file (${MENU})"

  # warn that ${DATA} will be ignored
  [ ! "${DATA}" ] || printf "\n%s" "WARNING -- variable \${DATA}=\"${DATA}\" will be ignored !"

  cat <<@EOF >> hlt.py

# change name of cms.Process
process.setName_("TEST$(date -u +'%Y%m%d%H%M%S')")

# disable HLT prescales
if hasattr(process, 'PrescaleService'):
  del process.PrescaleService

# set max number of input events
process.maxEvents.input = ${SIZE}
@EOF

  if [ "x${INPUT}" != "x" ]; then
    cat <<@EOF >> hlt.py

# set input EDM file
if hasattr(process, 'source') and hasattr(process.source, 'fileNames'):
  process.source.fileNames = [
    "${INPUT}",
  ]
else:
  raise RuntimeError("ERROR -- unsupported cfg file: process.source.fileNames does not exist")
@EOF
  fi

  # set MENU to name of ConfDB configuration (if any)
  MENU=$(python3 -c """
import sys
# redefine sys.argv (necessary to import
# cfg file if the latter uses VarParsing)
sys.argv = ['python3', 'hlt.py']
from hlt import cms,process
try:
  print(process.HLTConfigVersion.tableName.value())
except:
  print('')
""")

  # show name of ConfDB configuration (if available)
  [ ! "${MENU}" ] || log "ConfDB configuration: ${MENU}"

else
  # if ${DATA} is empty, set it to "--data"
  [ "${DATA}" ] || DATA="--data"
  # download HLT menu from ConfDB
  HLTGETCMD="hltGetConfiguration ${MENU}
    --process \"TEST$(date -u +'%Y%m%d%H%M%S')\"
    --max-events ${SIZE} --no-prescale --no-output
    ${DATA} --input ${INPUT} ${EXTRA} ${DBPROXYOPTS}"
  HLTGETCMD=$(echo "${HLTGETCMD}" | xargs)
  log "Creating HLT menu from ConfDB configuration\n> ${HLTGETCMD}"
  ${HLTGETCMD} > hlt.py
  # unset EXTRA environment variable (used later in cmsRun jobs)
  unset HLTGETCMD EXTRA
fi

# if missing, add a simplified HLTriggerFinalPath
if ! grep -q HLTriggerFinalPath hlt.py; then
  cat >> hlt.py << @EOF

# add (simplified) HLTriggerFinalPath if missing
process.hltTriggerSummaryAOD = cms.EDProducer( "TriggerSummaryProducerAOD",
    processName = cms.string( "@" )
)
process.hltTriggerSummaryRAW = cms.EDProducer( "TriggerSummaryProducerRAW",
    processName = cms.string( "@" )
)
process.HLTriggerFinalPath = cms.Path( process.hltTriggerSummaryAOD + process.hltTriggerSummaryRAW )

@EOF
fi

# select which events to run on
if [ "${SELECTION}" == "complex" ]; then
  cat >> hlt.py << @EOF

# event selection customised by hltIntegrationTests
process.source.eventsToProcess = cms.untracked.VEventRange( '$(echo $EVENTS | sed -e"s/,/','/g")' )
@EOF

elif (( ${SKIP} > 0 )); then
  cat >> hlt.py << @EOF

# event selection customised by hltIntegrationTests
process.source.skipEvents = cms.untracked.uint32( ${SKIP} )
@EOF
fi

# set the number of threads and streams for the whole hlt job
cat >> hlt.py << @EOF

# configure multithreading, and allocate 10 MB of stack space per thread
process.options.numberOfThreads = ${THREADS}
process.options.numberOfStreams = ${STREAMS}
process.options.sizeOfStackForThreadsInKB = 10*1024
# set allowed accelerators
process.options.accelerators = [ "$ACCELERATOR" ]

process.hltTriggerSummaryAOD.throw = cms.bool( True )
@EOF

# find number of cms.Paths in the HLT menu (incl. Dataset Paths)
NUM_PATHS=$(python3 -c """
import sys
# redefine sys.argv (necessary to import
# cfg file if the latter uses VarParsing)
sys.argv = ['python3', 'hlt.py']
from hlt import cms,process
try:
  print(len(process.paths_()))
except:
  print(0)
""")
log "\nThe HLT menu contains ${NUM_PATHS} Paths (incl. Dataset Paths)"

# list of trigger Paths to be tested standalone (always exclude HLTriggerFinalPath)
log "\nPreparing list of Paths to be tested standalone (paths.txt)"
[ "${PATHS}" ] || PATHS="*"
PATHS+=",-HLTriggerFinalPath"
log " - Path selection: \"${PATHS}\""
TRIGGERS=$(hltListPaths hlt.py -p --no-dep --select-paths "${PATHS}")
echo "${TRIGGERS[@]}" > paths.txt

# print some info
if [ "${SELECTION}" == "complex" ]; then
  log "\nWill run full menu and $(echo $TRIGGERS | wc -w) Paths standalone over $(echo ${EVENTS} | tr ',' '\n' | wc -l) events, with ${JOBS} jobs in parallel"
elif [ "${SIZE}" == "-1" ]; then
  log "\nWill run full menu and $(echo ${TRIGGERS} | wc -w) Paths standalone over all events, with ${JOBS} jobs in parallel"
else
  log "\nWill run full menu and $(echo ${TRIGGERS} | wc -w) Paths standalone over ${SIZE} events, with ${JOBS} jobs in parallel"
fi

# check the prescale modules
hltCheckPrescaleModules -w hlt.py

log "\nPreparing single-Path configurations"
for TRIGGER in $TRIGGERS; do
  cat > "${TRIGGER}".py << @EOF
from hlt import *

process.hltOutput = cms.OutputModule( "PoolOutputModule",
    fileName = cms.untracked.string( "${TRIGGER}.root" ),
    fastCloning = cms.untracked.bool( False ),
    compressionAlgorithm = cms.untracked.string( "ZSTD" ),
    compressionLevel = cms.untracked.int32( 3 ),
    outputCommands = cms.untracked.vstring(
      'drop *',
      'keep edmTriggerResults_*_*_*',
    )
)

process.Output = cms.EndPath( process.hltOutput )

process.schedule = cms.Schedule( process.${TRIGGER}, process.HLTriggerFinalPath, process.Output )

process.hltTriggerSummaryAOD.throw = cms.bool( True )
@EOF
done

# if a separate setup is requested, create the setup_cff.py file and patch all dumps to use it
if [ "${SETUP}" ]; then

  if [ "${MENU}" ]; then
    # we use ${MENU} here, not ${SETUP}, as we force the same DB / converter as the main menu
    # this is the hltGetConfiguration behaviour and would be confusing if you had to
    # specify converter/db on the setup menu on hltIntegrationTests but not on hltGetConfiguration
    read SETUP_Vx SETUP_DB _ <<< $(parse_HLT_menu "${MENU}")
    log "\nCreating setup_cff from ConfDB configuration: ${SETUP_Vx}/${SETUP_DB}:${SETUP}"
    hltConfigFromDB --${SETUP_Vx} --${SETUP_DB} ${DBPROXYOPTS} --cff --configName "$SETUP" \
      --nopaths --services -FUShmDQMOutputService,-PrescaleService,-EvFDaqDirector,-FastMonitoringService > setup_cff.py
    sed -i -e's/process = cms.Process(.*)/&\nprocess.load("setup_cff")/' hlt.py $(for TRIGGER in ${TRIGGERS}; do echo "${TRIGGER}".py; done)
  else
    printf "%s\n" "WARNING -- \"--setup ${SETUP}\" will be ignored (failed to deduce name of HLT menu from hlt.py)"
  fi
fi

# run all HLT dumps
cat > .makefile << @EOF
TRIGGERS=$(echo ${TRIGGERS})
CFGS=\$(TRIGGERS:%=%.py)
LOGS=\$(TRIGGERS:%=%.log)
DONE=\$(TRIGGERS:%=%.done)

.PHONY: all clean hlt \$(TRIGGERS)

all: hlt \$(TRIGGERS)

clean:
	@rm -f hlt.log hlt.done \$(LOGS) \$(DONE)

hlt: hlt.done

hlt.done: hlt.py
	@echo -e "\tfull menu dump"
	@cmsRun hlt.py ${EXTRA} >& hlt.log < /dev/zero && touch hlt.done

\$(TRIGGERS): %: %.done

\$(DONE): %.done: %.py
	@echo -e "\t\$*"
	@cmsRun \$*.py ${EXTRA} >& \$*.log < /dev/zero && touch \$*.done
@EOF

log "\nRunning..."
# if the whole hlt job runs with multithreading, run it by itself
# otherwise, run it in parallel with the single-trigger jobs
if ((THREADS > 0)); then
  make -f .makefile hlt
  make -f .makefile -j${JOBS} -k ${TRIGGERS}
else
  make -f .makefile -j${JOBS} -k
fi

# compare HLT results
log "\nComparing the results of running each path by itself with those from the full menu"
hltCompareResults
STATUS=$?
log "--------------------------"
if [ "${STATUS}" -eq 0 ]; then
  log "hltIntegrationTests PASSED"
else
  log "hltIntegrationTests FAILED"
fi
log "--------------------------"
log "exit status: $STATUS"
cd ..
exit ${STATUS}