Update ui_generators.py

ui_generators.py

@@ -147,7 +147,7 @@ def run_mentions_tab_display(token_state):
         html_parts.append(mentions_df_display[display_columns].head(20).to_html(escape=False, index=False, classes="table table-sm"))
 
     mentions_html_output = "\n".join(html_parts)
-    fig = None # Initialize fig to None
+    fig = None 
     if not mentions_df.empty and "sentiment_label" in mentions_df.columns:
         try:
             fig_plot, ax = plt.subplots(figsize=(6,4))
@@ -157,18 +157,15 @@ def run_mentions_tab_display(token_state):
             ax.set_ylabel("Count")
             plt.xticks(rotation=45, ha='right')
             plt.tight_layout()
-            fig = fig_plot # Assign the figure object
+            fig = fig_plot 
             logging.info("Mentions tab: Sentiment distribution plot generated.")
         except Exception as e:
             logging.error(f"Error generating mentions plot: {e}", exc_info=True)
-            # fig remains None
         finally:
             if fig is not None and fig is not plt: 
                  plt.close(fig) 
             elif fig is None and plt.get_fignums(): 
                  plt.close('all')
-
-
     return mentions_html_output, fig
 
 
@@ -191,7 +188,6 @@ def run_follower_stats_tab_display(token_state):
     plot_seniority_dist = None
     plot_industry_dist = None
 
