Refactor molecule input handling to implement two-way synchronization and replace ReducedRouteCGR with SB-CGR, enhancing state management and visualization consistency.

synplan/interfaces/gui.py

@@ -26,6 +26,7 @@ from synplan.utils.visualisation import (
     generate_results_html,
     html_top_routes_cluster,
     get_route_svg,
+    get_route_svg_from_json
 )
 from synplan.utils.config import TreeConfig, PolicyNetworkConfig
 from synplan.utils.loading import load_reaction_rules, load_building_blocks
@@ -169,8 +170,10 @@ def initialize_app():
         st.session_state.num_clusters_setting = 10
     if "route_cgrs_dict" not in st.session_state:
         st.session_state.route_cgrs_dict = None
-    if "r_route_cgrs_dict" not in st.session_state:
-        st.session_state.r_route_cgrs_dict = None
+    if "sb_cgrs_dict" not in st.session_state:
+        st.session_state.sb_cgrs_dict = None
+    if "route_json" not in st.session_state:
+        st.session_state.route_json = None
 
     # Subclustering state
     if "subclustering_done" not in st.session_state:
@@ -219,7 +222,7 @@ def setup_sidebar():
 
 
 def handle_molecule_input():
-    """3. Molecule Input: Managing the input area for molecule data."""
+    """3. Molecule Input: Managing the input area for molecule data with two-way synchronization."""
     st.header("Molecule input")
     st.markdown(
         """
@@ -228,42 +231,60 @@ def handle_molecule_input():
         * Draw it + Apply
         """
     )
