@weng-lab/genomebrowser-ui 0.1.6 → 0.1.8

package/src/TrackSelect/TrackSelect.tsx

@@ -1,32 +1,30 @@
 import {
   Box,
-  Stack,
-  TextField,
+  Button,
+  Dialog,
+  DialogActions,
+  DialogContent,
+  DialogContentText,
+  DialogTitle,
   FormControlLabel,
+  Stack,
   Switch,
-  Button,
+  TextField,
 } from "@mui/material";
+import { GridRowSelectionModel } from "@mui/x-data-grid";
+import { TreeViewBaseItem } from "@mui/x-tree-view";
+import React, { useEffect, useMemo, useRef, useState } from "react";
 import { DataGridWrapper } from "./DataGrid/DataGridWrapper";
-import { searchTracks, flattenIntoRow } from "./DataGrid/dataGridHelpers";
+import { flattenIntoRow, searchTracks } from "./DataGrid/dataGridHelpers";
 import { TreeViewWrapper } from "./TreeView/TreeViewWrapper";
 import {
   buildSortedAssayTreeView,
   buildTreeView,
   searchTreeItems,
 } from "./TreeView/treeViewHelpers";
-import { SearchTracksProps, ExtendedTreeItemProps, RowInfo } from "./types";
-import { rows, rowById } from "./consts";
-import React, { useState, useMemo, useEffect } from "react";
-import { TreeViewBaseItem } from "@mui/x-tree-view";
-import { GridRowSelectionModel } from "@mui/x-data-grid";
-import {
-  Dialog,
-  DialogTitle,
-  DialogContent,
-  DialogContentText,
-  DialogActions,
-} from "@mui/material";
+import { rowById, rows } from "./consts";
 import { SelectionStoreInstance } from "./store";
