From 6e4140ffd5bbc7a6448cd8bfa87eca49d1504a11 Mon Sep 17 00:00:00 2001
From: afonso <afonsobspinto@gmail.com>
Date: Thu, 25 Jul 2024 22:26:49 +0100
Subject: [PATCH] CELE-32 fix: Fix selection after split

---
 .../components/viewers/TwoD/ContextMenu.tsx   | 29 ++++++++++++++++++-
 1 file changed, 28 insertions(+), 1 deletion(-)

diff --git a/applications/visualizer/frontend/src/components/viewers/TwoD/ContextMenu.tsx b/applications/visualizer/frontend/src/components/viewers/TwoD/ContextMenu.tsx
index 4ef648ae..9a187427 100644
--- a/applications/visualizer/frontend/src/components/viewers/TwoD/ContextMenu.tsx
+++ b/applications/visualizer/frontend/src/components/viewers/TwoD/ContextMenu.tsx
@@ -67,9 +67,20 @@ const ContextMenu: React.FC<ContextMenuProps> = ({open, onClose, workspace, posi
             const newSplit = new Set(prevState.split);
             const newJoin = new Set(prevState.join);
 
+            const newSelectedNeurons = new Set(workspace.selectedNeurons);
+
             workspace.selectedNeurons.forEach(neuronId => {
                 if (isClass(neuronId, workspace)) {
                     newSplit.add(neuronId);
+                    newSelectedNeurons.delete(neuronId);
+
+                    // Add the individual neurons to the selected neurons
+                    Object.values(workspace.availableNeurons).forEach(neuron => {
+                        if (neuron.nclass === neuronId && neuron.nclass !== neuron.name) {
+                            newSelectedNeurons.add(neuron.name);
+                        }
+                    });
+
                     // Remove the corresponding class from the toJoin set
                     newJoin.forEach(joinNeuronId => {
                         if (workspace.availableNeurons[joinNeuronId].nclass === neuronId) {
@@ -79,6 +90,11 @@ const ContextMenu: React.FC<ContextMenuProps> = ({open, onClose, workspace, posi
                 }
             });
 
+            // Update the selected neurons in the workspace
+            workspace.batchUpdate(draft => {
+                draft.selectedNeurons = newSelectedNeurons;
+            });
+
             return {split: newSplit, join: newJoin};
         });
         onClose();
@@ -89,13 +105,19 @@ const ContextMenu: React.FC<ContextMenuProps> = ({open, onClose, workspace, posi
             const newJoin = new Set(prevState.join);
             const newSplit = new Set(prevState.split);
 
+            const newSelectedNeurons = new Set(workspace.selectedNeurons);
+
             workspace.selectedNeurons.forEach(neuronId => {
                 const neuronClass = workspace.availableNeurons[neuronId].nclass;
+
+                // Remove the individual neurons from the selected neurons and add the class neuron
                 Object.values(workspace.availableNeurons).forEach(neuron => {
-                    if (neuron.nclass === neuronClass && neuron.name !== neuron.nclass) {
+                    if (neuron.nclass === neuronClass && neuron.name !== neuronClass) {
+                        newSelectedNeurons.delete(neuron.name);
                         newJoin.add(neuron.name);
                     }
                 });
+                newSelectedNeurons.add(neuronClass);
 
                 // Remove the corresponding cells from the toSplit set
                 newSplit.forEach(splitNeuronId => {
@@ -105,6 +127,11 @@ const ContextMenu: React.FC<ContextMenuProps> = ({open, onClose, workspace, posi
                 });
             });
 
+            // Update the selected neurons in the workspace
+            workspace.batchUpdate(draft => {
+                draft.selectedNeurons = newSelectedNeurons;
+            });
+
             return {split: newSplit, join: newJoin};
         });
         onClose();