-    # Use st.session_state.ketcher to persist drawn molecule
-    molecule_text_input = st.text_input(
-        "SMILES:", value=st.session_state.ketcher, key="smiles_text_input_key"
+
+    if "shared_smiles" not in st.session_state:
+        st.session_state.shared_smiles = st.session_state.get("ketcher", DEFAULT_MOL)
+
+    if "ketcher_render_count" not in st.session_state:
+        st.session_state.ketcher_render_count = 0
+
+    def text_input_changed_callback():
+        new_text_value = (
+            st.session_state.smiles_text_input_key_for_sync
+        )  # Key of the text_input
+        if new_text_value != st.session_state.shared_smiles:
+            st.session_state.shared_smiles = new_text_value
+            st.session_state.ketcher = new_text_value
+            st.session_state.ketcher_render_count += 1
+
+    # SMILES Text Input
+    st.text_input(
+        "SMILES:",
+        value=st.session_state.shared_smiles,
+        key="smiles_text_input_key_for_sync",  # Unique key for this widget
+        on_change=text_input_changed_callback,
+        help="Enter SMILES string and press Enter. The drawing will update, and vice-versa.",
     )
 
-    smile_code_ketcher = st_ketcher(molecule_text_input, key="ketcher_widget")
-    # col_kethcer, col_info = st.columns([0.8, 0.2])
-    # with col_kethcer:
-    #     smile_code_ketcher = st_ketcher(molecule_text_input, key="ketcher_widget")
-    # with col_info:
-    #     st.subheader("Synthetic Complexity")
-    #     sascore = ()
-    #     st.markdown(f"SAScore: {sascore}")
-    #     syba_score = ()
-    #     st.markdown(f"SYBA: {sascore}")
-
-    current_smile_code = (
-        smile_code_ketcher  # The output from ketcher is the definitive SMILES
+    ketcher_key = f"ketcher_widget_for_sync_{st.session_state.ketcher_render_count}"
+    smile_code_output_from_ketcher = st_ketcher(
+        st.session_state.shared_smiles, key=ketcher_key
     )
 
+    if smile_code_output_from_ketcher != st.session_state.shared_smiles:
+        st.session_state.shared_smiles = smile_code_output_from_ketcher
+        st.session_state.ketcher = smile_code_output_from_ketcher
+        st.rerun()
+
+    current_smiles_for_planning = st.session_state.shared_smiles
+
+    last_planned_smiles = st.session_state.get("target_smiles")
     if (
-        "target_smiles" in st.session_state
-        and current_smile_code != st.session_state.target_smiles
+        last_planned_smiles
+        and current_smiles_for_planning != last_planned_smiles
+        and st.session_state.get("planning_done", False)
     ):
-        st.warning("Molecule structure changed. Please re-run planning.")
-        st.session_state.planning_done = False
-        st.session_state.clustering_done = False
-        st.session_state.subclustering_done = False
-        st.session_state.tree = None
-        st.session_state.res = None
-        st.session_state.clusters = None
-        st.session_state.reactions_dict = None
-        st.session_state.subclusters = None
-        st.session_state.ketcher = current_smile_code
+        st.warning(
+            "Molecule structure has changed since the last successful planning run. "
+            "Results shown below (if any) are for the previous molecule. "
+            "Please re-run planning for the current structure."
+        )
 
-    return current_smile_code
+    # Ensure st.session_state.ketcher is consistent for other parts of the app
+    if st.session_state.get("ketcher") != current_smiles_for_planning:
+        st.session_state.ketcher = current_smiles_for_planning
+
+    return current_smiles_for_planning
 
 
 def setup_planning_options():
@@ -272,9 +293,7 @@ def setup_planning_options():
     st.markdown(
         """If you modified the structure, please ensure you clicked on `Apply` (bottom right of the molecular editor)."""
     )
-    # This smile_code display will be updated if handle_molecule_input has run and returned a new smile_code
-    # However, to display it correctly, we need the current smile_code from the session or input handler.
-    # For simplicity, let's assume handle_molecule_input has updated st.session_state.ketcher
+
     st.markdown(
         f"The molecule SMILES is actually: ``{st.session_state.get('ketcher', DEFAULT_MOL)}``"
     )
@@ -361,7 +380,8 @@ def setup_planning_options():
         st.session_state.reactions_dict = None
         st.session_state.subclusters = None
         st.session_state.route_cgrs_dict = None
-        st.session_state.r_route_cgrs_dict = None
+        st.session_state.sb_cgrs_dict = None
+        st.session_state.route_json = None
         active_smile_code = st.session_state.get(
             "ketcher", DEFAULT_MOL
         )  # Get current SMILES
@@ -370,7 +390,7 @@ def setup_planning_options():
         )
 
         try:
-            target_molecule = mol_from_smiles(active_smile_code)
+            target_molecule = mol_from_smiles(active_smile_code, clean_stereo=True)
             if target_molecule is None:
                 st.error(f"Could not parse the input SMILES: {active_smile_code}")
             else:
@@ -419,7 +439,7 @@ def setup_planning_options():
 
                     mcts_progress_text = "Running MCTS iterations..."
                     mcts_bar = st.progress(0, text=mcts_progress_text)
-                    for step, (solved, node_id) in enumerate(tree):
+                    for step, (solved, route_id) in enumerate(tree):
                         progress_value = min(
                             1.0, (step + 1) / planning_params["max_iterations"]
                         )
@@ -464,34 +484,35 @@ def display_planning_results():
 
             st.subheader("Examples of found retrosynthetic routes")
             image_counter = 0
-            visualised_node_ids = set()
+            visualised_route_ids = set()
 
             if not winning_nodes:
                 st.warning(
                     "Planning solved, but no winning nodes found in the tree object."
                 )
             else:
-                for n, node_id in enumerate(winning_nodes):
+                for n, route_id in enumerate(winning_nodes):
                     if image_counter >= 3:
                         break
-                    if node_id not in visualised_node_ids:
+                    if route_id not in visualised_route_ids:
                         try:
-                            visualised_node_ids.add(node_id)
-                            num_steps = len(tree.synthesis_route(node_id))
-                            route_score = round(tree.route_score(node_id), 3)
-                            svg = get_route_svg(tree, node_id)
+                            visualised_route_ids.add(route_id)
+                            num_steps = len(tree.synthesis_route(route_id))
+                            route_score = round(tree.route_score(route_id), 3)
+                            svg = get_route_svg(tree, route_id)
+                            # svg = get_route_svg_from_json(st.session_state.route_json, route_id)
                             if svg:
                                 st.image(
                                     svg,
-                                    caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}",
+                                    caption=f"Route {route_id}; {num_steps} steps; Route score: {route_score}",
                                 )
                                 image_counter += 1
                             else:
                                 st.warning(
-                                    f"Could not generate SVG for route {node_id}."
+                                    f"Could not generate SVG for route {route_id}."
                                 )
                         except Exception as e:
-                            st.error(f"Error displaying route {node_id}: {e}")
+                            st.error(f"Error displaying route {route_id}: {e}")
         else:  # Not solved
             st.header("Planning Results")
             st.warning(
@@ -583,12 +604,6 @@ def setup_clustering():
         st.divider()
         st.header("Clustering the retrosynthetic routes")
 
-        # num_clusters_input = st.number_input( # This input was removed in the final user code, so omitting.
-        #     "Desired Number of Clusters (approximate):",
-        #     min_value=2, max_value=50, value=st.session_state.get("num_clusters_setting", 10),
-        #     key="num_clusters_input_key"
-        # )
-
         if st.button("Run Clustering", key="submit_clustering_button"):
             # st.session_state.num_clusters_setting = num_clusters_input
             st.session_state.clustering_done = False
@@ -597,7 +612,8 @@ def setup_clustering():
             st.session_state.reactions_dict = None
             st.session_state.subclusters = None
             st.session_state.route_cgrs_dict = None
-            st.session_state.r_route_cgrs_dict = None
+            st.session_state.sb_cgrs_dict = None
+            st.session_state.route_json = None
 
             with st.spinner("Performing clustering..."):
                 try:
@@ -608,19 +624,20 @@ def setup_clustering():
 
                     st.write("Calculating RoutesCGRs...")
                     route_cgrs_dict = compose_all_route_cgrs(current_tree)
-                    st.write("Processing ReducedRoutesCGRs...")
-                    r_route_cgrs_dict = compose_all_reduced_route_cgrs(route_cgrs_dict)
+                    st.write("Processing SB-CGRs...")
+                    sb_cgrs_dict = compose_all_sb_cgrs(route_cgrs_dict)
 
                     results = cluster_routes(
-                        r_route_cgrs_dict, use_strat=False
+                        sb_cgrs_dict, use_strat=False
                     )  # num_clusters was removed from args
                     results = dict(sorted(results.items(), key=lambda x: float(x[0])))
 
                     st.session_state.clusters = results
                     st.session_state.route_cgrs_dict = route_cgrs_dict
-                    st.session_state.r_route_cgrs_dict = r_route_cgrs_dict
+                    st.session_state.sb_cgrs_dict = sb_cgrs_dict
                     st.write("Extracting reactions...")
                     st.session_state.reactions_dict = extract_reactions(current_tree)
+                    st.session_state.route_json = make_json(st.session_state.reactions_dict)
 
                     if (
                         st.session_state.clusters is not None
@@ -634,7 +651,7 @@ def setup_clustering():
                         st.error("Clustering failed or returned empty results.")
                         st.session_state.clustering_done = False
 
-                    del results  # route_cgrs_dict, r_route_cgrs_dict are stored
+                    del results  # route_cgrs_dict, sb_cgrs_dict are stored
                     gc.collect()
                     st.rerun()
                 except Exception as e:
@@ -667,49 +684,50 @@ def display_clustering_results():
         for cluster_num, group_data in first_items:
             if (
                 not group_data
-                or "node_ids" not in group_data
-                or not group_data["node_ids"]
+                or "route_ids" not in group_data
+                or not group_data["route_ids"]
             ):
-                st.warning(f"Cluster {cluster_num} has no data or node_ids.")
+                st.warning(f"Cluster {cluster_num} has no data or route_ids.")
                 continue
             st.markdown(
                 f"**Cluster {cluster_num}** (Size: {group_data.get('group_size', 'N/A')})"
             )
-            node_id = group_data["node_ids"][0]
+            route_id = group_data["route_ids"][0]
             try:
-                num_steps = len(tree.synthesis_route(node_id))
-                route_score = round(tree.route_score(node_id), 3)
-                svg = get_route_svg(tree, node_id)
-                r_route_cgr = group_data.get("r_route_cgr")  # Safely get r_route_cgr
-                r_route_cgr_svg = None
-                if r_route_cgr:
-                    r_route_cgr.clean2d()
-                    r_route_cgr_svg = cgr_display(r_route_cgr)
-
-                if svg and r_route_cgr_svg:
+                num_steps = len(tree.synthesis_route(route_id))
+                route_score = round(tree.route_score(route_id), 3)
+                # svg = get_route_svg(tree, route_id)
+                svg = get_route_svg_from_json(st.session_state.route_json, route_id)
+                sb_cgr = group_data.get("sb_cgr")  # Safely get sb_cgr
+                sb_cgr_svg = None
+                if sb_cgr:
+                    sb_cgr.clean2d()
+                    sb_cgr_svg = cgr_display(sb_cgr)
+
+                if svg and sb_cgr_svg:
                     col1, col2 = st.columns([0.2, 0.8])
                     with col1:
-                        st.image(r_route_cgr_svg, caption="ReducedRouteCGR")
+                        st.image(sb_cgr_svg, caption="SB-CGR")
                     with col2:
                         st.image(
                             svg,
-                            caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}",
+                            caption=f"Route {route_id}; {num_steps} steps; Route score: {route_score}",
                         )
                 elif svg:  # Only route SVG available
                     st.image(
                         svg,
-                        caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}",
+                        caption=f"Route {route_id}; {num_steps} steps; Route score: {route_score}",
                     )
                     st.warning(
-                        f"ReducedRouteCGR could not be displayed for cluster {cluster_num}."
+                        f"SB-CGR could not be displayed for cluster {cluster_num}."
                     )
                 else:
                     st.warning(
-                        f"Could not generate SVG for route {node_id} or its ReducedRouteCGR."
+                        f"Could not generate SVG for route {route_id} or its SB-CGR."
                     )
             except Exception as e:
                 st.error(
-                    f"Error displaying route {node_id} for cluster {cluster_num}: {e}"
+                    f"Error displaying route {route_id} for cluster {cluster_num}: {e}"
                 )
 
         if remaining_items:
@@ -717,51 +735,52 @@ def display_clustering_results():
                 for cluster_num, group_data in remaining_items:
                     if (
                         not group_data
-                        or "node_ids" not in group_data
-                        or not group_data["node_ids"]
+                        or "route_ids" not in group_data
+                        or not group_data["route_ids"]
                     ):
                         st.warning(
-                            f"Cluster {cluster_num} in expansion has no data or node_ids."
+                            f"Cluster {cluster_num} in expansion has no data or route_ids."
                         )
                         continue
                     st.markdown(
                         f"**Cluster {cluster_num}** (Size: {group_data.get('group_size', 'N/A')})"
                     )
-                    node_id = group_data["node_ids"][0]
+                    route_id = group_data["route_ids"][0]
                     try:
-                        num_steps = len(tree.synthesis_route(node_id))
-                        route_score = round(tree.route_score(node_id), 3)
-                        svg = get_route_svg(tree, node_id)
-                        r_route_cgr = group_data.get("r_route_cgr")
-                        r_route_cgr_svg = None
-                        if r_route_cgr:
-                            r_route_cgr.clean2d()
-                            r_route_cgr_svg = cgr_display(r_route_cgr)
-
-                        if svg and r_route_cgr_svg:
+                        num_steps = len(tree.synthesis_route(route_id))
+                        route_score = round(tree.route_score(route_id), 3)
+                        # svg = get_route_svg(tree, route_id)
+                        svg = get_route_svg_from_json(st.session_state.route_json, route_id)
+                        sb_cgr = group_data.get("sb_cgr")
+                        sb_cgr_svg = None
+                        if sb_cgr:
+                            sb_cgr.clean2d()
+                            sb_cgr_svg = cgr_display(sb_cgr)
+
+                        if svg and sb_cgr_svg:
                             col1, col2 = st.columns([0.2, 0.8])
                             with col1:
