Refactoring till the unabstracted subclustering

app.py

@@ -20,8 +20,8 @@ from synplan.utils.loading import load_reaction_rules, load_building_blocks
 from synplan.utils.visualisation import generate_results_html, get_route_svg
 
 
-from cluster.super_cgr import *
-from cluster.rs_cgr import *
+from cluster.generalized_cgr import *
+from cluster.reduced_g_cgr import *
 from cluster.clustering import *
 from cluster.visualize import *
 from cluster.utils import *
@@ -119,15 +119,29 @@ if "tree" not in st.session_state:
     st.session_state.tree = None
 if "res" not in st.session_state:
     st.session_state.res = None
-if "num_clusters" not in st.session_state:
-    st.session_state.num_clusters = 10
-
-if 'clustering_started' not in st.session_state:
-    st.session_state.clustering_started = False
-if 'clusters_downloaded' not in st.session_state:
+if "target_smiles" not in st.session_state:
+    st.session_state.target_smiles = ''
+
+# Clustering state
+if "clustering_done" not in st.session_state:
+    st.session_state.clustering_done = False
+if "clusters" not in st.session_state:
+    st.session_state.clusters = None
+if "reactions_dict" not in st.session_state:
+    st.session_state.reactions_dict = None
+if "num_clusters_setting" not in st.session_state: # Store the setting used
+    st.session_state.num_clusters_setting = 10
+
+# Subclustering state
+if "subclustering_done" not in st.session_state:
+    st.session_state.subclustering_done = False
+if "sub" not in st.session_state:
+    st.session_state.sub = None
+
+# Download state (less critical now with direct download links)
+if 'clusters_downloaded' not in st.session_state: # Example, might not be needed
     st.session_state.clusters_downloaded = False
 
