From f494e436dbb3c7a8225d5d393652d23936f3c435 Mon Sep 17 00:00:00 2001
From: Sandro Wenzel <sandro.wenzel@cern.ch>
Date: Sun, 5 Oct 2025 07:26:05 +0200
Subject: [PATCH] Generators: Embedd into correct event when using
 collision-context

This fixes a problem with embedding and when option `--embedIntoFile`
is used. Here, we had an index problem in the primary generator class, which
was not taking into account the actual event index according the collision context.
This could lead to wrong vertices or using wrong MCHeader information for the embedding decision.
---
 .../include/Generators/PrimaryGenerator.h     |  3 +++
 Generators/src/PrimaryGenerator.cxx           | 25 +++++++++---------
 run/O2PrimaryServerDevice.h                   | 26 +++++++++++++++++++
 3 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/Generators/include/Generators/PrimaryGenerator.h b/Generators/include/Generators/PrimaryGenerator.h
index 097bfabf4ef57..cd4d842761418 100644
--- a/Generators/include/Generators/PrimaryGenerator.h
+++ b/Generators/include/Generators/PrimaryGenerator.h
@@ -80,6 +80,9 @@ class PrimaryGenerator : public FairPrimaryGenerator
   /** Public embedding methods **/
   Bool_t embedInto(TString fname);
 
+  /// sets the embedding index
+  void setEmbedIndex(int idx) { mEmbedIndex = idx; }
+
   void setExternalVertexForNextEvent(double x, double y, double z);
 
   // sets the vertex mode; if mode is kCCDB, a valid MeanVertexObject pointer must be given at the same time
diff --git a/Generators/src/PrimaryGenerator.cxx b/Generators/src/PrimaryGenerator.cxx
index ee7c6a16330f9..e7b70108a5d4b 100644
--- a/Generators/src/PrimaryGenerator.cxx
+++ b/Generators/src/PrimaryGenerator.cxx
@@ -81,7 +81,7 @@ Bool_t PrimaryGenerator::GenerateEvent(FairGenericStack* pStack)
   /** generate event **/
 
   /** normal generation if no embedding **/
-  if (!mEmbedTree) {
+  if (!mEmbedTree || mEmbedIndex < 0) {
     fixInteractionVertex(); // <-- always fixes vertex outside of FairROOT
     auto ret = FairPrimaryGenerator::GenerateEvent(pStack);
     if (ret) {
@@ -91,17 +91,18 @@ Bool_t PrimaryGenerator::GenerateEvent(FairGenericStack* pStack)
   }
 
   /** this is for embedding **/
-
-  /** setup interaction vertex **/
-  mEmbedTree->GetEntry(mEmbedIndex);
-  setInteractionVertex(mEmbedEvent);
-
-  /** notify event generators **/
-  auto genList = GetListOfGenerators();
-  for (int igen = 0; igen < genList->GetEntries(); ++igen) {
-    auto o2gen = dynamic_cast<Generator*>(genList->At(igen));
-    if (o2gen) {
-      o2gen->notifyEmbedding(mEmbedEvent);
+  if (mEmbedIndex >= 0) {
+    /** setup interaction vertex **/
+    mEmbedTree->GetEntry(mEmbedIndex);
+    setInteractionVertex(mEmbedEvent);
+
+    /** notify event generators **/
+    auto genList = GetListOfGenerators();
+    for (int igen = 0; igen < genList->GetEntries(); ++igen) {
+      auto o2gen = dynamic_cast<Generator*>(genList->At(igen));
+      if (o2gen) {
+        o2gen->notifyEmbedding(mEmbedEvent);
+      }
     }
   }
 
diff --git a/run/O2PrimaryServerDevice.h b/run/O2PrimaryServerDevice.h
index ece3747b2c94a..4eecf79cffdb6 100644
--- a/run/O2PrimaryServerDevice.h
+++ b/run/O2PrimaryServerDevice.h
@@ -46,6 +46,7 @@
 #include <chrono>
 #include <CCDB/BasicCCDBManager.h>
 #include <TRandom3.h>
+#include <regex>
 
 namespace o2
 {
@@ -135,6 +136,17 @@ class O2PrimaryServerDevice final : public fair::mq::Device
 
       auto embedinto_filename = conf.getEmbedIntoFileName();
       if (!embedinto_filename.empty()) {
+        // determine the sim prefix from the embedding filename
+        // the filename should be an MCHeader file ... so it should match SOME_PATH/prefix_MCHeader.root
+        std::regex re(R"((.*/)?([^/]+)_MCHeader\.root$)");
+        std::smatch match;
+
+        if (std::regex_search(embedinto_filename, match, re)) {
+          std::cout << "Extracted embedding prefix : " << match[2] << '\n';
+          mEmbeddIntoPrefix = match[2];
+        } else {
+          LOG(fatal) << "Embedding asked but no suitable embedding prefix extractable from " << embedinto_filename;
+        }
         mPrimGen->embedInto(embedinto_filename);
       }
 
@@ -197,6 +209,19 @@ class O2PrimaryServerDevice final : public fair::mq::Device
             auto& vertex = vertices.at(collisionindex);
             LOG(info) << "Setting vertex " << vertex << " for event " << mEventCounter << " for prefix " << mSimConfig.getOutPrefix() << " from CollContext";
             mPrimGen->setExternalVertexForNextEvent(vertex.X(), vertex.Y(), vertex.Z());
+
+            // set correct embedding index for PrimaryGenerator ... based on collision context for embedding
+            auto& collisionParts = mCollissionContext->getEventParts()[collisionindex];
+            int background_index = -1; // -1 means no embedding taking place for this signal
+
+            // find the part that corresponds to the event embeded into
+            for (auto& part : collisionParts) {
+              if (mCollissionContext->getSimPrefixes()[part.sourceID] == mEmbeddIntoPrefix) {
+                background_index = part.entryID;
+                LOG(info) << "Setting embedding index to " << background_index;
+              }
+            }
+            mPrimGen->setEmbedIndex(background_index);
           }
         }
         mPrimGen->GenerateEvent(mStack);
@@ -696,6 +721,7 @@ class O2PrimaryServerDevice final : public fair::mq::Device
   // some information specific to use case when we have a collision context
   o2::steer::DigitizationContext* mCollissionContext = nullptr; //!
   std::unordered_map<int, int> mEventID_to_CollID;              //!
+  std::string mEmbeddIntoPrefix;                                //! sim prefix of background events
 
   TRandom3 mSeedGenerator; //! specific random generator for seed generation for work chunks
 };