From 5b7b0c093ca6044d0e31d1702198b5b2384ea77c Mon Sep 17 00:00:00 2001
From: tmadlener <thomas.madlener@desy.de>
Date: Wed, 3 Jul 2024 20:30:07 +0200
Subject: [PATCH] Add utility class to reverse RecDqdx - Track relations

---
 utils/CMakeLists.txt                     |  1 +
 utils/include/edm4hep/utils/TrackUtils.h | 27 ++++++++++++++++++++++++
 utils/src/TrackUtils.cc                  | 26 +++++++++++++++++++++++
 3 files changed, 54 insertions(+)
 create mode 100644 utils/include/edm4hep/utils/TrackUtils.h
 create mode 100644 utils/src/TrackUtils.cc

diff --git a/utils/CMakeLists.txt b/utils/CMakeLists.txt
index 9c6a4654d..8c5ff3ee5 100644
--- a/utils/CMakeLists.txt
+++ b/utils/CMakeLists.txt
@@ -11,6 +11,7 @@ target_compile_features(kinematics INTERFACE cxx_std_17)
 set(utils_sources
   src/ParticleIDUtils.cc
   src/MurmurHash3.cpp
+  src/TrackUtils.cc
 )
 
 add_library(utils SHARED ${utils_sources})
diff --git a/utils/include/edm4hep/utils/TrackUtils.h b/utils/include/edm4hep/utils/TrackUtils.h
new file mode 100644
index 000000000..725d537d6
--- /dev/null
+++ b/utils/include/edm4hep/utils/TrackUtils.h
@@ -0,0 +1,27 @@
+#ifndef EDM4HEP_UTILS_TRACKUTILS_H
+#define EDM4HEP_UTILS_TRACKUTILS_H
+
+#include <edm4hep/RecDqdxCollection.h>
+#include <edm4hep/Track.h>
+
+#include <podio/Frame.h>
+
+#include <map>
+#include <vector>
+
+namespace edm4hep::utils {
+/// Utility class to invert the relations between RecDqdx to Track relation
+class TrackPIDHandler {
+  using TrackMapT = std::multimap<edm4hep::Track, edm4hep::RecDqdx>;
+
+  TrackMapT m_trackDqMap{}; ///< The internal map from tracks to RecDqdx
+public:
+  /// Add the information from the passed collection to the handler
+  void addColl(const edm4hep::RecDqdxCollection& coll);
+
+  /// Get all RecDqdx objects for the given track
+  std::vector<edm4hep::RecDqdx> getDqdxValues(const edm4hep::Track& track) const;
+};
+} // namespace edm4hep::utils
+
+#endif // EDM4HEP_UTILS_TRACKUTILS_H
diff --git a/utils/src/TrackUtils.cc b/utils/src/TrackUtils.cc
new file mode 100644
index 000000000..3d5e8eca4
--- /dev/null
+++ b/utils/src/TrackUtils.cc
@@ -0,0 +1,26 @@
+#include "edm4hep/utils/TrackUtils.h"
+#include "edm4hep/RecDqdxCollection.h"
+
+#include <iterator>
+
+namespace edm4hep::utils {
+
+void TrackPIDHandler::addColl(const edm4hep::RecDqdxCollection& coll) {
+  for (const auto dqdx : coll) {
+    m_trackDqMap.emplace(dqdx.getTrack(), dqdx);
+  }
+}
+
+std::vector<edm4hep::RecDqdx> TrackPIDHandler::getDqdxValues(const edm4hep::Track& track) const {
+  const auto& [begin, end] = m_trackDqMap.equal_range(track);
+  std::vector<edm4hep::RecDqdx> dqdxs{};
+  dqdxs.reserve(std::distance(begin, end));
+
+  for (auto it = begin; it != end; ++it) {
+    dqdxs.emplace_back(it->second);
+  }
+
+  return dqdxs;
+}
+
+} // namespace edm4hep::utils