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 = ({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 = ({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 = ({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 = ({open, onClose, workspace, posi }); }); + // Update the selected neurons in the workspace + workspace.batchUpdate(draft => { + draft.selectedNeurons = newSelectedNeurons; + }); + return {split: newSplit, join: newJoin}; }); onClose();