From dd2821065e654e4b4bac0b0c3dc0f6d7cbadd1b7 Mon Sep 17 00:00:00 2001
From: Braulio Rivas Abad <brauliorivas2004@gmail.com>
Date: Thu, 15 Aug 2024 15:34:11 -0500
Subject: [PATCH] reconnect mcparticle tree, go to level where objects aren't
 filtered

---
 js/filters/reconnect/mcparticletree.js | 121 +++++++++++--------------
 1 file changed, 55 insertions(+), 66 deletions(-)

diff --git a/js/filters/reconnect/mcparticletree.js b/js/filters/reconnect/mcparticletree.js
index d92828f..58395c2 100644
--- a/js/filters/reconnect/mcparticletree.js
+++ b/js/filters/reconnect/mcparticletree.js
@@ -1,89 +1,78 @@
 import { linkTypes } from "../../types/links.js";
 
-const findParentRow = (object, uniqueRows, rowToIndex) => {
-  const thisRowIndex = rowToIndex[object.row];
-  if (thisRowIndex > 0 && thisRowIndex < uniqueRows.length) {
-    return uniqueRows[thisRowIndex - 1];
+const findParticles = (otherObject, relationName, ids) => {
+  let oneToManyRelations;
+  if (otherObject.relations) {
+    oneToManyRelations = otherObject.relations.oneToManyRelations;
+  } else {
+    oneToManyRelations = otherObject.oneToManyRelations;
   }
-  return NaN;
-};
-
-const findDaughterRow = (object, uniqueRows, rowToIndex) => {
-  const thisRowIndex = rowToIndex[object.row];
-  if (thisRowIndex >= 0 && thisRowIndex < uniqueRows.length - 1) {
-    return uniqueRows[thisRowIndex + 1];
-  }
-  return NaN;
-};
 
-export function reconnectMCParticleTree(viewCurrentObjects, ids) {
-  const { collection, oneToMany } =
-    viewCurrentObjects.datatypes["edm4hep::MCParticle"];
+  const relations = oneToManyRelations[relationName];
 
-  const sortedCollection = collection.sort((a, b) => a.row - b.row);
+  if (relations.length === 0) return [];
 
-  // const beginRowsIndex = {};
-  // sortedCollection.forEach((object, index) => {
-  //   if (beginRowsIndex[object.row] === undefined) {
-  //     beginRowsIndex[object.row] = index;
-  //   }
-  // });
+  let hasAny = 0;
 
-  // const rows = sortedCollection.map((object) => object.row);
-  // const uniqueRows = [...new Set(rows)];
+  relations.forEach((object) =>
+    ids.has(`${object.index}-${object.collectionId}`) ? (hasAny += 1) : null
+  );
 
-  // const rowToIndex = {};
-  // for (const [index, row] of uniqueRows.entries()) {
-  //   rowToIndex[row] = index;
-  // }
-
-  // const rowToObjectsCount = {};
+  return hasAny > 0
+    ? relations
+    : relations
+        .map((parentLink) => findParticles(parentLink.to, relationName, ids))
+        .flat();
+};
 
-  // sortedCollection.forEach((object) => {
-  //   if (rowToObjectsCount[object.row] === undefined) {
-  //     rowToObjectsCount[object.row] = 1;
-  //     return;
-  //   }
-  //   rowToObjectsCount[object.row] += 1;
-  // });
+export function reconnectMCParticleTree(viewCurrentObjects, ids) {
+  const { collection, oneToMany } =
+    viewCurrentObjects.datatypes["edm4hep::MCParticle"];
 
-  for (const object of sortedCollection) {
+  for (const object of collection) {
     const { oneToManyRelations } = object;
     object.saveRelations();
 
+    const parentRelations = oneToManyRelations["parents"];
+    const daughterRelations = oneToManyRelations["daughters"];
+
     object.oneToManyRelations = {
       "parents": [],
       "daughters": [],
     };
 
-    // const parentRow = findParentRow(object, uniqueRows, rowToIndex);
-    // if (parentRow !== NaN) {
-    //   const beginIndex = beginRowsIndex[parentRow];
-    //   const endIndex = beginIndex + rowToObjectsCount[parentRow];
+    for (const parentRelation of parentRelations) {
+      const parent = parentRelation.to;
+      const parentIndex = `${parent.index}-${parent.collectionId}`;
 
-    //   for (let i = beginIndex; i < endIndex; i++) {
-    //     const newParentLink = new linkTypes["parents"](
-    //       object,
-    //       sortedCollection[i]
-    //     );
-    //     object.oneToManyRelations["parents"].push(newParentLink);
-    //     oneToMany["parents"].push(newParentLink);
-    //   }
-    // }
+      if (ids.has(parentIndex)) {
+        object.oneToManyRelations["parents"].push(parentRelation);
+        oneToMany["parents"].push(parentRelation);
+      } else {
+        const newParents = findParticles(parent, "parents", ids);
+        for (const newParent of newParents) {
+          const link = new linkTypes["parents"](object, newParent);
+          object.oneToManyRelations["parents"].push(link);
+          oneToMany["parents"].push(link);
+        }
+      }
+    }
 
-    // const daughterRow = findDaughterRow(object, uniqueRows, rowToIndex);
-    // if (daughterRow !== NaN) {
-    //   const beginIndex = beginRowsIndex[daughterRow];
-    //   const endIndex = beginIndex + rowToObjectsCount[daughterRow];
+    for (const daughterRelation of daughterRelations) {
+      const daughter = daughterRelation.to;
+      const daughterIndex = `${daughter.index}-${daughter.collectionId}`;
 
-    //   for (let i = beginIndex; i < endIndex; i++) {
-    //     const newDaughterLink = new linkTypes["daughters"](
-    //       object,
-    //       sortedCollection[i]
-    //     );
-    //     object.oneToManyRelations["daughters"].push(newDaughterLink);
-    //     oneToMany["daughters"].push(newDaughterLink);
-    //   }
-    // }
+      if (ids.has(daughterIndex)) {
+        object.oneToManyRelations["daughters"].push(daughterRelation);
+        oneToMany["daughters"].push(daughterRelation);
+      } else {
+        const newDaughters = findParticles(daughter, "daughters", ids);
+        for (const newDaughter of newDaughters) {
+          const link = new linkTypes["daughters"](object, newDaughter);
+          object.oneToManyRelations["daughters"].push(link);
+          oneToMany["daughters"].push(link);
+        }
+      }
+    }
   }
 }