Show circuit count when filtering by brain region (top menu)

src/components/explore-section/Circuit/ListView/ExploreCircuitTable.tsx

@@ -1,11 +1,11 @@
 'use client';
 
-import { Key, useMemo, useState } from 'react';
+import { Key, useEffect, useMemo, useState } from 'react';
 import { Table, Tooltip } from 'antd';
 import { TableRowSelection } from 'antd/es/table/interface';
 import { InfoCircleOutlined } from '@ant-design/icons';
-import { useAtomValue } from 'jotai';
-import circuitsFlat from '../content/circuits_flat';
+import { useAtomValue, useSetAtom } from 'jotai';
+import circuitsFlat, { circuitCountAtom } from '../content/circuits_flat';
 import { ArrowSmall } from '../icon/ArrowSubcircuitIcon';
 import CIRCUITS from '../content/circuits_tree';
 import { CircuitColumn, CircuitSchemaProps } from '../type';
@@ -26,23 +26,50 @@ const getExpandableRowKeys = (data: CircuitSchemaProps[]): string[] => {
   }, [] as string[]);
 };
 
-const brainRegionFilter = memoize(
-  (brainRegions: { brainRegionId: string | undefined; brainRegionSet: Set<string> }) =>
-    memoize((node: CircuitSchemaProps): boolean => {
-      return (
-        brainRegions.brainRegionSet.has(node.brainRegion.trim().toLocaleLowerCase()) ||
-        (node.subcircuits?.some((sc) => brainRegionFilter(brainRegions)(sc)) ?? false)
-      );
-    }),
-  (a) => `${a.brainRegionId}`
-);
+function brainRegionFilterFunction({
+  circuits,
+  regionSet,
+}: {
+  circuits: CircuitSchemaProps[];
+  region: string | undefined;
+  regionSet: Set<string>;
+}) {
+  let count = 0;
+
+  function recurse(circuit: CircuitSchemaProps): CircuitSchemaProps | null {
+    const brainRegionMatches = regionSet.has(circuit.brainRegion.trim().toLowerCase());
+
+    const filteredSubs = (circuit.subcircuits?.map(recurse).filter(Boolean) ??
+      []) as CircuitSchemaProps[];
+
+    if (brainRegionMatches || filteredSubs.length > 0) {
+      count += 1;
+      return {
+        ...circuit,
+        subcircuits: filteredSubs,
+      };
+    }
+
+    return null;
+  }
+
+  const filteredTree = circuits.map(recurse).filter(Boolean) as CircuitSchemaProps[];
+
+  return {
+    filteredTree,
+    count,
+  };
+}
+
+const brainRegionFilter = memoize(brainRegionFilterFunction, (param) => param.region ?? '');
 
 export default function ExploreCircuitTable() {
   const [expandedRowKeys, setExpandedRowKeys] = useState<string[]>(getExpandableRowKeys(CIRCUITS));
   const [selectedRowKeys, setSelectedRowKeys] = useState<Key[]>([]);
   const selectedBrainRegion = useAtomValue(selectedBrainRegionAtom);
   const selectedBrainRegions = useAtomValue(selectedBrainRegionWithDescendantsAndAncestorsAtom);
   const brainRegionNotationById = useAtomValue(brainRegionByIdMapAtom);
+  const setCircuitCount = useSetAtom(circuitCountAtom);
 
   const brainRegionSet = useMemo(() => {
     return new Set(
@@ -52,7 +79,15 @@ export default function ExploreCircuitTable() {
     );
   }, [selectedBrainRegions, brainRegionNotationById]);
 
-  const filter = brainRegionFilter({ brainRegionId: selectedBrainRegion?.id, brainRegionSet });
+  const filteredCircuits = brainRegionFilter({
+    region: selectedBrainRegion?.id,
+    regionSet: brainRegionSet,
+    circuits: CIRCUITS,
+  });
+
+  useEffect(() => {
+    setCircuitCount(filteredCircuits.count);
+  }, [filteredCircuits.count, setCircuitCount]);
 
   const handleExpandRow = (row: CircuitSchemaProps, _index: number) => {
     if (!row.hasSubcircuits) return;
@@ -162,8 +197,7 @@ export default function ExploreCircuitTable() {
 
   // SUBCIRCUIT TABLE - LEVEL 1
   const expandedRowRender = (circuit: CircuitSchemaProps) => {
-    const subcircuits = circuit.subcircuits?.filter((sc) => filter(sc));
-    if (!subcircuits || subcircuits.length === 0) return null;
+    if (!circuit.subcircuits || circuit.subcircuits.length === 0) return null;
 
     return (
       <div className="relative flex flex-col">
@@ -189,7 +223,7 @@ export default function ExploreCircuitTable() {
             styles.circuitTable
           )}
           columns={columns}
-          dataSource={circuit.subcircuits?.filter((sc) => filter(sc)) || []}
+          dataSource={circuit.subcircuits}
           pagination={false}
           rowSelection={rowSelection}
           expandable={{
@@ -234,7 +268,7 @@ export default function ExploreCircuitTable() {
           styles.circuitTable
         )}
         style={{ '--ant-table-expand-icon-col-width': '0px' } as React.CSSProperties}
-        dataSource={CIRCUITS.filter((c) => filter(c))}
+        dataSource={filteredCircuits.filteredTree}
         columns={columns}
         pagination={false}
         rowSelection={rowSelection}

src/components/explore-section/Circuit/content/circuits_flat.ts

@@ -1,3 +1,4 @@
+import { atom } from 'jotai';
 import { CircuitSchemaProps } from '../type';
 import CIRCUITS from './circuits_tree';
 
@@ -8,4 +9,8 @@ export const flattenRows = (data: CircuitSchemaProps[]): CircuitSchemaProps[] =>
   }, [] as CircuitSchemaProps[]);
 };
 
-export default flattenRows(CIRCUITS);
+const circuitsFlat = flattenRows(CIRCUITS);
+const circuitCountAtom = atom(circuitsFlat.length);
+
+export default circuitsFlat;
+export { circuitCountAtom };

src/components/explore-section/ExploreListingLayout/index.tsx

@@ -8,7 +8,7 @@ import { useParams, usePathname, useRouter } from 'next/navigation';
 import { CSSProperties, ReactNode } from 'react';
 import { ErrorBoundary } from 'react-error-boundary';
 
-import circuitsFlat from '../Circuit/content/circuits_flat';
+import { circuitCountAtom } from '../Circuit/content/circuits_flat';
 
 import SimpleErrorComponent from '@/components/GenericErrorFallback';
 import BackToInteractiveExplorationBtn from '@/components/explore-section/BackToInteractiveExplorationBtn';
@@ -24,7 +24,6 @@ import { totalByExperimentAndRegionsAtom } from '@/state/explore-section/list-vi
 import { ExploreDataScope } from '@/types/explore-section/application';
 import { VirtualLabInfo } from '@/types/virtual-lab/common';
 import { ensureString } from '@/util/type-guards';
-import { classNames } from '@/util/utils';
 import { DataTypeGroup } from '@/types/explore-section/data-types';
 
 const dataScope = ExploreDataScope.SelectedBrainRegion;
@@ -78,6 +77,7 @@ export default function ExploreListingLayout({
 
   const showCircuitMenu = dataTypeGroup === DataTypeGroup.ModelData;
   const activePath = pathname?.split('/').pop() || 'morphology';
+  const circuitCount = useAtomValue(circuitCountAtom);
 
   const onClick: MenuProps['onClick'] = async (info) => {
     const { key, domEvent } = info;
@@ -132,7 +132,7 @@ export default function ExploreListingLayout({
     items.push({
       key: 'circuit',
       title: 'Circuit',
-      label: `Circuit (${circuitsFlat.length})`,
+      label: `Circuit (${circuitCount})`,
       className: 'text-center font-semibold',
       style: {
         backgroundColor: circuitActive ? 'white' : '#002766',