WIP: allow user to override the nbatch, nbx per batch, bx spacing.

#!/usr/bin/env python
import os
import random
from Gaudi.Configuration import *

from Configurables import k4DataSvc
dsvc = k4DataSvc("EventDataSvc", input="simsignal_v01.root")

from Configurables import PodioInput
podioinput = PodioInput("PodioReader", collections=[
    "MCParticle",
    "VXDCollection",
    "ITKBarrelCollection",
    "ITKEndcapCollection",
    "TPCCollection",
    "OTKBarrelCollection",
    "OTKEndcapCollection",
    "EcalBarrelCollection", "EcalBarrelContributionCollection",
    "EcalEndcapsCollection", "EcalEndcapsContributionCollection",
    "EcalEndcapRingCollection", "EcalEndcapRingContributionCollection",
    "HcalBarrelCollection", "HcalBarrelContributionCollection",
    "HcalEndcapsCollection", "HcalEndcapsContributionCollection",
    "HcalEndcapRingCollection", "HcalEndcapRingContributionCollection",
    "MuonBarrelCollection",
    "MuonEndcapCollection"
])

from Configurables import DetSimMixingAlg
evtmixing = DetSimMixingAlg("EvtMixingAlg")
# evtmixing.BackgroundRates = {"testing": 1.0e6}
# evtmixing.BackgroundFileLists = {"testing": ["test.root"]}

bkg_timings = {
    "Pair": -2770,
    "BGB": 4145.41,
    "BGC": 87781.62,
    # "BTH": 357.20,
    # "TSC": 103.77
}

# following data is produced by Haoyu Shi
def bkg_files(bkg_type, bkg_version="250414"):
    bkg_topdir = "/cefs/higgs/shihy/tools/CEPCSW_TDR2536/CEPCSW/Test/Higgs"
    bkg_typedir = os.path.join(bkg_topdir, "%s_%s"%(bkg_type, bkg_version))
    filename = "test-bgsimdose.root"

    # todo: will load files randomly from the directory
    # subdir = "0"
    # bkg_files = [ os.path.join(bkg_typedir, subdir, filename) ]

    bkg_files = [os.path.join(bkg_typedir, d, filename) 
                    for d in os.listdir(bkg_typedir) 
                    if os.path.isdir(os.path.join(bkg_typedir, d))]
    
    # randomize the input files
    random.shuffle(bkg_files)

    return bkg_files[:10]

evtmixing.BackgroundTimings = bkg_timings
evtmixing.BackgroundFileLists = {
    key: bkg_files(key)
    for key in bkg_rates.keys()
}

# output
from Configurables import PodioOutput
out = PodioOutput("outputalg")
out.filename = "evtmix_v01.root"
out.outputCommands = ["drop *", "keep MCParticle", "keep Mixed*"]

from Configurables import ApplicationMgr
ApplicationMgr().TopAlg += [podioinput, evtmixing, out]
ApplicationMgr().ExtSvc += [dsvc]
ApplicationMgr().EvtSel = "NONE"
ApplicationMgr().EvtMax = 10
ApplicationMgr().OutputLevel = INFO