From e1969ed84031e4ab49843d5eab827b71ec507bbf Mon Sep 17 00:00:00 2001
From: Diogo Correia <diogo@metacell.us>
Date: Mon, 18 Nov 2024 15:30:28 +0000
Subject: [PATCH] ESCKAN-75 Filter list of nerves according to neurons.json
 provided by the client

---
 src/context/DataContextProvider.tsx | 73 ++++++++++++++++++++---------
 1 file changed, 50 insertions(+), 23 deletions(-)

diff --git a/src/context/DataContextProvider.tsx b/src/context/DataContextProvider.tsx
index f97041e..5f07f08 100644
--- a/src/context/DataContextProvider.tsx
+++ b/src/context/DataContextProvider.tsx
@@ -1,4 +1,10 @@
-import { PropsWithChildren, useEffect, useMemo, useState } from 'react';
+import {
+  PropsWithChildren,
+  useCallback,
+  useEffect,
+  useMemo,
+  useState,
+} from 'react';
 import { DataContext, Filters, ConnectionSummary } from './DataContext';
 import {
   HierarchicalNode,
@@ -68,27 +74,48 @@ export const DataContextProvider = ({
     return colorMap;
   }, [phenotypes]);
 
-  const updateSelectedConnectionSummary = (
-    summary:
-      | Omit<ConnectionSummary, 'filteredKnowledgeStatements'>
-      | ConnectionSummary
-      | null,
-    filters: Filters,
-    hierarchicalNodes: Record<string, HierarchicalNode>,
-  ) => {
-    if (summary) {
-      const filteredKnowledgeStatements = filterKnowledgeStatements(
-        summary.connections,
-        hierarchicalNodes,
-        filters,
-      );
-      return {
-        ...summary,
-        filteredKnowledgeStatements,
-      };
-    }
-    return null;
-  };
+  const updateSelectedConnectionSummary = useCallback(
+    (
+      summary:
+        | Omit<ConnectionSummary, 'filteredKnowledgeStatements'>
+        | ConnectionSummary
+        | null,
+      filters: Filters,
+      hierarchicalNodes: Record<string, HierarchicalNode>,
+    ) => {
+      if (summary) {
+        let filteredKnowledgeStatements = filterKnowledgeStatements(
+          summary.connections,
+          hierarchicalNodes,
+          filters,
+        );
+
+        filteredKnowledgeStatements = Object.fromEntries(
+          Object.entries(filteredKnowledgeStatements).map(
+            ([key, statement]) => [
+              key,
+              {
+                ...statement,
+                vias: statement.vias.map((via) => ({
+                  ...via,
+                  anatomical_entities: via.anatomical_entities.filter(
+                    (entity) => majorNerves.has(entity.id),
+                  ),
+                })),
+              },
+            ],
+          ),
+        );
+
+        return {
+          ...summary,
+          filteredKnowledgeStatements,
+        };
+      }
+      return null;
+    },
+    [majorNerves],
+  );
 
   const handleSetSelectedConnectionSummary = (
     summary: Omit<ConnectionSummary, 'filteredKnowledgeStatements'> | null,
@@ -105,7 +132,7 @@ export const DataContextProvider = ({
     setSelectedConnectionSummary((prevSummary) =>
       updateSelectedConnectionSummary(prevSummary, filters, hierarchicalNodes),
     );
-  }, [filters, hierarchicalNodes]);
+  }, [filters, hierarchicalNodes, updateSelectedConnectionSummary]);
 
   const dataContextValue = {
     filters,