From d6045c49f2246df900f5e81b34b1e2064ed4ebf2 Mon Sep 17 00:00:00 2001
From: Diana <70915994+diana0x0f@users.noreply.github.com>
Date: Sun, 5 Jan 2025 18:14:42 +0100
Subject: [PATCH] [PWGUD] Fix global forward tracks selection and store tan
 lambda information (#9185)

Co-authored-by: Diana Krupova <diana@Dianas-MacBook-Pro.local>
---
 PWGUD/TableProducer/UPCCandidateProducer.cxx | 34 +++++++++++++-------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/PWGUD/TableProducer/UPCCandidateProducer.cxx b/PWGUD/TableProducer/UPCCandidateProducer.cxx
index f93bea732b4..360f82113bb 100644
--- a/PWGUD/TableProducer/UPCCandidateProducer.cxx
+++ b/PWGUD/TableProducer/UPCCandidateProducer.cxx
@@ -1600,28 +1600,41 @@ struct UpcCandProducer {
 
     std::vector<int> selTrackIdsGlobal{};
 
-    // storing n-prong matches
     int32_t candID = 0;
-    auto midIt = bcsMatchedTrIdsMID.begin();
+
     for (auto& pair : bcsMatchedTrIdsGlobal) { // candidates with MFT
       auto globalBC = static_cast<int64_t>(pair.first);
-      const auto& fwdTrackIDs = pair.second;
+      const auto& fwdTrackIDs = pair.second; // Forward tracks
       uint32_t nMFTs = fwdTrackIDs.size();
-      if (nMFTs > fNFwdProngs) // too many tracks
-        continue;
       std::vector<int64_t> trkCandIDs{};
-      const auto& midTrackIDs = midIt->second; // to retrieve corresponding MCH-MID tracks
+
+      // Find corresponding midTrackIDs using std::find_if
+      auto midIt = std::find_if(bcsMatchedTrIdsMID.begin(), bcsMatchedTrIdsMID.end(),
+                                [globalBC](const auto& midPair) {
+                                  return midPair.first == static_cast<uint64_t>(globalBC);
+                                });
+
+      const auto* midTrackIDs = (midIt != bcsMatchedTrIdsMID.end()) ? &midIt->second : nullptr;
+
+      if (nMFTs > fNFwdProngs) // Skip if too many tracks
+        continue;
+
       if (nMFTs == fNFwdProngs) {
+
         for (auto iMft : fwdTrackIDs) {
           auto trk = fwdTracks.iteratorAt(iMft);
           auto trkEta = trk.eta();
-          if (trkEta > fMinEtaMFT && trkEta < fMaxEtaMFT) { // If the track is in the MFT acceptance, store the global track
-            trkCandIDs.insert(trkCandIDs.end(), fwdTrackIDs.begin(), fwdTrackIDs.end());
-          } else { // If the track is not in the MFT acceptance, store the MCH-MID track
-            trkCandIDs.insert(trkCandIDs.end(), midTrackIDs.begin(), midTrackIDs.end());
+
+          if (trkEta > fMinEtaMFT && trkEta < fMaxEtaMFT) {
+            // Track is within MFT acceptance, store forward track IDs
+            trkCandIDs.push_back(iMft);
+          } else if (midTrackIDs) {
+            // Track is outside MFT acceptance, store corresponding MID track IDs
+            trkCandIDs.insert(trkCandIDs.end(), midTrackIDs->begin(), midTrackIDs->end());
           }
         }
       }
+
       uint64_t closestBcMCH = 0;
       upchelpers::FITInfo fitInfo{};
       fitInfo.timeFT0A = -999.f;
@@ -1732,7 +1745,6 @@ struct UpcCandProducer {
                              amplitudesV0A,
                              relBCsV0A);
       candID++;
-      midIt++;
       trkCandIDs.clear();
     }