-                                st.image(r_route_cgr_svg, caption="ReducedRouteCGR")
+                                st.image(sb_cgr_svg, caption="SB-CGR")
                             with col2:
                                 st.image(
                                     svg,
-                                    caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}",
+                                    caption=f"Route {route_id}; {num_steps} steps; Route score: {route_score}",
                                 )
                         elif svg:
                             st.image(
                                 svg,
-                                caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}",
+                                caption=f"Route {route_id}; {num_steps} steps; Route score: {route_score}",
                             )
                             st.warning(
-                                f"ReducedRouteCGR could not be displayed for cluster {cluster_num}."
+                                f"SB-CGR could not be displayed for cluster {cluster_num}."
                             )
                         else:
                             st.warning(
-                                f"Could not generate SVG for route {node_id} or its ReducedRouteCGR."
+                                f"Could not generate SVG for route {route_id} or its SB-CGR."
                             )
                     except Exception as e:
                         st.error(
-                            f"Error displaying route {node_id} for cluster {cluster_num}: {e}"
+                            f"Error displaying route {route_id} for cluster {cluster_num}: {e}"
                         )
 
 
@@ -770,8 +789,8 @@ def download_clustering_results():
     if st.session_state.get("clustering_done", False):
         tree_for_html = st.session_state.get("tree")
         clusters_for_html = st.session_state.get("clusters")
-        r_route_cgrs_for_html = st.session_state.get(
-            "r_route_cgrs_dict"
+        sb_cgrs_for_html = st.session_state.get(
+            "sb_cgrs_dict"
         )  # This was used instead of reactions_dict in the original for report
 
         if not tree_for_html:
@@ -780,7 +799,7 @@ def download_clustering_results():
         if not clusters_for_html:
             st.warning("Cluster data not found. Cannot generate cluster reports.")
             return
-        # r_route_cgrs_for_html is optional for routes_clustering_report if not essential
+        # sb_cgrs_for_html is optional for routes_clustering_report if not essential
 
         st.subheader("Cluster Reports")  # Changed subheader in original
         st.write("Generate downloadable HTML reports for each cluster:")
@@ -799,7 +818,7 @@ def download_clustering_results():
                     tree_for_html,
                     clusters_for_html,  # Pass the whole dict
                     str(cluster_idx),  # Pass the key of the cluster
-                    r_route_cgrs_for_html,  # Pass the r_route_cgrs dict
+                    sb_cgrs_for_html,  # Pass the sb_cgrs dict
                     aam=False,
                 )
                 st.download_button(
@@ -826,7 +845,7 @@ def download_clustering_results():
                             tree_for_html,
                             clusters_for_html,
                             str(group_index),
-                            r_route_cgrs_for_html,
+                            sb_cgrs_for_html,
                             aam=False,
                         )
                         st.download_button(
@@ -851,7 +870,7 @@ def download_clustering_results():
                         tree_for_html,
                         clusters_for_html,
                         str(idx),
-                        r_route_cgrs_for_html,
+                        sb_cgrs_for_html,
                         aam=False,
                     )
                     filename = f"cluster_{idx}_{st.session_state.target_smiles}.html"
@@ -883,19 +902,17 @@ def setup_subclustering():
             with st.spinner("Performing subclustering analysis..."):
                 try:
                     clusters_for_sub = st.session_state.get("clusters")
-                    r_route_cgrs_dict_for_sub = st.session_state.get(
-                        "r_route_cgrs_dict"
-                    )
+                    sb_cgrs_dict_for_sub = st.session_state.get("sb_cgrs_dict")
                     route_cgrs_dict_for_sub = st.session_state.get("route_cgrs_dict")
 
                     if (
                         clusters_for_sub
-                        and r_route_cgrs_dict_for_sub
+                        and sb_cgrs_dict_for_sub
                         and route_cgrs_dict_for_sub
                     ):  # Ensure all are present
                         all_subgroups = subcluster_all_clusters(
                             clusters_for_sub,
-                            r_route_cgrs_dict_for_sub,
+                            sb_cgrs_dict_for_sub,
                             route_cgrs_dict_for_sub,
                         )
                         st.session_state.subclusters = all_subgroups
@@ -907,8 +924,8 @@ def setup_subclustering():
                         missing = []
                         if not clusters_for_sub:
                             missing.append("clusters")