+import { ExtendedTreeItemProps, SearchTracksProps } from "./types";
 
 export interface TrackSelectProps {
   store: SelectionStoreInstance;
@@ -37,21 +35,19 @@ export default function TrackSelect({ store }: TrackSelectProps) {
   const [sortedAssay, setSortedAssay] = useState(false);
   const [searchQuery, setSearchQuery] = useState("");
   const [isSearchResult, setIsSearchResult] = useState(false);
-  const selectedTracks = store((s) => s.selectedTracks);
+  const selectedIds = store((s) => s.selectedIds);
+  const getTrackIds = store((s) => s.getTrackIds);
   const setSelected = store((s) => s.setSelected);
   const clear = store((s) => s.clear);
   const MAX_ACTIVE = store((s) => s.maxTracks);
 
-  // Derive active tracks from selectedTracks (all keys are real track IDs)
-  const activeTracks = useMemo(
-    () => new Set(selectedTracks.keys()),
-    [selectedTracks],
-  );
+  // Get only real track IDs (no auto-generated group IDs)
+  const trackIds = useMemo(() => getTrackIds(), [selectedIds, getTrackIds]);
 
   const treeItems = useMemo(() => {
     return sortedAssay
       ? buildSortedAssayTreeView(
-          Array.from(selectedTracks.keys()),
+          Array.from(trackIds),
           {
             id: "1",
             isAssayItem: false,
@@ -63,7 +59,7 @@ export default function TrackSelect({ store }: TrackSelectProps) {
           rowById,
         )
       : buildTreeView(
-          Array.from(selectedTracks.keys()),
+          Array.from(trackIds),
           {
             id: "1",
             isAssayItem: false,
@@ -74,7 +70,7 @@ export default function TrackSelect({ store }: TrackSelectProps) {
           },
           rowById,
         );
-  }, [selectedTracks, sortedAssay]);
+  }, [trackIds, sortedAssay]);
 
   const [filteredRows, setFilteredRows] = useState(rows);
   const [filteredTreeItems, setFilteredTreeItems] = useState([
@@ -87,116 +83,168 @@ export default function TrackSelect({ store }: TrackSelectProps) {
       allRowInfo: [],
     },
   ] as TreeViewBaseItem<ExtendedTreeItemProps>[]);
+  const searchTimeoutRef = useRef<ReturnType<typeof setTimeout> | null>(null);
+  const searchResultIdsRef = useRef<Set<string>>(new Set());
 
   useEffect(() => {
     if (searchQuery === "") {
       setFilteredTreeItems(treeItems);
       setFilteredRows(rows);
       setIsSearchResult(false);
+      searchResultIdsRef.current = new Set();
+    } else if (searchResultIdsRef.current.size > 0) {
+      // When selection changes during search, rebuild tree from selected items that match search
+      const matchingTrackIds = Array.from(trackIds).filter((id) =>
+        searchResultIdsRef.current.has(id),
+      );
+
+      const newTreeItems = sortedAssay
+        ? buildSortedAssayTreeView(
+            matchingTrackIds,
+            {
+              id: "1",
+              isAssayItem: false,
+              label: "Biosamples",
+              icon: "folder",
+              children: [],
+              allRowInfo: [],
+            },
+            rowById,
+          )
+        : buildTreeView(
+            matchingTrackIds,
+            {
+              id: "1",
+              isAssayItem: false,
+              label: "Biosamples",
+              icon: "folder",
+              children: [],
+              allRowInfo: [],
+            },
+            rowById,
+          );
+
+      setFilteredTreeItems(newTreeItems);
     }
-  }, [treeItems, searchQuery]);
+  }, [treeItems, searchQuery, trackIds, sortedAssay]);
 
   const handleToggle = () => {
     setSortedAssay(!sortedAssay);
   };
 
   const handleSearch = (e: React.ChangeEvent<HTMLInputElement>) => {
-    setSearchQuery(e.target.value);
+    const query = e.target.value;
+    setSearchQuery(query);
+
+    // Clear previous timeout
+    if (searchTimeoutRef.current) {
+      clearTimeout(searchTimeoutRef.current);
+    }
 
-    const dataGridSearchProps: SearchTracksProps = {
-      jsonStructure: "tracks",
-      query: e.target.value,
-      keyWeightMap: [
-        "displayname",
-        "ontology",
-        "lifeStage",
-        "sampleType",
-        "type",
-        "experimentAccession",
-        "fileAccession",
-      ],
-    };
+    // Debounce the search
+    searchTimeoutRef.current = setTimeout(() => {
+      if (query === "") {
+        return; // useEffect handles empty query
+      }
 
-    const treeSearchProps: SearchTracksProps = {
-      treeItems: treeItems,
-      query: e.target.value,
-      keyWeightMap: [
-        "displayname",
-        "ontology",
-        "lifeStage",
-        "sampleType",
-        "type",
-        "experimentAccession",
-        "fileAccession",
-      ],
-    };
-    const newDataGridRows = searchTracks(dataGridSearchProps)
-      .map((t) => t.item)
-      .map(flattenIntoRow);
+      const dataGridSearchProps: SearchTracksProps = {
+        jsonStructure: "tracks",
+        query: query,
+        keyWeightMap: [
+          "displayname",
+          "ontology",
+          "lifeStage",
+          "sampleType",
+          "type",
+          "experimentAccession",
+          "fileAccession",
+        ],
+      };
 
-    // we only want the intersection of filtered tracks displayed on the DataGrid and user-selected tracks to be displayed on the tree
-    const newDataGridIds = newDataGridRows.map((r) => r.experimentAccession);
-    const retIds = searchTreeItems(treeSearchProps).map(
-      (r) => r.item.experimentAccession,
-    );
-    const newTreeIds = retIds.filter((i) => newDataGridIds.includes(i));
+      const treeSearchProps: SearchTracksProps = {
+        treeItems: treeItems,
+        query: query,
+        keyWeightMap: [
+          "displayname",
+          "ontology",
+          "lifeStage",
+          "sampleType",
+          "type",
+          "experimentAccession",
+          "fileAccession",
+        ],
+      };
+      const newDataGridRows = searchTracks(dataGridSearchProps)
+        .map((t) => t.item)
+        .map(flattenIntoRow);
 
-    // build new tree from the newTreeIds...maybe it would be faster to prune the current tree instead of rebuilding it?
-    const newTreeItems = sortedAssay
-      ? buildSortedAssayTreeView(
-          newTreeIds,
-          {
-            id: "1",
-            isAssayItem: false,
-            label: "Biosamples",
-            icon: "folder",
-            children: [],
-            allRowInfo: [],
-          },
-          rowById,
-        )
-      : buildTreeView(
-          newTreeIds,
-          {
-            id: "1",
-            isAssayItem: false,
-            label: "Biosamples",
-            icon: "folder",
-            children: [],
-            allRowInfo: [],
-          },
-          rowById,
-        );
+      // we only want the intersection of filtered tracks displayed on the DataGrid and user-selected tracks to be displayed on the tree
+      const newDataGridIds = newDataGridRows.map((r) => r.experimentAccession);
+      const retIds = searchTreeItems(treeSearchProps).map(
+        (r) => r.item.experimentAccession,
+      );
+      const newTreeIds = retIds.filter((i) => newDataGridIds.includes(i));
+
+      // build new tree from the newTreeIds
+      const newTreeItems = sortedAssay
+        ? buildSortedAssayTreeView(
+            newTreeIds,
+            {
+              id: "1",
+              isAssayItem: false,
+              label: "Biosamples",
+              icon: "folder",
+              children: [],
+              allRowInfo: [],
+            },
+            rowById,
+          )
+        : buildTreeView(
+            newTreeIds,
+            {
+              id: "1",
+              isAssayItem: false,
+              label: "Biosamples",
+              icon: "folder",
+              children: [],
+              allRowInfo: [],
+            },
+            rowById,
+          );
+
+      // Store search result IDs in ref for use in useEffect
+      searchResultIdsRef.current = new Set(newDataGridIds);
 
-    setFilteredRows(newDataGridRows);
-    setIsSearchResult(true);
-    setFilteredTreeItems(newTreeItems);
+      setFilteredRows(newDataGridRows);
+      setIsSearchResult(true);
+      setFilteredTreeItems(newTreeItems);
+    }, 300);
   };
 
   const handleSelection = (newSelection: GridRowSelectionModel) => {
-    const idsSet =
+    const allIds: Set<string> =
       (newSelection && (newSelection as any).ids) ?? new Set<string>();
 
-    // Build a Map of only real track IDs (filter out auto-generated group IDs)
-    const newTracks = new Map<string, RowInfo>();
-    idsSet.forEach((id: string) => {
-      const row = rowById.get(id);
-      if (row) {
-        newTracks.set(id, row);
+    // Count only real track IDs for the limit check
+    let realTrackCount = 0;
+    allIds.forEach((id: string) => {
+      if (rowById.has(id)) {
+        realTrackCount++;
       }
     });
 
     // Block only if the new selection would exceed the limit
-    if (newTracks.size > MAX_ACTIVE) {
+    if (realTrackCount > MAX_ACTIVE) {
       setLimitDialogOpen(true);
       return;
     }
 
-    setSelected(newTracks);
+    // Store ALL IDs (including auto-generated group IDs)
+    setSelected(allIds);
   };
 
   return (
-    <Box sx={{ flex: 1 }}>
+    <Box sx={{ flex: 1, pt: 1 }}>
       <Box display="flex" justifyContent="space-between" sx={{ mb: 3 }}>
         <TextField
           id="outlined-suffix-shrink"
@@ -222,7 +270,7 @@ export default function TrackSelect({ store }: TrackSelectProps) {
                 ? `${filteredRows.length} Search Results`
                 : `${rows.length} Available Tracks`
             }
-            selectedTracks={selectedTracks}
+            selectedIds={selectedIds}
             handleSelection={handleSelection}
             sortedAssay={sortedAssay}
           />
@@ -231,8 +279,7 @@ export default function TrackSelect({ store }: TrackSelectProps) {
           <TreeViewWrapper
             store={store}
             items={filteredTreeItems}
-            selectedTracks={selectedTracks}
-            activeTracks={activeTracks}
+            trackIds={trackIds}
             isSearchResult={isSearchResult}
           />
         </Box>

package/src/TrackSelect/TreeView/TreeViewWrapper.tsx

@@ -12,7 +12,7 @@ import { Avatar } from "@mui/material";
 export function TreeViewWrapper({
   store,
   items,
-  activeTracks,
+  trackIds,
   isSearchResult,
 }: TreeViewWrapperProps) {
   const removeIds = store((s) => s.removeIds);
@@ -28,11 +28,13 @@ export function TreeViewWrapper({
       removedIds.forEach((id) => {
         const row = rowById.get(id);
         if (row) {
-          // Add the auto-generated group IDs for this track's ontology and assay
+          // Add the auto-generated group IDs for this track's grouping hierarchy
+          // Default view: ontology -> displayname
           idsToRemove.add(`auto-generated-row-ontology/${row.ontology}`);
           idsToRemove.add(
-            `auto-generated-row-ontology/${row.ontology}-assay/${row.assay}`,
+            `auto-generated-row-ontology/${row.ontology}-displayname/${row.displayname}`,
           );
+          // Sorted by assay view: assay -> ontology -> displayname
           idsToRemove.add(`auto-generated-row-assay/${row.assay}`);
           idsToRemove.add(
             `auto-generated-row-assay/${row.assay}-ontology/${row.ontology}`,
@@ -76,7 +78,7 @@ export function TreeViewWrapper({
             color: "text.primary",
           }}
         >
-          {activeTracks.size}
+          {trackIds.size}
         </Avatar>
         <Typography fontWeight="bold">
           Active Tracks

package/src/TrackSelect/TreeView/treeViewHelpers.tsx

@@ -99,8 +99,10 @@ export function buildSortedAssayTreeView(
         id: row.experimentAccession,
         isAssayItem: false,
         label: row.experimentAccession,
-        icon: row.assay,
+        icon: "removeable",
+        assayName: row.assay,
         children: [],
+        allExpAccessions: [row.experimentAccession],
       };
       sampleAssayMap.set(row.displayname + row.assay, expNode);
       displayNameNode.children!.push(expNode);
@@ -183,8 +185,10 @@ export function buildTreeView(
       expNode = {
         id: row.experimentAccession,
         label: row.experimentAccession,
-        icon: row.assay,
+        icon: "removeable",
+        assayName: row.assay,
         children: [],
+        allExpAccessions: [row.experimentAccession],
       };
       sampleAssayMap.set(row.displayname + row.assay, expNode);
       displayNameNode.children!.push(expNode);
@@ -282,7 +286,7 @@ const TreeItemLabelText = styled(Typography)({
   fontFamily: "inherit",
 });
 
-function CustomLabel({ icon: Icon, children, isAssayItem, ...other }: CustomLabelProps) {
+function CustomLabel({ icon: Icon, children, isAssayItem, assayName, ...other }: CustomLabelProps) {
   const variant = isAssayItem ? "subtitle2" : "body2";
   const fontWeight = isAssayItem ? "bold" : 500;
   return (
@@ -305,8 +309,9 @@ function CustomLabel({ icon: Icon, children, isAssayItem, ...other }: CustomLabe
           sx={{ mr: 1, fontSize: "1.2rem" }}
         />
       )}
-      <Stack direction="row" spacing={2} alignItems="center">
-        { isAssayItem && AssayIcon(other.id) }
+      <Stack direction="row" spacing={1} alignItems="center">
+        {isAssayItem && AssayIcon(other.id)}
+        {assayName && AssayIcon(assayName)}
         <TreeItemLabelText fontWeight={fontWeight} variant={variant}>{children}</TreeItemLabelText>
       </Stack>
     </TreeItemLabel>
@@ -330,14 +335,7 @@ const TreeItemContent = styled("div")(({ theme }) => ({
   marginBottom: theme.spacing(0.5),
   marginTop: theme.spacing(0.5),
   fontWeight: 500,
-  [`&[data-focused], &[data-selected]`]: {
-    backgroundColor: theme.palette.primary.dark,
-    color: theme.palette.primary.contrastText,
-    ...theme.applyStyles("light", {
-      backgroundColor: theme.palette.primary.main,
-    }),
-  },
-  "&:not([data-focused], [data-selected]):hover": {
+  "&:hover": {
     backgroundColor: alpha(theme.palette.primary.main, 0.1),
     color: "white",
     ...theme.applyStyles("light", {
@@ -417,6 +415,7 @@ export const CustomTreeItem = React.forwardRef(function CustomTreeItem(
                 ),
               expandable: (status.expandable && status.expanded).toString(),
               isAssayItem: item.isAssayItem,
+              assayName: item.assayName,
               id: item.id
             })}
           />

package/src/TrackSelect/store.ts

@@ -1,27 +1,50 @@
 import { create, StoreApi, UseBoundStore } from "zustand";
-import { SelectionState, SelectionAction, RowInfo } from "./types";
+import { rowById } from "./consts";
+import { RowInfo, SelectionAction, SelectionState } from "./types";
 
 export type SelectionStoreInstance = UseBoundStore<
   StoreApi<SelectionState & SelectionAction>
 >;
 
-export function createSelectionStore(initialTracks?: Map<string, RowInfo>) {
+// Helper to check if an ID is auto-generated by DataGrid grouping
+const isAutoGeneratedId = (id: string) => id.startsWith("auto-generated-row-");
+
+export function createSelectionStore(initialIds?: Set<string>) {
   return create<SelectionState & SelectionAction>((set, get) => ({
     maxTracks: 30,
-    selectedTracks: initialTracks ? new Map(initialTracks) : new Map<string, RowInfo>(),
-    selectedIds: () => new Set(get().selectedTracks.keys()),
-    setSelected: (tracks: Map<string, RowInfo>) =>
+    // Stores ALL selected IDs, including auto-generated group IDs
+    selectedIds: initialIds ? new Set(initialIds) : new Set<string>(),
+    // Returns only real track IDs (filters out auto-generated group IDs)
+    getTrackIds: () => {
+      const all = get().selectedIds;
+      return new Set([...all].filter((id) => !isAutoGeneratedId(id)));
+    },
+    // Returns a Map of track IDs to RowInfo (no auto-generated IDs)
+    getTrackMap: () => {
+      const all = get().selectedIds;
+      const map = new Map<string, RowInfo>();
+      all.forEach((id) => {
+        if (!isAutoGeneratedId(id)) {
+          const row = rowById.get(id);
+          if (row) {
+            map.set(id, row);
+          }
+        }
+      });
+      return map;
+    },
+    setSelected: (ids: Set<string>) =>
       set(() => ({
-        selectedTracks: new Map(tracks),
+        selectedIds: new Set(ids),
       })),
     removeIds: (removedIds: Set<string>) =>
       set((state) => {
-        const next = new Map(state.selectedTracks);
+        const next = new Set(state.selectedIds);
         removedIds.forEach((id) => {
           next.delete(id);
         });
-        return { selectedTracks: next };
+        return { selectedIds: next };
       }),
-    clear: () => set(() => ({ selectedTracks: new Map<string, RowInfo>() })),
+    clear: () => set(() => ({ selectedIds: new Set<string>() })),
   }));
 }

package/src/TrackSelect/types.ts

@@ -59,6 +59,10 @@ export type ExtendedTreeItemProps = {
   label: string;
   icon: string;
   isAssayItem?: boolean;
+  /**
+   * The assay name for leaf nodes (experiment accession items)
+   */
+  assayName?: string;
   /**
    * list of all the experimentAccession values in the children/grandchildren of the item, or the accession of the item itself
    * this is used in updating the rowSelectionModel when removing items from the Tree View panel
@@ -71,8 +75,7 @@ export type ExtendedTreeItemProps = {
 export type TreeViewWrapperProps = {
   store: SelectionStoreInstance;
   items: TreeViewBaseItem<ExtendedTreeItemProps>[];
-  selectedTracks: Map<string, RowInfo>;
-  activeTracks: Set<string>; // doesn't have the autogenerated row groupings to provide accurate number of tracks
+  trackIds: Set<string>; // real track IDs only (no auto-generated)
   isSearchResult: boolean;
 };
 
@@ -80,6 +83,7 @@ export interface CustomLabelProps {
   id: string;
   children: React.ReactNode;
   isAssayItem?: boolean;
+  assayName?: string;
   icon: React.ElementType | React.ReactElement;
 }
 
@@ -94,12 +98,16 @@ export interface CustomTreeItemProps
  */
 export type SelectionState = {
   maxTracks: number;
-  selectedTracks: Map<string, RowInfo>;
+  // All selected IDs including auto-generated group IDs from DataGrid
+  selectedIds: Set<string>;
 };
 
 export type SelectionAction = {
-  selectedIds: () => Set<string>;
-  setSelected: (tracks: Map<string, RowInfo>) => void;
+  // Returns only real track IDs (filters out auto-generated group IDs)
+  getTrackIds: () => Set<string>;
+  // Returns a Map of track IDs to RowInfo (no auto-generated IDs)
+  getTrackMap: () => Map<string, RowInfo>;
+  setSelected: (ids: Set<string>) => void;
   removeIds: (removedIds: Set<string>) => void;
   clear: () => void;
 };
@@ -126,7 +134,7 @@ interface BaseTableProps extends Omit<DataGridPremiumProps, "columns"> {
 
 type DataGridWrapperProps = {
   rows: RowInfo[];
-  selectedTracks: Map<string, RowInfo>;
+  selectedIds: Set<string>; // all IDs including auto-generated group IDs
   handleSelection: (newSelection: GridRowSelectionModel) => void;
   sortedAssay: boolean;
 };

package/test/main.tsx

@@ -1,13 +1,10 @@
-import TrackSelect from "../src/TrackSelect/TrackSelect"
+import { createSelectionStore } from "../src/TrackSelect/store";
+import TrackSelect from "../src/TrackSelect/TrackSelect";
 import { createRoot } from "react-dom/client";
 
 function Main() {
-    return (
-        <TrackSelect>
-        </TrackSelect>
-    )
+  const store = createSelectionStore();
+  return <TrackSelect store={store} />;
 }
 
-createRoot(document.getElementById("root")!).render(
-  <Main/>
-);
+createRoot(document.getElementById("root")!).render(<Main />);