From ecb1375cbf9c027cc0a6b4817db2d1cfd1e81acc Mon Sep 17 00:00:00 2001
From: swenzel <sandro.wenzel@cern.ch>
Date: Tue, 21 Nov 2023 11:36:27 +0100
Subject: [PATCH] Improvements for generating events from O2 kine file

* optionally allow round-robin
* be able to give file name via config key params
---
 Generators/include/Generators/GeneratorFromFile.h        | 3 ++-
 Generators/include/Generators/GeneratorFromO2KineParam.h | 2 ++
 Generators/src/GeneratorFactory.cxx                      | 4 +++-
 Generators/src/GeneratorFromFile.cxx                     | 5 +++++
 4 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/Generators/include/Generators/GeneratorFromFile.h b/Generators/include/Generators/GeneratorFromFile.h
index 13b6239189a38..b75279cf3319f 100644
--- a/Generators/include/Generators/GeneratorFromFile.h
+++ b/Generators/include/Generators/GeneratorFromFile.h
@@ -92,7 +92,8 @@ class GeneratorFromO2Kine : public o2::eventgen::Generator
   int mEventsAvailable = 0;
   bool mSkipNonTrackable = true; //! whether to pass non-trackable (decayed particles) to the MC stack
   bool mContinueMode = false;    //! whether we want to continue simulation of previously inhibited tracks
-  ClassDefOverride(GeneratorFromO2Kine, 1);
+  bool mRoundRobin = false;      //! whether we want to take events from file in a round robin fashion
+  ClassDefOverride(GeneratorFromO2Kine, 2);
 };
 
 } // end namespace eventgen
diff --git a/Generators/include/Generators/GeneratorFromO2KineParam.h b/Generators/include/Generators/GeneratorFromO2KineParam.h
index 56642ccdd52db..c0e032f740aba 100644
--- a/Generators/include/Generators/GeneratorFromO2KineParam.h
+++ b/Generators/include/Generators/GeneratorFromO2KineParam.h
@@ -31,6 +31,8 @@ namespace eventgen
 struct GeneratorFromO2KineParam : public o2::conf::ConfigurableParamHelper<GeneratorFromO2KineParam> {
   bool skipNonTrackable = true;
   bool continueMode = false;
+  bool roundRobin = false;   // read events with period boundary conditions
+  std::string fileName = ""; // filename to read from - takes precedence over SimConfig if given
   O2ParamDef(GeneratorFromO2KineParam, "GeneratorFromO2Kine");
 };
 
diff --git a/Generators/src/GeneratorFactory.cxx b/Generators/src/GeneratorFactory.cxx
index 80e80dbefa091..5815543ad7100 100644
--- a/Generators/src/GeneratorFactory.cxx
+++ b/Generators/src/GeneratorFactory.cxx
@@ -149,7 +149,9 @@ void GeneratorFactory::setPrimaryGenerator(o2::conf::SimConfig const& conf, Fair
     LOG(info) << "using external kinematics";
   } else if (genconfig.compare("extkinO2") == 0) {
     // external kinematics from previous O2 output
-    auto extGen = new o2::eventgen::GeneratorFromO2Kine(conf.getExtKinematicsFileName().c_str());
+    auto name1 = GeneratorFromO2KineParam::Instance().fileName;
+    auto name2 = conf.getExtKinematicsFileName();
+    auto extGen = new o2::eventgen::GeneratorFromO2Kine(name1.size() > 0 ? name1.c_str() : name2.c_str());
     extGen->SetStartEvent(conf.getStartEvent());
     primGen->AddGenerator(extGen);
     if (GeneratorFromO2KineParam::Instance().continueMode) {
diff --git a/Generators/src/GeneratorFromFile.cxx b/Generators/src/GeneratorFromFile.cxx
index f3729bddc1bdf..39fd809ac5560 100644
--- a/Generators/src/GeneratorFromFile.cxx
+++ b/Generators/src/GeneratorFromFile.cxx
@@ -196,6 +196,7 @@ bool GeneratorFromO2Kine::Init()
   LOG(info) << param;
   mSkipNonTrackable = param.skipNonTrackable;
   mContinueMode = param.continueMode;
+  mRoundRobin = param.roundRobin;
 
   return true;
 }
@@ -259,6 +260,10 @@ bool GeneratorFromO2Kine::importParticles()
       particlecounter++;
     }
     mEventCounter++;
+    if (mRoundRobin) {
+      LOG(info) << "Resetting event counter to 0; Reusing events from file";
+      mEventCounter = mEventCounter % mEventsAvailable;
+    }
 
     if (tracks) {
       delete tracks;