-                        if not r_route_cgrs_dict_for_sub:
-                            missing.append("ReducedRouteCGRs dictionary")
+                        if not sb_cgrs_dict_for_sub:
+                            missing.append("SB-CGRs dictionary")
                         if not route_cgrs_dict_for_sub:
                             missing.append("RouteCGRs dictionary")
                         st.error(
@@ -970,17 +987,15 @@ def display_subclustering_results():
                         current_subcluster_data = sub[user_input_cluster_num_display][
                             selected_subcluster_idx
                         ]
-                        if "r_route_cgr" in current_subcluster_data:
-                            cluster_r_route_cgr_display = current_subcluster_data[
-                                "r_route_cgr"
-                            ]
-                            cluster_r_route_cgr_display.clean2d()
+                        if "sb_cgr" in current_subcluster_data:
+                            cluster_sb_cgr_display = current_subcluster_data["sb_cgr"]
+                            cluster_sb_cgr_display.clean2d()
                             st.image(
-                                cluster_r_route_cgr_display.depict(),
-                                caption=f"ReducedRouteCGR of parent Cluster {user_input_cluster_num_display}",
+                                cluster_sb_cgr_display.depict(),
+                                caption=f"SB-CGR of parent Cluster {user_input_cluster_num_display}",
                             )
                         else:
-                            st.warning("ReducedRouteCGR for this subcluster not found.")
+                            st.warning("SB-CGR for this subcluster not found.")
             else:
                 st.warning(
                     f"Selected cluster {user_input_cluster_num_display} not found in subclustering results."
@@ -1002,14 +1017,14 @@ def display_subclustering_results():
                 subcluster_to_display = subcluster_content
                 if (
                     not subcluster_to_display
-                    or "nodes_data" not in subcluster_to_display
-                    or not subcluster_to_display["nodes_data"]
+                    or "routes_data" not in subcluster_to_display
+                    or not subcluster_to_display["routes_data"]
                 ):
                     st.info("No routes or data found for this subcluster selection.")
                 else:
                     MAX_ROUTES_PER_SUBCLUSTER = 5
                     all_route_ids_in_subcluster = list(
-                        subcluster_to_display["nodes_data"].keys()
+                        subcluster_to_display["routes_data"].keys()
                     )
                     routes_to_display_direct = all_route_ids_in_subcluster[
                         :MAX_ROUTES_PER_SUBCLUSTER
@@ -1025,6 +1040,7 @@ def display_subclustering_results():
                     if "synthon_reaction" in subcluster_to_display:
                         synthon_reaction = subcluster_to_display["synthon_reaction"]
                         try:
+                            synthon_reaction.clean2d()
                             st.image(
                                 depict_custom_reaction(synthon_reaction),
                                 caption=f"Markush-like pseudo reaction of subcluster",
@@ -1033,48 +1049,50 @@ def display_subclustering_results():
                             st.warning(f"Could not depict synthon reaction: {e_depict}")
                     else:
                         st.info("No synthon reaction data for this subcluster.")
-
-                    for route_id in routes_to_display_direct:
-                        try:
-                            route_score_sub = round(tree.route_score(route_id), 3)
-                            svg_sub = get_route_svg(tree, route_id)
-                            if svg_sub:
-                                st.image(
-                                    svg_sub,
-                                    caption=f"Route {route_id}; Score: {route_score_sub}",
-                                )
-                            else:
-                                st.warning(
-                                    f"Could not generate SVG for route {route_id}."
+                    with st.container(height=500):
+                        for route_id in routes_to_display_direct:
+                            try:
+                                route_score_sub = round(tree.route_score(route_id), 3)
+                                # svg_sub = get_route_svg(tree, route_id)
+                                svg_sub = get_route_svg_from_json(st.session_state.route_json, route_id)
+                                if svg_sub:
+                                    st.image(
+                                        svg_sub,
+                                        caption=f"Route {route_id}; Score: {route_score_sub}",
+                                    )
+                                else:
+                                    st.warning(
+                                        f"Could not generate SVG for route {route_id}."
+                                    )
+                            except Exception as e:
+                                st.error(
+                                    f"Error displaying route {route_id} in subcluster: {e}"
                                 )
-                        except Exception as e:
-                            st.error(
-                                f"Error displaying route {route_id} in subcluster: {e}"
-                            )
 
-                    if remaining_routes_sub:
-                        with st.expander(
-                            f"... and {len(remaining_routes_sub)} more routes in this subcluster"
-                        ):
-                            for route_id in remaining_routes_sub:
-                                try:
-                                    route_score_sub = round(
-                                        tree.route_score(route_id), 3
-                                    )
-                                    svg_sub = get_route_svg(tree, route_id)
-                                    if svg_sub:
-                                        st.image(
-                                            svg_sub,
-                                            caption=f"Route {route_id}; Score: {route_score_sub}",
+                        if remaining_routes_sub:
+                            with st.expander(
+                                f"... and {len(remaining_routes_sub)} more routes in this subcluster"
+                            ):
+                                for route_id in remaining_routes_sub:
+                                    try:
+                                        route_score_sub = round(
+                                            tree.route_score(route_id), 3
                                         )
-                                    else:
-                                        st.warning(
-                                            f"Could not generate SVG for route {route_id}."
+                                        # svg_sub = get_route_svg(tree, route_id)
+                                        svg_sub = get_route_svg_from_json(st.session_state.route_json, route_id)
+                                        if svg_sub:
+                                            st.image(
+                                                svg_sub,
+                                                caption=f"Route {route_id}; Score: {route_score_sub}",
+                                            )
+                                        else:
+                                            st.warning(
+                                                f"Could not generate SVG for route {route_id}."
+                                            )
+                                    except Exception as e:
+                                        st.error(
+                                            f"Error displaying route {route_id} in subcluster (expanded): {e}"
                                         )
-                                except Exception as e:
-                                    st.error(
-                                        f"Error displaying route {route_id} in subcluster (expanded): {e}"
-                                    )
             else:
                 st.info("Select a valid cluster and subcluster index to see details.")
 
@@ -1089,16 +1107,16 @@ def download_subclustering_results():
 
         sub = st.session_state.get("subclusters")
         tree = st.session_state.get("tree")
-        r_route_cgrs_for_report = st.session_state.get(
-            "r_route_cgrs_dict"
+        sb_cgrs_for_report = st.session_state.get(
+            "sb_cgrs_dict"
         )  # Used by routes_subclustering_report
 
         user_input_cluster_num_display = st.session_state.subcluster_num_select_key
         selected_subcluster_idx = st.session_state.subcluster_index_select_key
 
-        if not tree or not sub or not r_route_cgrs_for_report:
+        if not tree or not sub or not sb_cgrs_for_report:
             st.warning(
-                "Missing data for subclustering report generation (tree, subclusters, or ReducedRouteCGRs)."
+                "Missing data for subclustering report generation (tree, subclusters, or SB-CGRs)."
             )
             return
 
@@ -1114,11 +1132,11 @@ def download_subclustering_results():
             processed_subcluster_data = post_process_subgroup(
                 subcluster_data_for_report
             )
-            if "nodes_data" in subcluster_data_for_report and isinstance(
-                subcluster_data_for_report["nodes_data"], dict
+            if "routes_data" in subcluster_data_for_report and isinstance(
+                subcluster_data_for_report["routes_data"], dict
             ):
                 processed_subcluster_data["group_lgs"] = group_by_identical_values(
-                    subcluster_data_for_report["nodes_data"]
+                    subcluster_data_for_report["routes_data"]
                 )
             else:
                 processed_subcluster_data["group_lgs"] = {}
@@ -1129,7 +1147,7 @@ def download_subclustering_results():
                     processed_subcluster_data,  # Pass the specific post-processed subcluster data
                     user_input_cluster_num_display,
                     selected_subcluster_idx,
-                    r_route_cgrs_for_report,  # Pass the whole r_route_cgrs dict
+                    sb_cgrs_for_report,  # Pass the whole sb_cgrs dict
                     if_lg_group=True,  # This parameter was in the original call
                 )
                 st.download_button(
@@ -1162,7 +1180,8 @@ def implement_restart():
             "reactions_dict",
             "num_clusters_setting",
             "route_cgrs_dict",
-            "r_route_cgrs_dict",
+            "sb_cgrs_dict",
+            "route_json",
             "subclustering_done",
             "subclusters",  # "sub" was renamed
             "clusters_downloaded",