-    # --- Monthly Gains Table & Plot ---
     monthly_gains_df = follower_stats_df[
         (follower_stats_df[FOLLOWER_STATS_TYPE_COLUMN] == 'follower_gains_monthly') &
         (follower_stats_df[FOLLOWER_STATS_CATEGORY_COLUMN].notna()) &
@@ -240,7 +236,6 @@ def run_follower_stats_tab_display(token_state):
         html_parts.append("<p>No monthly follower gain data available or required columns missing.</p>")
     html_parts.append("<hr/>")
 
-    # --- Seniority Table & Plot ---
     seniority_df = follower_stats_df[
         (follower_stats_df[FOLLOWER_STATS_TYPE_COLUMN] == 'follower_seniority') &
         (follower_stats_df[FOLLOWER_STATS_CATEGORY_COLUMN].notna()) &
@@ -274,7 +269,6 @@ def run_follower_stats_tab_display(token_state):
         html_parts.append("<p>No follower seniority data available or required columns missing.</p>")
     html_parts.append("<hr/>")
 
-    # --- Industry Table & Plot ---
     industry_df = follower_stats_df[
         (follower_stats_df[FOLLOWER_STATS_TYPE_COLUMN] == 'follower_industry') &
         (follower_stats_df[FOLLOWER_STATS_CATEGORY_COLUMN].notna()) &
@@ -304,7 +298,6 @@ def run_follower_stats_tab_display(token_state):
                  plt.close(plot_industry_dist)
             elif plot_industry_dist is None and plt.get_fignums():
                  plt.close('all')
-
     else:
         html_parts.append("<p>No follower industry data available or required columns missing.</p>")
 
@@ -327,8 +320,8 @@ def create_analytics_plot_panel(label, plot_id_str):
     """
     with gr.Column() as panel_col: 
         with gr.Row(equal_height=False, variant="panel"): 
-            plot_component = gr.Plot(label=label, scale=8) # Plot takes more space
-            with gr.Column(scale=2, min_width=100): # Column for all buttons
+            plot_component = gr.Plot(label=label, scale=8) 
+            with gr.Column(scale=2, min_width=100): 
                 bomb_button = gr.Button(BOMB_ICON, variant="secondary", size="sm", elem_id=f"bomb_{plot_id_str}")
                 explore_button = gr.Button(EXPLORE_ICON, variant="secondary", size="sm", elem_id=f"explore_{plot_id_str}")
                 formula_button = gr.Button(FORMULA_ICON, variant="secondary", size="sm", elem_id=f"formula_{plot_id_str}")
@@ -340,50 +333,51 @@ def build_analytics_tab_plot_area(plot_configs):
     Returns a dictionary of plot UI objects.
     """
     logging.info(f"Building plot area for {len(plot_configs)} analytics plots.")
-    # Stores {"plot_id": {"plot_component": gr.Plot, "bomb_button": gr.Button, "explore_button": gr.Button, "formula_button": gr.Button, "label": str, "panel_component": gr.Column}}
     plot_ui_objects = {} 
     
     current_section_title = ""
-    
-    for i in range(0, len(plot_configs), 2): 
-        if plot_configs[i]["section"] != current_section_title:
-            current_section_title = plot_configs[i]["section"]
+    i = 0
+    while i < len(plot_configs):
+        config1 = plot_configs[i]
+
+        # Start a new section if necessary
+        if config1["section"] != current_section_title:
+            current_section_title = config1["section"]
             gr.Markdown(f"### {current_section_title}")
 
-        with gr.Row(equal_height=False): # Main row for two plot panels
-            # First plot in the pair
-            config1 = plot_configs[i]
+        # Create a new row for each pair of plots or a single plot if it's the last one
+        # or if the next plot is in a different section.
+        with gr.Row(equal_height=False):
+            # Process the first plot of the potential pair
             panel_col1, plot_comp1, bomb_btn1, explore_btn1, formula_btn1 = \
                 create_analytics_plot_panel(config1["label"], config1["id"])
             plot_ui_objects[config1["id"]] = {
-                "plot_component": plot_comp1,
-                "bomb_button": bomb_btn1,
-                "explore_button": explore_btn1,
-                "formula_button": formula_btn1,
-                "label": config1["label"],
-                "panel_component": panel_col1 # Store the panel itself for visibility/scale changes
+                "plot_component": plot_comp1, "bomb_button": bomb_btn1,
+                "explore_button": explore_btn1, "formula_button": formula_btn1,
+                "label": config1["label"], "panel_component": panel_col1
             }
             logging.debug(f"Created UI panel for plot_id: {config1['id']}")
+            i += 1 # Move to the next plot configuration
 
-            # Second plot in the pair, if it exists
-            if i + 1 < len(plot_configs):
-                config2 = plot_configs[i+1]
-                # If the second plot is in a new section, this row will only contain the first plot.
-                # The new section header for config2 will be handled in the next iteration of the outer loop.
+            # Check if there's a second plot for this row
+            if i < len(plot_configs):
+                config2 = plot_configs[i]
+                # Add to the same row ONLY if it's in the same section
                 if config2["section"] == current_section_title:
                     panel_col2, plot_comp2, bomb_btn2, explore_btn2, formula_btn2 = \
                         create_analytics_plot_panel(config2["label"], config2["id"])
                     plot_ui_objects[config2["id"]] = {
-                        "plot_component": plot_comp2,
-                        "bomb_button": bomb_btn2,
-                        "explore_button": explore_btn2,
-                        "formula_button": formula_btn2,
-                        "label": config2["label"],
-                        "panel_component": panel_col2
+                        "plot_component": plot_comp2, "bomb_button": bomb_btn2,
+                        "explore_button": explore_btn2, "formula_button": formula_btn2,
+                        "label": config2["label"], "panel_component": panel_col2
                     }
-                    logging.debug(f"Created UI panel for plot_id: {config2['id']}")
-                # If config2 is in a new section, its panel will be created in the next iteration's new row.
-            # else: panel for config1 is already in the row.
+                    logging.debug(f"Created UI panel for plot_id: {config2['id']} in same row")
+                    i += 1 # Move to the next plot configuration
+                # If config2 is in a new section, the row ends with config1.
+                # config2 will be processed as the first plot in a new row in the next iteration.
+            # If no more plots or next is in new section, row is complete.
         
     logging.info(f"Finished building plot area. Total plot objects: {len(plot_ui_objects)}")
+    if len(plot_ui_objects) != len(plot_configs):
+        logging.error(f"MISMATCH: Expected {len(plot_configs)} plot objects, but created {len(plot_ui_objects)}.")
     return plot_ui_objects