feat(undo): Implement comprehensive undo support for visual states

Author: AmirMohammad Cheraghali

src/App.tsx

@@ -188,35 +188,10 @@ function App() {
   const [showIons, setShowIons] = useState(false);
 
   const [showSurface, setShowSurface] = useState(initialUrlState.showSurface || false);
+  const [measurements, setMeasurements] = useState<Measurement[]>([]);
+  const [proteinTitle, setProteinTitle] = useState<string | null>(null);
 
-  // Undo/Redo Stack
-  const visualState: VisualState = useMemo(() => ({
-    representation,
-    coloring,
-    colorPalette,
-    showLigands,
-    showIons,
-    showSurface,
-    customBackgroundColor: customBackgroundColor || '',
-    customColors
-  }), [representation, coloring, colorPalette, showLigands, showIons, showSurface, customBackgroundColor, customColors]);
 
-  const handleVisualStateChange = useCallback((newState: VisualState) => {
-    setRepresentation(newState.representation);
-    setColoring(newState.coloring);
-    setColorPalette(newState.colorPalette);
-    setShowLigands(newState.showLigands);
-    setShowIons(newState.showIons);
-    setShowSurface(newState.showSurface);
-    setCustomBackgroundColor(newState.customBackgroundColor || null);
-    setCustomColors(newState.customColors);
-  }, []);
-
-  const { undo, redo, canUndo, canRedo } = useVisualStack({
-    state: visualState,
-    onChange: handleVisualStateChange,
-    resetTrigger: pdbId // Reset stack when PDB ID changes to avoid cross-structure confusing undos
-  });
 
   // ... (lines 53-343) ...
 
@@ -319,29 +294,75 @@ function App() {
   // Store previous theme to restore after exiting Publication Mode
   const previousThemeRef = useRef(isLightMode);
 
-  // Effect to apply Publication Mode settings
-  useEffect(() => {
-    if (isPublicationMode) {
-      // Save current theme before overriding
-      previousThemeRef.current = isLightMode;
+  const togglePublicationMode = useCallback((shouldBeEnabled?: boolean) => {
+    const nextState = shouldBeEnabled !== undefined ? shouldBeEnabled : !isPublicationMode;
+
+    if (nextState === isPublicationMode) return;
 
-      // Auto-set High Quality Defaults
+    setIsPublicationMode(nextState);
+
+    if (nextState) {
+      // Enable
+      previousThemeRef.current = isLightMode;
       setRepresentation('cartoon');
       setIsCleanMode(true);
       setColoring('chainid');
       setCustomBackgroundColor('#ffffff'); // White background
       setIsLightMode(true); // Ensure light mode for paper look
     } else {
-      // Restore previous configuration
+      // Disable
       setIsCleanMode(false);
       setCustomBackgroundColor(null); // Revert to theme
       setIsLightMode(previousThemeRef.current); // Restore original theme
     }
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [isPublicationMode]);
+  }, [isPublicationMode, isLightMode]);
+
+
+
 
   // ... (fetchTitle logic) ... 
 
+  // Undo/Redo Stack (Moved here to access all state variables)
+  const visualState: VisualState = useMemo(() => ({
+    representation,
+    coloring,
+    colorPalette,
+    showLigands,
+    showIons,
+    showSurface,
+    customBackgroundColor: customBackgroundColor || '',
+    customColors,
+    isSpinning,
+    isCleanMode,
+    showContactMap,
+    isPublicationMode,
+    highlightedResidue,
+    measurements
+  }), [representation, coloring, colorPalette, showLigands, showIons, showSurface, customBackgroundColor, customColors, isSpinning, isCleanMode, showContactMap, isPublicationMode, highlightedResidue, measurements]);
+
+  const handleVisualStateChange = useCallback((newState: VisualState) => {
+    setRepresentation(newState.representation);
+    setColoring(newState.coloring);
+    setColorPalette(newState.colorPalette);
+    setShowLigands(newState.showLigands);
+    setShowIons(newState.showIons);
+    setShowSurface(newState.showSurface);
+    setCustomBackgroundColor(newState.customBackgroundColor || null);
+    setCustomColors(newState.customColors);
+    setIsSpinning(newState.isSpinning);
+    setIsCleanMode(newState.isCleanMode);
+    setShowContactMap(newState.showContactMap);
+    setIsPublicationMode(newState.isPublicationMode);
+    setHighlightedResidue(newState.highlightedResidue);
+    setMeasurements(newState.measurements);
+  }, []);
+
+  const { undo, redo, canUndo, canRedo } = useVisualStack({
+    state: visualState,
+    onChange: handleVisualStateChange,
+    resetTrigger: pdbId
+  });
+
   // --- DERIVED STATE (Dr. AI V4) ---
   const structureStats = useMemo(() => {
     const chainCount = chains.length;
@@ -430,7 +451,7 @@ function App() {
     }
   }, [showLigands, initialUrlState.showLigands, pdbId, dataSource, addToHistory]);
 
-  const [proteinTitle, setProteinTitle] = useState<string | null>(null);
+
 
   // Consolidate Title & Metadata Fetching
   // 1. Remove separate 'fetchTitle' effect that was specific to RCSB PDB
@@ -815,7 +836,7 @@ function App() {
   const [hoveredResidue, setHoveredResidue] = useState<ResidueInfo | null>(null);
 
   // --- MEASUREMENT STATE ---
-  const [measurements, setMeasurements] = useState<Measurement[]>([]);
+
   const [isMeasurementPanelOpen, setIsMeasurementPanelOpen] = useState(false);
   const [measurementTextColorMode, setMeasurementTextColorMode] = useState<MeasurementTextColor>('auto');
 
@@ -968,7 +989,7 @@ function App() {
       label: isPublicationMode ? 'Exit Publication Mode' : 'Enter Publication Mode',
       icon: Camera,
       category: 'View',
-      perform: () => setIsPublicationMode(prev => !prev)
+      perform: () => togglePublicationMode()
     },
     {
       id: 'take-snapshot',
@@ -1344,7 +1365,7 @@ function App() {
             isMeasurementMode={isMeasurementMode}
             setIsMeasurementMode={setIsMeasurementMode}
             isPublicationMode={isPublicationMode}
-            setIsPublicationMode={setIsPublicationMode}
+            onTogglePublicationMode={togglePublicationMode}
             onClearMeasurements={() => {
               setMeasurements([]);
               viewerRef.current?.clearMeasurements();

src/components/Controls.tsx

@@ -277,7 +277,7 @@ interface ControlsProps {
     isMeasurementMode: boolean;
     setIsMeasurementMode: (mode: boolean) => void;
     isPublicationMode: boolean;
-    setIsPublicationMode: (val: boolean) => void;
+    onTogglePublicationMode: () => void;
     onShare: () => void;
     onToggleMeasurement?: () => void;
     onClearMeasurements: () => void;
@@ -368,7 +368,7 @@ export const Controls: React.FC<ControlsProps> = ({
     onToggleMeasurement,
     onClearMeasurements,
     isPublicationMode,
-    setIsPublicationMode,
+    onTogglePublicationMode,
     pdbMetadata,
     isLightMode,
     setIsLightMode,
@@ -1118,7 +1118,7 @@ export const Controls: React.FC<ControlsProps> = ({
                                     {/* Tools: Publication & Measure */}
                                     <div className="pt-2 border-t border-white/5 space-y-2">
                                         <button
-                                            onClick={() => setIsPublicationMode(!isPublicationMode)}
+                                            onClick={onTogglePublicationMode}
                                             className={`w-full flex items-center justify-center gap-2 px-3 py-2 rounded-lg text-xs font-semibold transition-all ${isPublicationMode
                                                 ? 'bg-purple-600 text-white shadow-lg shadow-purple-500/20'
                                                 : (isLightMode ? 'bg-neutral-100 hover:bg-purple-50 text-neutral-600 hover:text-purple-600 border border-neutral-200' : 'bg-neutral-800 hover:bg-neutral-700 text-neutral-300 border border-white/10')

src/hooks/useVisualStack.ts

@@ -1,5 +1,5 @@
 import { useState, useEffect, useCallback, useRef } from 'react';
-import type { RepresentationType, ColoringType, ColorPalette, CustomColorRule } from '../types';
+import type { RepresentationType, ColoringType, ColorPalette, CustomColorRule, ResidueInfo, Measurement } from '../types';
 
 export interface VisualState {
     representation: RepresentationType;
@@ -10,6 +10,12 @@ export interface VisualState {
     showSurface: boolean;
     customBackgroundColor: string;
     customColors: CustomColorRule[];
+    isSpinning: boolean;
+    isCleanMode: boolean;
+    showContactMap: boolean;
+    isPublicationMode: boolean;
+    highlightedResidue: ResidueInfo | null;
+    measurements: Measurement[];
 }
 
 interface UseVisualStackProps {
@@ -55,7 +61,13 @@ export function useVisualStack({ state, onChange, resetTrigger }: UseVisualStack
             prev.showIons !== state.showIons ||
             prev.showSurface !== state.showSurface ||
             prev.customBackgroundColor !== state.customBackgroundColor ||
-            prev.customColors !== state.customColors;
+            prev.customColors !== state.customColors ||
+            prev.isSpinning !== state.isSpinning ||
+            prev.isCleanMode !== state.isCleanMode ||
+            prev.showContactMap !== state.showContactMap ||
+            prev.isPublicationMode !== state.isPublicationMode ||
+            prev.highlightedResidue !== state.highlightedResidue ||
+            prev.measurements !== state.measurements;
 
         if (hasChanged) {
             setPast(prevPast => {