#!/usr/bin/env python3

from pathlib import Path
from multiprocessing import Pool
from functools import partial

# This script runs a Geant4 simulation in parallel by passing chunks of events to subprocesses.
# This is a workaround to achieve parallel processing even though Geant4 is not thread-save
# and thus the internal parallelism of the ACTS examples framework cannot be used.
#
# Note that:
# ==========
#
# * This should give equivalent results to a sequential run if the RNG is initialized with
#   the same seed in all runs
#
# * So far this works only for csv outputs, since they write the results in one file per event
#   (the naming of the csv-files should be equivalent to a sequential run)
#
# * In principle it is not difficult to extend this for ROOT files as well. One would need to
#   write the root-files into separate directory per chunk, and then use ROOT's hadd to combine
#   the output files.
#


def runGeant4EventRange(detector, trackingGeometry, beginEvent, endEvent, outputDir):
    import acts
    import acts.examples
    from acts.examples.simulation import addParticleGun, addGeant4, EtaConfig

    u = acts.UnitConstants

    field = acts.ConstantBField(acts.Vector3(0, 0, 2 * u.T))
    rnd = acts.examples.RandomNumbers(seed=42)

    s = acts.examples.Sequencer(
        events=endEvent - beginEvent, skip=beginEvent, numThreads=1
    )

    outputDir = Path(outputDir)
    addParticleGun(
        s,
        EtaConfig(-2.0, 2.0),
        rnd=rnd,
        outputDirCsv=outputDir / "csv",
        outputDirRoot=None,
    )
    addGeant4(
        s,
        detector,
        trackingGeometry,
        field,
        outputDirCsv=outputDir / "csv",
        outputDirRoot=None,
        rnd=rnd,
    )

    s.run()


if "__main__" == __name__:
    from acts.examples.odd import getOpenDataDetector

    detector = getOpenDataDetector()
    trackingGeometry = detector.trackingGeometry()
    decorators = detector.contextDecorators()

    n_events = 100
    n_jobs = 8

    chunksize = n_events // (n_jobs - 1)
    begins = range(0, n_events, chunksize)
    ends = [min(b + chunksize, n_events) for b in begins]

    outputDir = Path.cwd()

    with Pool(n_jobs) as p:
        p.starmap(partial(runGeant4EventRange, outputDir=outputDir), zip(begins, ends))