From eb57f57a8a7c1a2fd52aa5f687aeeeb484aa5dea Mon Sep 17 00:00:00 2001
From: tmadlener <thomas.madlener@desy.de>
Date: Fri, 12 Apr 2024 20:22:32 +0200
Subject: [PATCH] Make it possible to add meta information post-hoc

---
 test/utils/test_PIDHandler.cpp                | 33 +++++++++++++++
 utils/include/edm4hep/utils/ParticleIDUtils.h |  7 ++++
 utils/src/ParticleIDUtils.cc                  | 42 +++++++++++--------
 3 files changed, 64 insertions(+), 18 deletions(-)

diff --git a/test/utils/test_PIDHandler.cpp b/test/utils/test_PIDHandler.cpp
index 412bcedb0..1b018a08e 100644
--- a/test/utils/test_PIDHandler.cpp
+++ b/test/utils/test_PIDHandler.cpp
@@ -120,6 +120,39 @@ TEST_CASE("PIDHandler from variadic list of collections", "[pid_utils]") {
   }
 }
 
+TEST_CASE("PIDHandler w/ addMetaInfo", "[pid_utils]") {
+  using namespace edm4hep;
+  auto handler = utils::PIDHandler();
+
+  const auto recoColl = createRecos();
+  auto pidColl1 = createParticleIDs(recoColl, 1.0f);
+  for (auto pid : pidColl1) {
+    pid.setAlgorithmType(42);
+  }
+  const auto pidInfo1 = utils::ParticleIDMeta{"fancyAlgo", 42, {"p1", "p2"}};
+
+  handler.addColl(pidColl1, pidInfo1);
+
+  REQUIRE(handler.getAlgoType("fancyAlgo").value_or(0) == 42);
+  REQUIRE(handler.getParamIndex(42, "p2").value_or(-1) == 1);
+  REQUIRE(handler.getPID(recoColl[0], 42).value() == pidColl1[0]);
+
+  // Technically, we can even just add meta data without having a corresponding
+  // ParticleID collection to match
+  handler.addMetaInfo(utils::ParticleIDMeta{"anotherAlgo", 123, {}});
+  REQUIRE(handler.getAlgoType("anotherAlgo").value() == 123);
+
+  // Expected exceptions also get thrown
+  REQUIRE_THROWS_AS(handler.addMetaInfo(utils::ParticleIDMeta{"anotherAlgo", 321, {"param"}}), std::runtime_error);
+  // No information about this meta data can be obtained
+  REQUIRE_FALSE(handler.getParamIndex(321, "param").has_value());
+
+  REQUIRE_THROWS_AS(handler.addMetaInfo(utils::ParticleIDMeta{"newAlgo", 42, {"PARAM"}}), std::runtime_error);
+  // Existing meta info is unchanged
+  REQUIRE_FALSE(handler.getParamIndex(42, "PARAM").has_value());
+  REQUIRE(handler.getParamIndex(42, "p2").value_or(-1) == 1);
+}
+
 TEST_CASE("PIDHandler from Frame w/ metadata", "[pid_utils]") {
   using namespace edm4hep;
   const auto& [event, metadata] = createEventAndMetadata();
diff --git a/utils/include/edm4hep/utils/ParticleIDUtils.h b/utils/include/edm4hep/utils/ParticleIDUtils.h
index a1ceabb19..a9fa18e21 100644
--- a/utils/include/edm4hep/utils/ParticleIDUtils.h
+++ b/utils/include/edm4hep/utils/ParticleIDUtils.h
@@ -58,6 +58,13 @@ class PIDHandler {
   /// Add the information from one ParticleIDCollection to the handler
   void addColl(const edm4hep::ParticleIDCollection& coll);
 
+  /// Add the information from one ParticleIDCollection to the handler together
+  /// with its meta data
+  void addColl(const edm4hep::ParticleIDCollection& coll, const edm4hep::utils::ParticleIDMeta& pidInfo);
+
+  /// Add meta information for a collection
+  void addMetaInfo(const edm4hep::utils::ParticleIDMeta& pidInfo);
+
   /// Retrieve all ParticleIDs that are related to the passed
   /// ReconstructedParticle
   std::vector<edm4hep::ParticleID> getPIDs(const edm4hep::ReconstructedParticle& reco) const;
diff --git a/utils/src/ParticleIDUtils.cc b/utils/src/ParticleIDUtils.cc
index 70d14b947..c24b98458 100644
--- a/utils/src/ParticleIDUtils.cc
+++ b/utils/src/ParticleIDUtils.cc
@@ -1,7 +1,6 @@
 #include <edm4hep/utils/ParticleIDUtils.h>
 
 #include "edm4hep/Constants.h"
-#include "edm4hep/ReconstructedParticle.h"
 
 #include <podio/FrameCategories.h>
 
@@ -18,6 +17,27 @@ void PIDHandler::addColl(const edm4hep::ParticleIDCollection& coll) {
   }
 }
 
+void PIDHandler::addColl(const edm4hep::ParticleIDCollection& coll, const edm4hep::utils::ParticleIDMeta& pidInfo) {
+  addColl(coll);
+  addMetaInfo(pidInfo);
+}
+
+void PIDHandler::addMetaInfo(const edm4hep::utils::ParticleIDMeta& pidInfo) {
+  const auto [algoIt, inserted] = m_algoTypes.emplace(pidInfo.algoName, pidInfo.algoType);
+  if (!inserted) {
+    throw std::runtime_error("Cannot have duplicate algorithm names (" + pidInfo.algoName + " already exists)");
+  }
+
+  const auto [__, parInserted] = m_algoParamNames.emplace(pidInfo.algoType, pidInfo.paramNames);
+  if (!parInserted) {
+    if (inserted) {
+      m_algoTypes.erase(algoIt);
+    }
+    throw std::runtime_error("Cannot have duplicate algorithm types (" + std::to_string(pidInfo.algoType) +
+                             " already exists)");
+  }
+}
+
 PIDHandler PIDHandler::from(const podio::Frame& event, const podio::Frame& metadata) {
   PIDHandler handler{};
   for (const auto& name : event.getAvailableCollections()) {
@@ -25,23 +45,9 @@ PIDHandler PIDHandler::from(const podio::Frame& event, const podio::Frame& metad
     if (const auto pidColl = dynamic_cast<const edm4hep::ParticleIDCollection*>(coll)) {
       handler.addColl(*pidColl);
 
-      const auto& algoName =
-          metadata.getParameter<std::string>(podio::collMetadataParamName(name, edm4hep::pidAlgoName));
-      const auto algoType = metadata.getParameter<int>(podio::collMetadataParamName(name, edm4hep::pidAlgoType));
-
-      // Here we have to assume that an empty name and an algoType of 0 simply
-      // mean that this has not been set at all
-      if (!algoName.empty() && algoType != 0) {
-        const auto [_, inserted] = handler.m_algoTypes.emplace(algoName, algoType);
-        if (!inserted) {
-          throw std::runtime_error("Cannot have duplicate algorithm names");
-        }
-
-        const auto& paramNames = metadata.getParameter<std::vector<std::string>>(
-            podio::collMetadataParamName(name, edm4hep::pidParameterNames));
-        if (!paramNames.empty()) {
-          handler.m_algoParamNames.emplace(algoType, paramNames);
-        }
+      auto maybeMetaInfo = PIDHandler::getAlgoInfo(metadata, name);
+      if (maybeMetaInfo) {
+        handler.addMetaInfo(std::move(maybeMetaInfo.value()));
       }
     }
   }