-
 intro_text = '''
 This is a demo of the graphical user interface of 
 [SynPlanner](https://github.com/Laboratoire-de-Chemoinformatique/SynPlanner/).
@@ -154,26 +168,41 @@ molecule = st.text_input("SMILES:", DEFAULT_MOL)
 smile_code = st_ketcher(molecule)
 target_molecule = mol_from_smiles(smile_code)
 
-building_blocks_path = hf_hub_download(
-        repo_id="Laboratoire-De-Chemoinformatique/SynPlanner",
-        filename="building_blocks_em_sa_ln.smi",
-        subfolder="building_blocks",
-        local_dir="."
-    )
-
-ranking_policy_weights_path = hf_hub_download(
-        repo_id="Laboratoire-De-Chemoinformatique/SynPlanner",
-        filename="ranking_policy_network.ckpt",
-        subfolder="uspto/weights",
-        local_dir="."
-    )
-
-reaction_rules_path = hf_hub_download(
-        repo_id="Laboratoire-De-Chemoinformatique/SynPlanner",
-        filename="uspto_reaction_rules.pickle",
-        subfolder="uspto",
-        local_dir="."
-    )
+if 'target_smiles' in st.session_state and smile_code != st.session_state.target_smiles:
+    # If the SMILES changes, invalidate previous results
+    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.sub = None
+
+@st.cache_resource
+def load_planning_resources():
+    building_blocks_path = hf_hub_download(
+            repo_id="Laboratoire-De-Chemoinformatique/SynPlanner",
+            filename="building_blocks_em_sa_ln.smi",
+            subfolder="building_blocks",
+            local_dir="."
+        )
+    ranking_policy_weights_path = hf_hub_download(
+            repo_id="Laboratoire-De-Chemoinformatique/SynPlanner",
+            filename="ranking_policy_network.ckpt",
+            subfolder="uspto/weights",
+            local_dir="."
+        )
+    reaction_rules_path = hf_hub_download(
+            repo_id="Laboratoire-De-Chemoinformatique/SynPlanner",
+            filename="uspto_reaction_rules.pickle",
+            subfolder="uspto",
+            local_dir="."
+        )
+    return building_blocks_path, ranking_policy_weights_path, reaction_rules_path
+
+building_blocks_path, ranking_policy_weights_path, reaction_rules_path = load_planning_resources()
 
 st.header('Launch calculation')
 st.markdown(
@@ -198,7 +227,7 @@ with col_options_1:
     c_ucb = st.number_input("C coefficient of UCB", value=0.1, placeholder="Type a number...")
 
 with col_options_2:
-    max_iterations = st.slider('Total number of MCTS iterations', min_value=50, max_value=300, value=100)
+    max_iterations = st.slider('Total number of MCTS iterations', min_value=50, max_value=1000, value=300)
     max_depth = st.slider('Maximal number of reaction steps', min_value=3, max_value=9, value=6)
     min_mol_size = st.slider('Minimum size of a molecule to be precursor', min_value=0, max_value=7, value=0)
 
@@ -210,85 +239,147 @@ search_strategy = search_strategy_translator[search_strategy_input]
 
 submit_planning = st.button('Start retrosynthetic planning')
 
-# if submit_planning:
-if submit_planning and not st.session_state.planning_done:
-    with st.status("Downloading data"):
-        st.write("Downloading building blocks")
-        building_blocks = load_building_blocks(building_blocks_path, standardize=False)
-        st.write('Downloading reaction rules')
-        reaction_rules = load_reaction_rules(reaction_rules_path)
-        st.write('Loading policy network')
-        policy_config = PolicyNetworkConfig(weights_path=ranking_policy_weights_path)
-        policy_function = PolicyNetworkFunction(policy_config=policy_config)
-
-    tree_config = TreeConfig(
-        search_strategy=search_strategy,
-        evaluation_type="rollout",
-        max_iterations=max_iterations,
-        max_depth=max_depth,
-        min_mol_size=min_mol_size,
-        init_node_value=0.5,
-        ucb_type=ucb_type,
-        c_ucb=c_ucb,
-        silent=True
-    )
-
-    tree = Tree(
-        target=target_molecule,
-        config=tree_config,
-        reaction_rules=reaction_rules,
-        building_blocks=building_blocks,
-        expansion_function=policy_function,
-        evaluation_function=None,
-    )
-
-    mcts_progress_text = "Running retrosynthetic planning"
-    mcts_bar = st.progress(0, text=mcts_progress_text)
-    for step, (solved, node_id) in enumerate(tree):
-        mcts_bar.progress(step / max_iterations, text=mcts_progress_text)
-
-    res = extract_tree_stats(tree, target_molecule)
-
-    # Store planning outputs in session_state so they persist
-    st.session_state['tree'] = tree
-    st.session_state['res'] = res
-    st.session_state.planning_done = True
-
-    # Display results if planning has been completed
-if st.session_state.planning_done and st.session_state.res is not None:
+if submit_planning:
+    # Reset downstream states if replanning
+    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.sub = None
+    st.session_state.target_smiles = smile_code # Store the SMILES used for this run
+
+    try:
+        target_molecule = mol_from_smiles(smile_code)
+        if target_molecule is None:
+             st.error(f"Could not parse the input SMILES: {smile_code}")
+        else:
+            with st.spinner("Running retrosynthetic planning..."):
+                with st.status("Loading resources...", expanded=False) as status:
+                    st.write("Loading building blocks...")
+                    building_blocks = load_building_blocks(building_blocks_path, standardize=False)
+                    st.write('Loading reaction rules...')
+                    reaction_rules = load_reaction_rules(reaction_rules_path)
+                    st.write('Loading policy network...')
+                    policy_config = PolicyNetworkConfig(weights_path=ranking_policy_weights_path)
+                    policy_function = PolicyNetworkFunction(policy_config=policy_config)
+                    status.update(label="Resources loaded!", state="complete")
+
+                tree_config = TreeConfig(
+                    search_strategy=search_strategy,
+                    evaluation_type="rollout",
+                    max_iterations=max_iterations,
+                    max_depth=max_depth,
+                    min_mol_size=min_mol_size,
+                    init_node_value=0.5,
+                    ucb_type=ucb_type,
+                    c_ucb=c_ucb,
+                    silent=True
+                )
+
+                tree = Tree(
+                    target=target_molecule,
+                    config=tree_config,
+                    reaction_rules=reaction_rules,
+                    building_blocks=building_blocks,
+                    expansion_function=policy_function,
+                    evaluation_function=None,
+                )
+
+                mcts_progress_text = "Running MCTS iterations..."
+                mcts_bar = st.progress(0, text=mcts_progress_text)
+                for step, (solved, node_id) in enumerate(tree):
+                    progress_value = min(1.0, (step + 1) / max_iterations)
+                    mcts_bar.progress(progress_value, text=f"{mcts_progress_text} ({step+1}/{max_iterations})")
+
+                res = extract_tree_stats(tree, target_molecule)
+
+                # Store planning outputs in session_state
+                st.session_state['tree'] = tree
+                st.session_state['res'] = res
+                st.session_state.planning_done = True
+                st.rerun() # Rerun to display results cleanly
+
+    except Exception as e:
+        st.error(f"An error occurred during planning: {e}")
+        st.session_state.planning_done = False # Ensure state reflects failure
+
+# Display results if planning has been completed
+if st.session_state.get('planning_done', False):
     res = st.session_state.res
     tree = st.session_state.tree
 
-    st.header('Results')
-    if res["solved"]:
-        # st.balloons()
-        
+    if res is None or tree is None:
+        st.error("Planning results are missing from session state. Please re-run planning.")
+        st.session_state.planning_done = False # Reset state
+    elif res["solved"]:
+        st.header('Planning Results')
+        # st.balloons() # Optional fun
+        winning_nodes = sorted(set(tree.winning_nodes)) if hasattr(tree, 'winning_nodes') and tree.winning_nodes else []
+        st.subheader(f"Number of unique routes found: {len(winning_nodes)}")
+
         st.subheader("Examples of found retrosynthetic routes")
         image_counter = 0
         visualised_node_ids = set()
-        for n, node_id in enumerate(sorted(set(tree.winning_nodes))):
-            if image_counter == 3:
-                break
-            if n % 2 == 0 and node_id not in visualised_node_ids:
-                visualised_node_ids.add(node_id)
-                image_counter += 1
-                num_steps = len(tree.synthesis_route(node_id))
-                route_score = round(tree.route_score(node_id), 3)
-                st.image(get_route_svg(tree, node_id), 
-                        caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}")
-
+        # Ensure winning_nodes is iterable and not empty
+
+        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):
+                if image_counter >= 3: # Use >= for clarity
+                    break
+                # Simple display logic: show first 3 unique routes
+                if node_id not in visualised_node_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)
+                        if svg:
+                            st.image(svg, caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}")
+                            image_counter += 1
+                        else:
+                            st.warning(f"Could not generate SVG for route {node_id}.")
+                    except Exception as e:
+                        st.error(f"Error displaying route {node_id}: {e}")
+    
         stat_col, download_col = st.columns(2, gap="medium")
         with stat_col:
             st.subheader("Statistics")
-            df = pd.DataFrame(res, index=[0])
-            st.write(df[["target_smiles", "num_routes", "num_nodes", "num_iter", "search_time"]])
+            try:
+                # Ensure 'target_smiles' exists in res, if not, use the stored one
+                if 'target_smiles' not in res:
+                    res['target_smiles'] = st.session_state.target_smiles
+                # Select only existing columns safely
+                cols_to_show = [col for col in ["target_smiles", "num_routes", "num_nodes", "num_iter", "search_time"] if col in res]
+                df = pd.DataFrame(res, index=[0])[cols_to_show]
+                st.dataframe(df) # Use dataframe for better display
+            except Exception as e:
+                st.error(f"Error displaying statistics: {e}")
+                st.write(res) # Show raw dict if DataFrame fails
+        
         with download_col:
             st.subheader("Downloads")
-            html_body = generate_results_html(tree, html_path=None, extended=True)
-            dl_html = download_button(html_body, 'results_synplanner.html', 'Download results as a HTML file')
-            dl_csv = download_button(pd.DataFrame(res, index=[0]), 
-                                    'results_synplanner.csv', 'Download statistics as a csv file')
-            st.markdown(dl_html + dl_csv, unsafe_allow_html=True)
+            try:
+                html_body = generate_results_html(tree, html_path=None, extended=True)
+                dl_html = download_button(html_body, 'results_synplanner.html', 'Download results (HTML)')
+                if dl_html: st.markdown(dl_html, unsafe_allow_html=True)
+
+                # Ensure res is suitable for DataFrame before creating/downloading
+                try:
+                    res_df = pd.DataFrame(res, index=[0])
+                    dl_csv = download_button(res_df, 'results_synplanner.csv', 'Download statistics (CSV)')
+                    if dl_csv: st.markdown(dl_csv, unsafe_allow_html=True)
+                except Exception as e:
+                     st.error(f"Could not prepare statistics CSV for download: {e}")
+
+            except Exception as e:
+                st.error(f"Error generating download links: {e}")
+
+        st.divider()
 
         st.header("Clustering the retrosynthetic routes")
 
@@ -297,92 +388,361 @@ if st.session_state.planning_done and st.session_state.res is not None:
 
         cluster_box, _ = st.columns(2, gap="medium")
         with cluster_box:
-            num_clusters = st.slider(
-                'Number of clusters to display',
+            num_clusters_input = st.slider(
+                'Max number of clusters to generate',
                 min_value=2,
-                max_value=10,
-                value=st.session_state['num_clusters'],
+                max_value=min(50, res.get("num_routes", 50)), # Sensible max based on routes found
+                value=st.session_state.num_clusters_setting,
                 key='cluster_slider'
             )
-            st.session_state['num_clusters'] = num_clusters
-
-        if st.button('Start clustering', key='submit_clustering'):
-            st.session_state.clustering_started = True
-            st.subheader("Examples of clusters")
-            super_cgrs_dict = reassign_nums(tree)
-
-            reduced_super_cgrs_dict = process_all_rs_cgrs(super_cgrs_dict)
-
-            mfp = MorganFingerprint()
-
-            results = cluster_molecules(reduced_super_cgrs_dict, mfp, max_clusters=num_clusters)
-
-            clusters = results['clusters_dict']
-
-            for cluster_num, node_id_list in clusters.items():
-                st.markdown(f"Cluster's number: {cluster_num}; Size {len(node_id_list)}")
-                node_id = node_id_list[0]
-                num_steps = len(tree.synthesis_route(node_id))
-                route_score = round(tree.route_score(node_id), 3)
-                st.image(get_route_svg(tree, node_id), caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}")
-
-            cluster_sizes = [len(cluster) for cluster in clusters.values()]
-
-            cluster_stat_col, cluster_download_col = st.columns(2, gap="medium")
-            with cluster_stat_col:
-                st.subheader("Statistics")
-                cluster_df = pd.DataFrame({'Cluster': range(len(cluster_sizes)), 'Routes': cluster_sizes})
-                st.write(cluster_df)
-
-            def on_download_click():
-                st.session_state.clusters_downloaded = True
-                st.write("Download clusters button pressed via on_click. Updated session state:", dict(st.session_state))
-                save_route_images(tree, reactions_dict, cluster_dict=clusters_converted)
-                
-            with cluster_download_col:
-                st.subheader("Downloads: Don't work. Resets evey time")
-                reactions_dict = extract_reactions(tree)
-                clusters_converted = {int(key): value for key, value in clusters.items()} if clusters else clusters
-                
-                # Use on_click to capture the click event reliably.
-                st.button('Download clusters', key='download_clusters_button', on_click=on_download_click)
-                
-                st.write("Clusters downloaded flag (from session_state):", st.session_state.get("clusters_downloaded"))
-
-            col1, _ = st.columns([.2, .8])
-            with col1:
-                fig = pie_chart(cluster_sizes)
-                st.pyplot(fig)
-            st.header("Sub Clustering the retrosynthetic routes - Resets every time when i interact with input widget")
-            sub = sublcuster_all(clusters, reactions_dict)
-            col2, _ = st.columns([.2, .8])
-            with col2:
-                user_input_cluster_num = st.number_input("Enter a number:", min_value=1,
-                                                          max_value=max(clusters.keys()), value=1, step=1)
-
-                st.write(f"You entered the # cluster: {user_input_cluster_num}")
-                sub_step_cluster = sub[user_input_cluster_num]
-                allowed_numbers = sub_step_cluster.keys()
-                selected_number = st.selectbox("Choose a number:", allowed_numbers)
-                st.write(f"You entered number of steps: {selected_number}")
-                subclusters = sub_step_cluster[selected_number]
-
-            st.subheader(f"Found number of subclusters: {len(subclusters)}")
-            for subcluster_num, subcluster_set in enumerate(subclusters):
-                st.write(f"Subcluster #: {subcluster_num + 1}")
-                for route_id in subcluster_set:
-                    st.write(f"Node_ID: {route_id}")
-                    st.image(get_route_svg(tree, route_id), caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}")
-
-else:
-        st.write("Found no reaction path.")
-
+        if st.button('Run Clustering', key='submit_clustering'):
+             # Update the setting in session state when the button is clicked
+            st.session_state.num_clusters_setting = num_clusters_input
+            # Reset downstream states
+            st.session_state.clustering_done = False
+            st.session_state.subclustering_done = False
+            st.session_state.clusters = None
+            st.session_state.reactions_dict = None
+            st.session_state.sub = None
+
+            with st.spinner("Performing clustering..."):
+                try:
+                    # Ensure tree is available from session state
+                    current_tree = st.session_state.tree
+                    if not current_tree:
+                        st.error("Tree object not found. Please re-run planning.")
+                    else:
+                        st.write("Calculating Generalized CGRs...")
+                        g_cgrs_dict = reassign_nums(current_tree) # Assuming this needs the tree
+                        st.write("Processing RG-CGRs...")
+                        reduced_g_cgrs_dict = process_all_rg_cgrs(g_cgrs_dict) # Assuming this uses the previous output
+
+                        mfp = MorganFingerprint()
+                        st.write(f"Clustering into max {st.session_state.num_clusters_setting} clusters...")
+                        results = cluster_molecules(reduced_g_cgrs_dict, mfp, max_clusters=st.session_state.num_clusters_setting)
+
+                        st.session_state.clusters = results.get('clusters_dict')
+                        st.session_state.rg_cgrs_dict = reduced_g_cgrs_dict
+                        # Extract reactions *after* clustering if needed, ensure tree is passed
+                        st.write("Extracting reactions...")
+                        st.session_state.reactions_dict = extract_reactions(current_tree)
+
+                        if st.session_state.clusters and st.session_state.reactions_dict:
+                             st.session_state.clustering_done = True
+                             st.success(f"Clustering complete. Found {len(st.session_state.clusters)} clusters.")
+                        else:
+                             st.error("Clustering failed or returned empty results.")
+                             st.session_state.clustering_done = False
+
+                        # Clean up large intermediate objects if possible
+                        del g_cgrs_dict
+                        del results
+                        gc.collect()
+
+                        st.rerun() # Rerun to display clustering results cleanly
+                except Exception as e:
+                    st.error(f"An error occurred during clustering: {e}")
+                    st.session_state.clustering_done = False
+
+        # --- Display Clustering Results (if done) ---
+        if st.session_state.get('clustering_done', False):
+            clusters = st.session_state.clusters
+            reactions_dict = st.session_state.reactions_dict # Needed for download
+            
+            tree = st.session_state.tree # Needed for display and download
+
+            if not clusters or not reactions_dict or not tree:
+                 st.error("Clustering results are missing from session state. Please re-run clustering.")
+                 st.session_state.clustering_done = False # Reset flag
+            else:
+                st.subheader(f"Best routes from {len(clusters)} Found Clusters")
+                # Display first route from first few clusters
+                displayed_clusters = 0
+                for cluster_num, node_id_list in clusters.items():
+                    if displayed_clusters >= 10: # Limit displayed clusters
+                         st.write(f"... and {len(clusters) - displayed_clusters} more clusters.")
+                         break
+                    if not node_id_list: continue # Skip empty clusters
+
+                    st.markdown(f"**Cluster {cluster_num}** (Size: {len(node_id_list)}) - Example Route:")
+                    node_id = node_id_list[0] # Display the first route as example
+                    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)
+                        if svg:
+                            st.image(svg, caption=f"Route {node_id}; {num_steps} steps; Route score: {route_score}")
+                        else:
+                            st.warning(f"Could not generate SVG for route {node_id}.")
+                        displayed_clusters += 1
+                    except Exception as e:
+                        st.error(f"Error displaying route {node_id} for cluster {cluster_num+1}: {e}")
+
+                cluster_sizes = [len(cluster) for cluster in clusters.values()]
+                cluster_stat_col, cluster_download_col = st.columns(2, gap="medium")
+
+                with cluster_stat_col:
+                    st.subheader("Cluster Statistics")
+                    if cluster_sizes:
+                        cluster_df = pd.DataFrame({'Cluster': range(1, len(cluster_sizes) + 1), 'Number of Routes': cluster_sizes})
+                        st.dataframe(cluster_df)
+
+                        # Display Pie Chart using Matplotlib
+                        # try:
+                        #     fig, ax = plt.subplots(figsize=(5, 4)) # Adjust size if needed
+                        #     ax.pie(cluster_sizes, labels=[f'C{i+1}' for i in range(len(cluster_sizes))], autopct='%1.1f%%', startangle=90)
+                        #     ax.axis('equal')  # Equal aspect ratio ensures that pie is drawn as a circle.
+                        #     st.pyplot(fig)
+                        #     plt.close(fig) # Close the figure to free memory
+                        # except Exception as e:
+                        #     st.error(f"Could not generate pie chart: {e}")
+                    else:
+                        st.write("No cluster data to display statistics for.")
+
+                with cluster_download_col:
+                    st.subheader("Cluster Reports") # Changed subheader
+
+                    # Retrieve necessary data from session state
+                    tree_for_html = st.session_state.get('tree')
+                    clusters_for_html = st.session_state.get('clusters')
+                    rg_cgrs_for_html = st.session_state.get('rg_cgrs_dict')
+
+                    if not tree_for_html:
+                        st.warning("MCTS Tree data not found. Cannot generate reports.")
+                    elif not clusters_for_html:
+                        st.warning("Cluster data not found. Cannot generate reports.")
+                    else:
+                        st.write("Generate downloadable HTML reports for each cluster:")
+
+                        # Limit the number of download links shown directly if there are many clusters
+                        MAX_DOWNLOAD_LINKS_DISPLAYED = 15 # Adjust as needed
+                        num_clusters_total = len(clusters_for_html)
+                        clusters_items = list(clusters_for_html.items()) # Get items to slice
+
+                        for i, (cluster_num_idx, node_ids) in enumerate(clusters_items):
+                            if i >= MAX_DOWNLOAD_LINKS_DISPLAYED:
+                                st.caption(f"... plus {num_clusters_total - MAX_DOWNLOAD_LINKS_DISPLAYED} more cluster reports available.")
+                                # Consider adding a button to download all as a zip if needed
+                                break
+
+                            cluster_num_display = int(cluster_num_idx) # Use 1-based index
+
+                            if not node_ids: # Skip empty clusters
+                                st.caption(f"Cluster {cluster_num_display} is empty, no report generated.")
+                                continue
+
+                            try:
+                                try:
+                                    cluster_html_content = generate_cluster_html(
+                                        tree=tree_for_html,
+                                        cluster_node_ids=node_ids,
+                                        cluster_num=cluster_num_display,
+                                        rg_cgrs_dict=rg_cgrs_for_html,
+                                        aam=False
+                                    )
+                                except Exception as e:
+                                    st.error(f"Error generating report/link for Cluster {cluster_num_display}: {e}")
+                                
+
+                                # --- Create the download button using the existing function ---
+                                download_filename = f"cluster_{cluster_num_display}_report.html"
+                                button_text = f"Cluster {cluster_num_display} Report (HTML)"
+                                dl_button_html = download_button(
+                                    object_to_download=cluster_html_content, 
+                                    download_filename=download_filename,
+                                    button_text=button_text
+                                    # pickle_it=False # Ensure it's not pickled
+                                )
+
+                                if dl_button_html:
+                                    st.markdown(dl_button_html, unsafe_allow_html=True)
+                                else:
+                                    # Error message if button creation failed (e.g., encoding error)
+                                    st.error(f"Failed to create download link for Cluster {cluster_num_display}.")
+
+                            except Exception as e:
+                                # Catch errors during HTML generation or button creation for a specific cluster
+                                st.error(f"Error generating report/link for Cluster {cluster_num_display}: {e}")
+                                # Optionally add more detailed logging here:
+                                # import traceback
+                                # st.error(traceback.format_exc())
+
+                        if num_clusters_total > MAX_DOWNLOAD_LINKS_DISPLAYED:
+                            # Optional: Add a button here to generate and download a ZIP file
+                            # containing all cluster reports. This requires more implementation
+                            # (using zipfile library in memory).
+                            # e.g., if st.button("Download All Reports as ZIP"): ...
+                            pass
+
+
+                st.divider()
+
+                # --- Subclustering Section ---
+                st.header("Sub-Clustering within a selected Cluster")
+
+                # Button to trigger the subclustering calculation
+                if st.button("Run Subclustering Analysis", key="submit_subclustering"):
+                    st.session_state.subclustering_done = False # Reset flag
+                    st.session_state.sub = None # Clear old results
+                    with st.spinner("Performing subclustering analysis..."):
+                         try:
+                              # Retrieve necessary data from session state
+                              clusters_for_sub = st.session_state.get('clusters')
+                              reactions_for_sub = st.session_state.get('reactions_dict')
+                              if clusters_for_sub and reactions_for_sub:
+                                   sub = sublcuster_all(clusters_for_sub, reactions_for_sub)
+                                   st.session_state.sub = sub
+                                   st.session_state.subclustering_done = True
+                                   st.success("Subclustering analysis complete.")
+                                   # Clean up intermediates if possible
+                                   gc.collect()
+                                   st.rerun() # Rerun to display results/inputs cleanly
+                              else:
+                                   st.error("Missing cluster or reaction data needed for subclustering.")
+                         except Exception as e:
+                              st.error(f"An error occurred during subclustering: {e}")
+                              st.session_state.subclustering_done = False
+
+
+                # Display subclustering inputs and results ONLY if subclustering is done
+                if st.session_state.get('subclustering_done', False):
+                    sub = st.session_state.sub
+                    tree = st.session_state.tree
+                    clusters_for_sub = st.session_state.get('clusters')
+
+                    if not sub or not tree:
+                        st.error("Subclustering results are missing from session state. Please re-run subclustering.")
+                        st.session_state.subclustering_done = False 
+                    else:
+                        sub_input_col, sub_display_col = st.columns([0.2, 0.8]) # Adjust column ratio if needed
+
+                        with sub_input_col:
+                            st.subheader("Select Cluster and Step")
+                            # Cluster selection (use cluster numbers as displayed, usually 1-based)
+                            available_cluster_nums = [int(k) for k in sub.keys()] # Use 1-based indexing for UI
+                            if not available_cluster_nums:
+                                st.warning("No clusters available in subclustering results.")
+                            else:
+                                # Key is essential here to maintain state across reruns
+                                user_input_cluster_num_display = st.selectbox(
+                                    "Select Cluster #:",
+                                    options=sorted(available_cluster_nums),
+                                    key='subcluster_num_select'
+                                )
+
+                                # Convert back to 0-based index for accessing 'sub' dictionary
+                                selected_cluster_idx = user_input_cluster_num_display
+
+                                if selected_cluster_idx in sub:
+                                    sub_step_cluster = sub[selected_cluster_idx]
+                                    allowed_step_numbers = sorted(list(sub_step_cluster.keys()))
+
+                                    if not allowed_step_numbers:
+                                        st.warning(f"No reaction steps found for Cluster {user_input_cluster_num_display}.")
+                                    else:
+                                        # Key is essential here
+                                        selected_step_number = st.selectbox(
+                                            "Select Number of Steps:",
+                                            options=allowed_step_numbers,
+                                            key='subcluster_steps_select'
+                                        )
+                                        # --- Display logic moved to the right column ---
+                                        rg_cgrs = st.session_state.get('rg_cgrs_dict')
+                                        cluster_rg_cgr = rg_cgrs[clusters_for_sub[user_input_cluster_num_display][0]]
+                                        cluster_rg_cgr.clean2d()
+                                        st.image(cluster_rg_cgr.depict(), caption=f"RG-CGR of cluster")
+                                    
+
+                                else:
+                                    st.warning(f"Selected cluster {user_input_cluster_num_display} (index {selected_cluster_idx}) not found in subclustering results.")
+
+                                
+
+                        with sub_display_col:
+                            st.subheader("Subcluster Results")
+                            # Check if inputs are valid before trying to display
+                            if 'user_input_cluster_num_display' in locals() and \
+                            'selected_cluster_idx' in locals() and \
+                            selected_cluster_idx in sub and \
+                            'selected_step_number' in locals() and \
+                            selected_step_number in sub[selected_cluster_idx]:
+
+                                subclusters = sub[selected_cluster_idx][selected_step_number]
+                                st.write(f"Displaying **{len(subclusters)}** subclusters for **Cluster {user_input_cluster_num_display}** with **{selected_step_number} steps**:")
+
+                                if not subclusters:
+                                    st.info("No subclusters found for this selection.")
+                                else:
+                                    # Limit the display if there are too many subclusters/routes
+                                    MAX_DISPLAY_SUBCLUSTERS = 20
+                                    MAX_ROUTES_PER_SUBCLUSTER = 10
+
+                                    for subcluster_num, subcluster_set in enumerate(subclusters):
+                                        if subcluster_num >= MAX_DISPLAY_SUBCLUSTERS:
+                                            st.write(f"... and {len(subclusters) - MAX_DISPLAY_SUBCLUSTERS} more subclusters.")
+                                            break
+
+                                        st.markdown(f"--- \n**Subcluster {subcluster_num + 1}** (Size: {len(subcluster_set)})")
+                                        routes_shown = 0
+                                        for route_id in subcluster_set:
+                                            if routes_shown >= MAX_ROUTES_PER_SUBCLUSTER:
+                                                st.write(f"(Showing first {MAX_ROUTES_PER_SUBCLUSTER} routes)")
+                                                break
+                                            try:
+                                                # Need num_steps and route_score for caption (optional but nice)
+                                                num_steps_sub = len(tree.synthesis_route(route_id))
+                                                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}; {num_steps_sub} steps; Score: {route_score_sub}")
+                                                else:
+                                                    st.warning(f"Could not generate SVG for route {route_id}.")
+                                                routes_shown += 1
+                                            except Exception as e:
+                                                st.error(f"Error displaying route {route_id} in subcluster {subcluster_num+1}: {e}")
+                            else:
+                                st.info("Select a cluster and step number to view subclusters.")
+
+
+    # --- Handling No Solution Case ---
+    elif not st.session_state.get('planning_done', False):
+        # Only show this if planning was attempted but failed (or not run yet)
+        # Avoid showing it if just molecule changed
+        if submit_planning: # Check if the button was just pressed
+            st.warning("Planning did not complete successfully or is still running.")
+
+    else: # Planning done, but res["solved"] is False
+        st.header('Planning Results')
+        st.warning("No reaction path found for the target molecule with the current settings.")
+        st.write("Consider adjusting planning options (e.g., increase iterations, adjust depth, check molecule validity).")
+        # Optionally display basic stats even if not solved
+        stat_col, _ = st.columns(2)
+        with stat_col:
+            st.subheader("Run Statistics (No Solution)")
+            try:
+                 if 'target_smiles' not in res: res['target_smiles'] = st.session_state.target_smiles
+                 cols_to_show = [col for col in ["target_smiles", "num_nodes", "num_iter", "search_time"] if col in res]
+                 df = pd.DataFrame(res, index=[0])[cols_to_show]
+                 st.dataframe(df)
+            except Exception as e:
+                 st.error(f"Error displaying statistics: {e}")
+                 st.write(res)
+
+
+# --- Restart Button ---
 st.divider()
-st.header('Restart from the beginning?')
-if st.button("Restart"):
-    st.session_state.planning_done = False
-    st.session_state.tree = None
-    st.session_state.res = None
-    st.session_state.clustering_started = False
-    st.session_state.clusters_downloaded = False
+st.header('Restart Application State')
+if st.button("Clear All Results & Restart"):
+    # Clear all relevant session state keys
+    keys_to_clear = [
+        "planning_done", "tree", "res", "target_smiles",
+        "clustering_done", "clusters", "reactions_dict", "num_clusters_setting", "rg_cgrs_dict",
+        "subclustering_done", "sub",
+        "clusters_downloaded" # Add any other state keys you use
+    ]
+    for key in keys_to_clear:
+        if key in st.session_state:
+            del st.session_state[key]
+    # Clear ketcher state by assigning a default value (or empty string)
+    st.session_state.ketcher = DEFAULT_MOL # Reset ketcher to default
     st.rerun()