Update app.py

app.py

@@ -1,580 +1,485 @@
-# -*- coding: utf-8 -*-
-"""
-Gradio App for AI Data Analysis, Model Comparison, and Explainability
-Requires: HF_TOKEN and WANDB_API_KEY environment variables.
-"""
-
 import os
 import gradio as gr
-import pandas as pd # Make sure pandas is imported
+import pandas as pd
 import numpy as np
 import matplotlib.pyplot as plt
 import shap
 import lime
 import lime.lime_tabular
-import optuna
+# import optuna # Removing Optuna for this simplified approach
 import wandb
 import json
 import time
 import psutil
 import shutil
 import ast
-from smolagents import HfApiModel, CodeAgent # Assuming smolagents is installed
+from smolagents import HfApiModel, CodeAgent
 from huggingface_hub import login
-from sklearn.model_selection import train_test_split, cross_val_score
+from sklearn.model_selection import train_test_split, cross_val_score # Keep cross_val_score if needed elsewhere, but not primary for comparison here
 from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
-from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
+# from sklearn.metrics import ConfusionMatrixDisplay # Not used currently
+from sklearn.ensemble import RandomForestClassifier # Keep RF
+# from sklearn.ensemble import GradientBoostingClassifier # Remove GB for simplicity now
 from sklearn.linear_model import LogisticRegression
-from sklearn.svm import SVC # Kept import in case you add it later
-from sklearn.preprocessing import LabelEncoder, StandardScaler
-from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import LabelEncoder, StandardScaler # Added StandardScaler
+from sklearn.pipeline import Pipeline # Added Pipeline
 from datetime import datetime
-# from PIL import Image
-import warnings
-import joblib # For saving models
-from typing import List, Tuple, Optional # Keep these
-
-# Suppress common warnings
-warnings.filterwarnings("ignore")
+# from PIL import Image # Likely not needed directly
 
-# --- Authentication and Setup ---
-# (Keep Authentication and Setup block as before)
-print("--- Initializing Setup ---")
+# --- Authentication and Setup (Keep as is) ---
 hf_token = os.getenv("HF_TOKEN")
-wandb_api_key = os.getenv("WANDB_API_KEY")
-wandb_run = None
-if not hf_token: print("Warning: HF_TOKEN environment variable not set.")
-else:
-    try: login(token=hf_token); print("Hugging Face login successful.")
-    except Exception as e: print(f"Hugging Face login failed: {e}")
-if not wandb_api_key:
-    print("Warning: WANDB_API_KEY environment variable not set. WandB logging will be disabled.")
-    if wandb.run is None:
-        try: wandb.init(mode="disabled"); print("WandB initialized in disabled mode.")
-        except Exception as e: print(f"Failed to initialize WandB in disabled mode: {e}")
+wandb_api_key = os.getenv("WANDB_API_KEY") # Get WandB key
+
+# Authenticate with Hugging Face
+if hf_token:
+    try:
+        login(token=hf_token)
+        print("HF Login successful.")
+    except Exception as e:
+        print(f"HF login failed: {e}")
 else:
-    try: wandb.login(key=wandb_api_key); print("WandB login successful.")
+    print("HF_TOKEN not found.")
+
+# Login to WandB if key exists
+if wandb_api_key:
+    try:
+        wandb.login(key=wandb_api_key)
+        print("WandB login successful.")
     except Exception as e:
-        print(f"WandB login failed: {e}. Disabling WandB.")
-        if wandb.run is None:
-             try: wandb.init(mode="disabled"); print("WandB initialized in disabled mode due to login failure.")
-             except Exception as e_init: print(f"Failed to initialize WandB in disabled mode: {e_init}")
-agent = None
+        print(f"WandB login failed: {e}. Logging will be disabled.")
+        wandb.init(mode="disabled") # Disable if login fails
+else:
+    print("WANDB_API_KEY not found. WandB logging disabled.")
+    wandb.init(mode="disabled") # Disable if no key
+
+# SmolAgent initialization (Keep as is)
 try:
-    print("Initializing SmolAgent...")
     model_api = HfApiModel("mistralai/Mixtral-8x7B-Instruct-v0.1", token=hf_token)
-    agent = CodeAgent(tools=[], model=model_api, additional_authorized_imports=["numpy", "pandas", "matplotlib.pyplot", "seaborn", "sklearn", "json", "os"])
-    print("SmolAgent initialized successfully.")
-except Exception as e: print(f"Error initializing SmolAgent: {e}. AI Agent features might fail.")
+    print("SmolAgent API Model initialized.")
+except Exception as e:
+    print(f"SmolAgent initialization failed: {e}")
+    model_api = None # Set to None if failed
+
 df_global = None
-split_data_global = None
-comparison_results_global = None
-best_model_details_global = None
-print("Global variables initialized.")
-print("--- Setup Complete ---")
-
-
-# --- Data Handling ---
-
-def clean_data(df: pd.DataFrame) -> pd.DataFrame: # Added type hints
-    """Cleans the input DataFrame."""
-    print("Starting data cleaning...")
-    df_cleaned = df.copy()
-    df_cleaned = df_cleaned.dropna(how='all', axis=1).dropna(how='all', axis=0)
-    print(f"Shape after dropping fully empty rows/cols: {df_cleaned.shape}")
-    object_cols = df_cleaned.select_dtypes(include='object').columns
-    if not object_cols.empty:
-        print(f"Encoding object columns: {list(object_cols)}")
-        for col in object_cols:
-            df_cleaned[col] = df_cleaned[col].astype(str)
-            df_cleaned[col] = LabelEncoder().fit_transform(df_cleaned[col])
-    numeric_cols = df_cleaned.select_dtypes(include=np.number).columns
+# --- NEW: Global variable for split data ---
+split_data_global = None # Will store (X_train, X_test, y_train, y_test, label_encoder)
+
+# --- clean_data, upload_file, AI Agent functions (Keep as is from your original code) ---
+def clean_data(df):
+    # Your original clean_data implementation
+    df = df.copy() # Work on copy
+    df = df.dropna(how='all', axis=1).dropna(how='all', axis=0)
+    for col in df.select_dtypes(include='object').columns:
+        df[col] = df[col].astype(str)
+        df[col] = LabelEncoder().fit_transform(df[col])
+    # Impute only if numeric columns exist
+    numeric_cols = df.select_dtypes(include=np.number).columns
     if not numeric_cols.empty:
-        cols_with_na = df_cleaned[numeric_cols].isnull().sum()
-        cols_to_impute = cols_with_na[cols_with_na > 0].index
-        if not cols_to_impute.empty:
-            print(f"Imputing NaNs with mean in columns: {list(cols_to_impute)}")
-            for col in cols_to_impute:
-                mean_val = df_cleaned[col].mean()
-                df_cleaned[col] = df_cleaned[col].fillna(mean_val)
-    print("Data cleaning finished.")
-    return df_cleaned
-
-# ADDED TYPE HINT HERE
-def upload_file(file) -> pd.DataFrame:
-    """Handles file upload, cleaning, and global state update."""
-    global df_global, split_data_global, comparison_results_global, best_model_details_global
+        df[numeric_cols] = df[numeric_cols].fillna(df[numeric_cols].mean())
+    return df
+
+def upload_file(file):
+    global df_global, split_data_global # Reset split data on new upload
     df_global = None
     split_data_global = None
-    comparison_results_global = None
-    best_model_details_global = None
-    print("Reset global data states on file change.")
-
     if file is None:
-        # Return an empty DataFrame or one with a status message, matching hint
-        return pd.DataFrame({"Status": ["No file uploaded or file removed."]})
-
-    print(f"Uploading file: {file.name}")
+        return pd.DataFrame({"Error": ["No file uploaded."]})
     try:
-        ext = os.path.splitext(file.name)[-1].lower()
-        if ext == ".csv":
-            df = pd.read_csv(file.name)
-        elif ext in [".xls", ".xlsx"]:
-            df = pd.read_excel(file.name)
-        else:
-             # Return DataFrame matching hint
-             return pd.DataFrame({"Error": [f"Unsupported file type: {ext}"]})
-
-        print(f"Original data shape: {df.shape}")
-        df_cleaned = clean_data(df)
-        print(f"Cleaned data shape: {df_cleaned.shape}")
-        df_global = df_cleaned
-        print("Global DataFrame updated with cleaned data.")
-        # Return DataFrame matching hint
-        return df_global.head()
+        ext = os.path.splitext(file.name)[-1].lower() # Use lower()
+        df = pd.read_csv(file.name) if ext == ".csv" else pd.read_excel(file.name)
+        df = clean_data(df)
+        df_global = df
+        print("File uploaded and cleaned.")
+        return df.head()
     except Exception as e:
-        print(f"Error processing file {file.name}: {e}")
-        df_global = None; split_data_global = None; comparison_results_global = None; best_model_details_global = None
-        # Return DataFrame matching hint
+        print(f"Error in upload_file: {e}")
         return pd.DataFrame({"Error": [f"Failed to process file: {e}"]})
 
-
-# --- AI Agent Analysis ---
-# (Keep format_observations, format_insights, format_analysis_report as before)
-def format_observations(observations):
-    """Formats the observations dictionary into HTML list items."""
-    if not isinstance(observations, dict): return f"<p style='color: orange;'>Observations data is not a dictionary: {type(observations)}</p>"
-    items_html = ""
-    for key, value in observations.items():
-        formatted_key = key.replace('_', ' ').title()
-        if isinstance(value, (dict, list)):
-             formatted_value = json.dumps(value, indent=2); value_html = f"<pre style='margin: 0; padding: 8px; background: #ffffff; border: 1px solid #ccc; border-radius: 4px; font-size: 0.9em; white-space: pre-wrap; word-wrap: break-word;'>{formatted_value}</pre>"
-        else: formatted_value = str(value); value_html = f"<p style='margin: 0; padding: 8px; background: #ffffff; border: 1px solid #ccc; border-radius: 4px; font-size: 0.9em;'>{formatted_value}</p>"
-        items_html += f"""<div style="margin-bottom: 12px; padding: 10px; background: #fdfefe; border-radius: 4px; box-shadow: 0 1px 3px rgba(0,0,0,0.1);"><h4 style="margin: 0 0 8px 0; color: #34495e;">{formatted_key}</h4>{value_html}</div>"""
-    return items_html if items_html else "<p>No observations found.</p>"
-
-def format_insights(insights, visuals):
-    """Formats insights and embeds corresponding visuals."""
-    if not isinstance(insights, dict): return f"<p style='color: orange;'>Insights data is not a dictionary: {type(insights)}</p>"
-    items_html = ""; visual_idx = 0; insight_keys = list(insights.keys())
-    for i, key in enumerate(insight_keys):
-        insight_text = str(insights[key]); formatted_key = key.replace('_', ' ').title()
-        items_html += f"""<div style="margin: 20px 0; padding: 15px; background: #ffffff; border-radius: 8px; box-shadow: 0 2px 5px rgba(0,0,0,0.1);"><h4 style='margin-top: 0; margin-bottom: 10px; color: #16a085;'>Insight {i+1}: {formatted_key}</h4><p style="margin-bottom: 15px;">{insight_text}</p>"""
-        if visual_idx < len(visuals):
-            img_path = visuals[visual_idx]
-            items_html += f'<img src="/file={img_path}" alt="Visualization for {formatted_key}" style="max-width: 95%; height: auto; display: block; margin-top: 10px; border-radius: 6px; border: 1px solid #eee; box-shadow: 0 1px 3px rgba(0,0,0,0.1);">'
-            visual_idx += 1
-        items_html += "</div>"
-    if visual_idx < len(visuals):
-         items_html += "<h4 style='margin-top: 25px; color: #2980b9;'>Additional Visualizations:</h4>"
-         for i in range(visual_idx, len(visuals)):
-             img_path = visuals[i]
-             items_html += f"""<div style="margin: 20px 0; padding: 15px; background: #ffffff; border-radius: 8px; box-shadow: 0 2px 5px rgba(0,0,0,0.1);"><img src="/file={img_path}" alt="Additional Visualization {i+1}" style="max-width: 95%; height: auto; display: block; margin: auto; border-radius: 6px; border: 1px solid #eee; box-shadow: 0 1px 3px rgba(0,0,0,0.1);"></div>"""
-    return items_html if (items_html or visuals) else "<p>No insights or visuals generated/found.</p>"
-
-def format_analysis_report(raw_output, visuals):
-    """Formats the AI agent's output into readable HTML."""
-    print("Formatting AI analysis report...")
-    report_html = ""; analysis_dict = {}; parsing_error = None
+# --- AI Agent functions (Keep your original format_*, analyze_data) ---
+# Placeholder for brevity - use your original functions
+def format_analysis_report(raw_output, visuals): return f"<h2>AI Report</h2><pre>{str(raw_output)}</pre>", visuals
+def format_observations(observations): return f"<pre>{str(observations)}</pre>"
+def format_insights(insights, visuals): return f"<pre>{str(insights)}</pre>"
+def analyze_data(csv_file, additional_notes=""):
+     print("Running AI Agent (stub)...")
+     # Your original analyze_data logic here
+     # Ensure it uses wandb.init(reinit=True) if called multiple times
+     # and finishes the run: wandb.finish()
+     if not model_api: return "AI Agent not initialized.", []
+     # Dummy result
+     analysis_result = {"observations": {"data": "desc"}, "insights": {"insight1": "text"}}
+     visuals = [] # Agent should save plots to './figures/'
+     return format_analysis_report(analysis_result, visuals)
+
+# --- MODIFIED: prepare_data ---
+def prepare_data(df, target_column=None) -> bool:
+    """Splits data and stores it globally. Returns True on success, False on failure."""
+    global split_data_global
+    print("Preparing data split...")
     try:
-        if isinstance(raw_output, str):
-            cleaned_output = raw_output.strip().removeprefix("```python").removeprefix("```json").removesuffix("```").strip()
-            dict_start_index = cleaned_output.find('{')
-            if dict_start_index != -1:
-                 try: analysis_dict = ast.literal_eval(cleaned_output[dict_start_index:])
-                 except (SyntaxError, ValueError, TypeError) as e: parsing_error = f"Error parsing agent output: {e}\nRaw output:\n{raw_output}"; print(parsing_error)
-            else: parsing_error = f"Could not find dictionary start '{{' in agent output.\nRaw output:\n{raw_output}"; print(parsing_error)
-        elif isinstance(raw_output, dict): analysis_dict = raw_output
-        else: parsing_error = f"Output is not a string or dictionary, type: {type(raw_output)}.\nRaw output:\n{str(raw_output)}"; print(parsing_error)
-        report_html = """<div style="font-family: Arial, sans-serif; line-height: 1.6; color: #333; padding: 15px; border: 1px solid #ddd; border-radius: 8px; background-color: #f9f9f9;"><h1 style="color: #2c3e50; border-bottom: 2px solid #3498db; padding-bottom: 10px; margin-top: 0;">📊 AI Data Analysis Report</h1>"""
-        if parsing_error: report_html += f"<div style='background-color: #f8d7da; color: #721c24; border: 1px solid #f5c6cb; padding: 10px; border-radius: 5px; margin-bottom: 15px;'><pre>{parsing_error}</pre></div>"
-        observations = analysis_dict.get('observations', {}); report_html += """<div style="margin-top: 20px; background: #ecf0f1; padding: 15px; border-radius: 5px;"><h2 style="color: #2980b9; margin-top: 0;">🔍 Key Observations</h2>"""
-        report_html += format_observations(observations) if observations else "<p>No 'observations' found or parsed.</p>"; report_html += "</div>"
-        insights = analysis_dict.get('insights', {}); report_html += """<div style="margin-top: 25px;"><h2 style="color: #2980b9;">💡 Insights & Visualizations</h2>"""
-        report_html += format_insights(insights, visuals) if (insights or visuals) else "<p>No 'insights' or visuals found or parsed.</p>"; report_html += "</div>"
-        report_html += "</div>"; print("Report formatting complete.")
-        return report_html, visuals
-    except Exception as e:
-        print(f"Critical error in format_analysis_report: {e}"); error_message = f"<p style='color: red; font-weight: bold;'>Error generating report:</p><pre>{str(e)}</pre>"; raw_display = f"<p style='font-weight: bold;'>Raw Agent Output:</p><pre>{str(raw_output)}</pre>"; return error_message + raw_display, visuals
-
+        if df is None or df.empty:
+            print("Error: DataFrame is empty in prepare_data.")
+            split_data_global = None
+            return False
+
+        # --- Target Column Logic ---
+        if target_column is None:
+            # Prioritize object columns if they exist and are not all unique
+            object_cols = df.select_dtypes(include=['object', 'category']).columns
+            potential_targets = [col for col in object_cols if df[col].nunique() < len(df)]
+            if potential_targets:
+                target_column = potential_targets[0] # Take the first suitable object col
+                print(f"Target column auto-selected (object): '{target_column}'")
+            else:
+                target_column = df.columns[-1] # Fallback to last column
+                print(f"Target column auto-selected (last): '{target_column}'")
+        elif target_column not in df.columns:
+            print(f"Error: Specified target column '{target_column}' not found.")
+            split_data_global = None
+            return False
+
+        X = df.drop(columns=[target_column])
+        y = df[target_column].copy()
+
+        # --- Feature Check (ensure numeric) ---
+        # (Should be handled by clean_data, but double-check)
+        non_numeric_features = X.select_dtypes(exclude=np.number).columns
+        if not non_numeric_features.empty:
+            print(f"Warning: Dropping non-numeric feature columns: {list(non_numeric_features)}")
+            X = X.drop(columns=non_numeric_features)
+        if X.empty:
+            print("Error: No numeric features left after dropping non-numeric ones.")
+            split_data_global = None
+            return False
+
+        # --- Target Encoding ---
+        label_encoder = None
+        if not pd.api.types.is_numeric_dtype(y):
+            print(f"Encoding target column '{target_column}' with LabelEncoder.")
+            label_encoder = LabelEncoder()
+            y = label_encoder.fit_transform(y)
+        else:
+             # Check if float target should be treated as classification (e.g., integers represented as float)
+             if pd.api.types.is_float_dtype(y) and np.all(y == y.astype(int)):
+                  print(f"Target '{target_column}' is float but looks like integer. Converting to int.")
+                  y = y.astype(int)
+
+        # --- Check for sufficient classes ---
+        if y.nunique() < 2:
+             print(f"Error: Target column '{target_column}' has less than 2 unique values after processing.")
+             split_data_global = None
+             return False
+
+        # --- Perform Split ---
+        try:
+            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42, stratify=y)
+            print("Performed stratified split.")
+        except ValueError: # Handle cases where stratification is not possible
+            print("Stratified split failed, using non-stratified split.")
+            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
 
-# ADDED TYPE HINT HERE
-def analyze_data(csv_file, additional_notes="") -> Tuple[str, List[str]]:
-    """Runs the SmolAgent for data analysis and visualization."""
-    global df_global, agent
-    # Default return values matching the type hint
-    default_error_html = "<p style='color:red;'>An error occurred.</p>"
-    default_visuals = []
+        split_data_global = (X_train, X_test, y_train, y_test, label_encoder)
+        print(f"Data split successfully: Train {X_train.shape}, Test {X_test.shape}")
+        return True
 
-    if df_global is None: return "<p style='color:red;'>Please upload a file first.</p>", default_visuals
-    if agent is None: return "<p style='color:red;'>AI Agent is not available (initialization failed).</p>", default_visuals
-    if csv_file is None: return "<p style='color:red;'>File object missing, please re-upload.</p>", default_visuals
+    except Exception as e:
+        print(f"Error during data preparation: {e}")
+        import traceback
+        traceback.print_exc()
+        split_data_global = None
+        return False
+
+# --- NEW: run_comparison_and_explainability ---
+def run_comparison_and_explainability():
+    """Compares models, explains the best one, and logs to WandB."""
+    global df_global, split_data_global
+
+    # --- 1. Check Prerequisites ---
+    if df_global is None:
+        return pd.DataFrame({"Error": ["No data uploaded."]}), None, None, "Error: Upload data first."
+    if split_data_global is None:
+        # Attempt to prepare data if not already done
+        print("Split data not found globally, attempting to prepare now...")
+        if not prepare_data(df_global):
+            return pd.DataFrame({"Error": ["Data preparation failed."]}), None, None, "Error: Failed to prepare data for comparison."
+        # If prepare_data succeeded, split_data_global is now populated
+
+    # Unpack the globally stored split data
+    X_train, X_test, y_train, y_test, label_encoder = split_data_global
+    class_names = getattr(label_encoder, 'classes_', [str(c) for c in np.unique(y_train)]) if label_encoder else [str(c) for c in np.unique(y_train)]
+    class_names = [str(c) for c in class_names] # Ensure strings
 
-    print("--- Starting AI Agent Analysis ---")
-    start_time = time.time(); process = psutil.Process(os.getpid()); initial_memory = process.memory_info().rss / 1024 ** 2
-    figures_dir = './figures'
-    try:
-        if os.path.exists(figures_dir): shutil.rmtree(figures_dir); print(f"Cleaned existing directory: {figures_dir}")
-        os.makedirs(figures_dir); print(f"Created directory: {figures_dir}")
-    except Exception as e: print(f"Error managing figures directory: {e}"); return f"<p style='color:red;'>Error setting up visualization directory: {e}</p>", default_visuals
+    print("--- Starting Model Comparison & Explainability ---")
 
-    wandb_run_agent = None; run_name = f"AgentAnalysis_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
-    # ... (WandB init logic as before) ...
-    if wandb.run is None or wandb.run.mode != "disabled":
-         try:
-             if wandb.run and wandb.run.id: wandb.finish()
-             wandb_run_agent = wandb.init(project="ai-data-analysis-gradio", name=run_name, config={...}, reinit=True) # Simplified config display
-             print(f"WandB run '{run_name}' initialized for Agent Analysis.")
-         except Exception as e: print(f"Error initializing WandB run for Agent Analysis: {e}"); wandb_run_agent = None
-    else: print("WandB disabled, skipping Agent run logging.")
-
-    analysis_result = None; visuals = []
-    try:
-        # ... (prompt definition as before) ...
-        prompt = f"""
-Analyze `df_global`. Tasks: 3 observations, 5 insights, 5 visualizations saved to './figures/'.
-Output Format: Python dictionary {{'observations':{{...}}, 'insights':{{...}}}}. Context: {additional_notes}
-Use `df_global`. Save plots with plt.savefig('./figures/unique_name.png', bbox_inches='tight') and plt.clf(). No plt.show().
-""" # Simplified prompt display
-        print("Running AI agent..."); analysis_result = agent.run(prompt, additional_args={"df_global": df_global.copy()}); print("AI agent finished.")
-        if os.path.exists(figures_dir):
-             visuals = [os.path.join(figures_dir, f) for f in os.listdir(figures_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
-             visuals = [v for v in visuals if os.path.isfile(v)]; print(f"Found {len(visuals)} visualizations in {figures_dir}.")
-        else: print(f"Warning: Figures directory '{figures_dir}' not found after agent run.")
-    except Exception as e:
-        print(f"Error during AI agent execution: {e}"); import traceback; traceback.print_exc()
-        if wandb_run_agent: wandb_run_agent.finish(exit_code=1)
-        # Return values matching type hint on error
-        return f"<p style='color:red;'>Error running AI agent: {e}</p>", default_visuals
-
-    execution_time = time.time() - start_time; final_memory = process.memory_info().rss / 1024 ** 2; memory_usage = final_memory - initial_memory
-    print(f"Agent execution time: {execution_time:.2f}s, Memory usage: {memory_usage:.2f} MB")
-    # ... (WandB logging logic as before) ...
-    if wandb_run_agent:
-        try:
-             wandb.log({"agent_execution_time_sec": execution_time, "agent_memory_usage_mb": memory_usage, "visualizations_generated": len(visuals)}, commit=False)
-             # Log visualizations, etc.
-             for viz_path in visuals:
-                  if os.path.exists(viz_path):
-                     try: wandb.log({f"agent_visualization_{os.path.basename(viz_path)}": wandb.Image(viz_path)}, commit=False)
-                     except Exception as log_e: print(f"Wandb img log error: {log_e}")
-             if analysis_result: wandb.log({"agent_raw_output": str(analysis_result)[:10000]}, commit=True)
-             print("Logged agent results to WandB.")
-        except Exception as e: print(f"Error logging agent results to WandB: {e}")
-        finally: wandb_run_agent.finish(); print(f"WandB run '{run_name}' finished.")
-
-    # Ensure return matches type hint
-    formatted_html, _ = format_analysis_report(analysis_result, visuals) # Get formatted HTML
-    return formatted_html, visuals
-
-
-# --- Model Training and Comparison ---
-
-def prepare_data(df, target_column=None): # -> Tuple[pd.DataFrame, pd.DataFrame, pd.Series, pd.Series] - Internal use, hint optional
-    """Prepares data for modeling (selects target, splits, handles encoding)."""
-    global split_data_global
-    print("--- Preparing Data for Modeling ---")
-    if df is None or df.empty: raise ValueError("Cannot prepare data: DataFrame is empty.")
-    # ... (logic for selecting target, dropping, encoding 'y' as before) ...
-    if target_column is None: target_column = df.columns[-1]; print(f"Target column automatically selected: '{target_column}'")
-    elif target_column not in df.columns: raise ValueError(f"Target column '{target_column}' not found.")
-    else: print(f"Using specified target column: '{target_column}'")
-    X = df.drop(columns=[target_column]); y = df[target_column].copy(); le = None
-    if y.dtype == 'object' or pd.api.types.is_categorical_dtype(y): le = LabelEncoder(); y = le.fit_transform(y); print(f"Encoded target. Classes: {le.classes_}")
-    non_numeric_cols = X.select_dtypes(exclude=np.number).columns
-    if not non_numeric_cols.empty: print(f"Warning: Non-numeric columns found in features: {list(non_numeric_cols)}. Dropping them."); X = X.drop(columns=non_numeric_cols)
-    if X.empty: raise ValueError("No features remaining to train the model.")
-    if y.nunique() < 2: raise ValueError("Target column must have at least two unique classes for classification.")
-    try: X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42, stratify=y)
-    except ValueError as split_e: print(f"Stratified split failed ({split_e}). Trying non-stratified split."); X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
-    split_data_global = (X_train, X_test, y_train, y_test, le); print("Data prepared and split stored globally.")
-    return X_train, X_test, y_train, y_test
-
-# --- Model Training and Comparison ---
-
-# ADDED TYPE HINT HERE
-def train_and_compare_models(tune_rf=True, tune_gb=True, n_trials_optuna=10) -> pd.DataFrame:
-    """Trains, (optionally) tunes, evaluates multiple models, and logs comparison."""
-    global df_global, split_data_global, comparison_results_global, best_model_details_global, wandb_run
-    # Default return DataFrame matching hint
-    default_error_df = pd.DataFrame({"Error": ["Comparison failed."]})
-
-    if df_global is None: print("Error: No data loaded for training/comparison."); return pd.DataFrame({"Error": ["Please upload data first."]})
-
-    print("--- Starting Model Training and Comparison ---")
-    run_name = f"CompareModels_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+    # --- 2. Define Models ---
+    # Using Pipelines where scaling is beneficial
     models_to_compare = {
-        "LogisticRegression": Pipeline([('scaler', StandardScaler()), ('logreg', LogisticRegression(max_iter=1000, random_state=42, class_weight='balanced'))]),
-        "RandomForest": RandomForestClassifier(random_state=42, class_weight='balanced'),
-        "GradientBoosting": GradientBoostingClassifier(random_state=42)
-    }
-    config = {
-        "task": "Model Comparison", "models": list(models_to_compare.keys()), "tune_rf": tune_rf,
-        "tune_gb": tune_gb, "optuna_trials": n_trials_optuna if (tune_rf or tune_gb) else 0,
-        "data_shape": df_global.shape if df_global is not None else "N/A", "test_size": 0.3, "stratify": True
+        "LogisticRegression": Pipeline([
+            ('scaler', StandardScaler()),
+            ('logreg', LogisticRegression(max_iter=1000, random_state=42, class_weight='balanced'))
+        ]),
+        "RandomForest": RandomForestClassifier(random_state=42, class_weight='balanced')
+        # Add more models here if desired, e.g., GradientBoostingClassifier
+        # "GradientBoosting": GradientBoostingClassifier(random_state=42)
     }
 
-    # --- WandB Setup ---
-    # (WandB init logic...)
+    # --- 3. Initialize WandB Run ---
+    run_name = f"CompareExplain_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+    wandb_run = None
     if wandb.run is None or wandb.run.mode != "disabled":
         try:
-            if wandb.run and wandb.run.id: wandb.finish()
-            wandb_run = wandb.init(project="ai-data-analysis-gradio", name=run_name, config=config, reinit=True)
-            print(f"WandB run '{run_name}' initialized for Model Comparison.")
-        except Exception as e: print(f"Error initializing WandB run for Comparison: {e}"); wandb_run = None
-    else: print("WandB disabled, skipping Comparison run logging."); wandb_run = None
-
-    results = []; best_f1 = -1; best_model_obj = None; best_model_name = None; best_model_params = None
-    try:
-        if split_data_global: print("Using previously split data."); X_train, X_test, y_train, y_test, _ = split_data_global
-        else: print("Preparing data for comparison..."); X_train, X_test, y_train, y_test = prepare_data(df_global)
-
-        # --- Optuna Objective Functions ---
-        def objective_rf(trial):
-            # --- CORRECTED PARAMETER DEFINITION ---
-            params = {
-                "n_estimators": trial.suggest_int("n_estimators", 50, 250, step=50),
-                "max_depth": trial.suggest_int("max_depth", 5, 20, log=True),
-                "min_samples_split": trial.suggest_int("min_samples_split", 2, 16),
-                "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 16),
-                "criterion": trial.suggest_categorical("criterion", ["gini", "entropy"]),
-                "class_weight": trial.suggest_categorical("class_weight", ["balanced", "balanced_subsample", None]),
-                "random_state": 42
-            }
-            # --- END CORRECTION ---
-            try:
-                 model = RandomForestClassifier(**params)
-                 score = cross_val_score(model, X_train, y_train, cv=3, scoring="f1_weighted", n_jobs=-1).mean()
-                 if wandb_run: wandb.log({"optuna_rf_trial": trial.number, "optuna_rf_cv_f1w": score, **params}, commit=False)
-                 return score
-            except ValueError as e: print(f"Optuna RF trial error (params {params}): {e}"); return -1
-
-        def objective_gb(trial):
-             # --- CORRECTED PARAMETER DEFINITION ---
-             params = {
-                "n_estimators": trial.suggest_int("n_estimators", 50, 250, step=50),
-                "learning_rate": trial.suggest_float("learning_rate", 0.01, 0.3, log=True),
-                "max_depth": trial.suggest_int("max_depth", 3, 10),
-                "min_samples_split": trial.suggest_int("min_samples_split", 2, 16),
-                "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 16),
-                "subsample": trial.suggest_float("subsample", 0.6, 1.0),
-                "random_state": 42
+            # Finish any existing run
+            if wandb.run: wandb.finish()
+            wandb_run = wandb.init(
+                project="huggingface-data-analysis", # Or your preferred project
+                name=run_name,
+                config={
+                    "task": "Comparison & Explainability",
+                    "models": list(models_to_compare.keys()),
+                    "data_shape_train": X_train.shape,
+                    "data_shape_test": X_test.shape,
+                },
+                reinit=True
+            )
+            print(f"WandB Run '{run_name}' started.")
+        except Exception as e:
+            print(f"WandB init failed: {e}")
+            wandb_run = None # Ensure it's None if failed
+
+    # --- 4. Train and Evaluate Models ---
+    results = []
+    trained_models = {} # Store fitted models
+    print("Comparing models...")
+    for name, model in models_to_compare.items():
+        print(f"  Training {name}...")
+        start_time = time.time()
+        try:
+            model.fit(X_train, y_train)
+            y_pred = model.predict(X_test)
+            duration = time.time() - start_time
+
+            # Calculate metrics
+            metrics = {
+                "Model": name,
+                "Accuracy": accuracy_score(y_test, y_pred),
+                "Precision (Weighted)": precision_score(y_test, y_pred, average="weighted", zero_division=0),
+                "Recall (Weighted)": recall_score(y_test, y_pred, average="weighted", zero_division=0),
+                "F1 Score (Weighted)": f1_score(y_test, y_pred, average="weighted", zero_division=0),
+                "Time (s)": duration
             }
-             # --- END CORRECTION ---
-             try:
-                 model = GradientBoostingClassifier(**params)
-                 score = cross_val_score(model, X_train, y_train, cv=3, scoring="f1_weighted", n_jobs=-1).mean()
-                 if wandb_run: wandb.log({"optuna_gb_trial": trial.number, "optuna_gb_cv_f1w": score, **params}, commit=False)
-                 return score
-             except ValueError as e: print(f"Optuna GB trial error (params {params}): {e}"); return -1
-
-        # --- Model Training Loop ---
-        for name, model_pipeline in models_to_compare.items():
-            print(f"--- Training and Evaluating: {name} ---")
-            start_time = time.time(); current_params = model_pipeline.get_params(); final_model = model_pipeline
-            try:
-                # Tuning logic as before...
-                if name == "RandomForest" and tune_rf:
-                     study_rf = optuna.create_study(direction="maximize"); study_rf.optimize(objective_rf, n_trials=n_trials_optuna, timeout=300)
-                     final_model = RandomForestClassifier(**study_rf.best_params, random_state=42); current_params = final_model.get_params()
-                     print(f"Best RF params (CV F1w: {study_rf.best_value:.4f})")
-                     if wandb_run: wandb.log({f"{name}_best_cv_f1w": study_rf.best_value, f"{name}_best_params": study_rf.best_params}, commit=False)
-                elif name == "GradientBoosting" and tune_gb:
-                     study_gb = optuna.create_study(direction="maximize"); study_gb.optimize(objective_gb, n_trials=n_trials_optuna, timeout=300)
-                     final_model = GradientBoostingClassifier(**study_gb.best_params, random_state=42); current_params = final_model.get_params()
-                     print(f"Best GB params (CV F1w: {study_gb.best_value:.4f})")
-                     if wandb_run: wandb.log({f"{name}_best_cv_f1w": study_gb.best_value, f"{name}_best_params": study_gb.best_params}, commit=False)
-
-                # Train/Eval logic as before...
-                final_model.fit(X_train, y_train)
-                y_pred = final_model.predict(X_test)
-                accuracy=accuracy_score(y_test, y_pred); precision=precision_score(y_test, y_pred, average="weighted", zero_division=0); recall=recall_score(y_test, y_pred, average="weighted", zero_division=0); f1=f1_score(y_test, y_pred, average="weighted", zero_division=0); duration = time.time() - start_time
-                print(f"{name} Test - F1(w): {f1:.4f}, Acc: {accuracy:.4f}, Time: {duration:.2f}s")
-                metrics = { "Model": name, "Test Accuracy": accuracy, "Test Precision (Weighted)": precision, "Test Recall (Weighted)": recall, "Test F1 Score (Weighted)": f1, "Training Time (s)": duration, "Tuned": (name == "RandomForest" and tune_rf) or (name == "GradientBoosting" and tune_gb) }
-                results.append(metrics)
-                # WandB logging logic...
-                if wandb_run:
-                    log_metrics = {f"{name}_{k.lower().replace(' (weighted)','_w').replace(' ','_')}": v for k, v in metrics.items() if k not in ["Model", "Tuned"]}; log_metrics[f"{name}_tuned_flag"] = metrics["Tuned"]
-                    wandb.log(log_metrics, commit=False)
-                # Update best model logic...
-                if f1 > best_f1: best_f1 = f1; best_model_name = name; best_model_obj = final_model; best_model_params = current_params; print(f"*** New best model: {name} ***")
-            except Exception as train_e: print(f"ERROR training/evaluating {name}: {train_e}"); results.append({"Model": name, "Error": str(train_e)}); import traceback; traceback.print_exc()
-
-        # --- Finalize Comparison ---
-        # (Logic as before: create DataFrame, sort, store globals, save artifact, log table)
-        if not results: print("No models trained."); return pd.DataFrame({"Status": ["Model training failed."]})
-        comparison_df = pd.DataFrame(results)
-        if "Test F1 Score (Weighted)" in comparison_df.columns: comparison_df = comparison_df.sort_values(by="Test F1 Score (Weighted)", ascending=False).reset_index(drop=True)
-        comparison_results_global = comparison_df
-        print("\n--- Model Comparison Summary ---"); print(comparison_df.to_string())
-        if best_model_obj is not None: best_model_details_global = {'name': best_model_name, 'model': best_model_obj, 'params': best_model_params, 'f1_score': best_f1}; print(f"Stored best model: {best_model_name}")
-        # Save artifact logic...
-        # Log table logic...
-
-        return comparison_df
-
-    except Exception as e:
-        print(f"Error during model comparison: {e}"); import traceback; traceback.print_exc()
-        if wandb_run: wandb_run.finish(exit_code=1)
-        return pd.DataFrame({"Error": [f"Comparison failed: {e}"]})
-    finally:
-        # Finish WandB run...
-        if wandb_run and wandb.run: wandb_run.finish(); print(f"WandB run '{run_name}' finished.")
-        wandb_run = None
-
-
-# --- Model Explainability ---
-
-# TYPE HINT ALREADY ADDED HERE
-def explainability(_=None) -> Tuple[List[str], Optional[str], str]:
-    """Generates SHAP and LIME explanations for the best performing model."""
-    global split_data_global, best_model_details_global, wandb_run
-    # Default returns match hint
-    default_shap_paths = []
-    default_lime_path = None
-    default_status = "Error: Explainability could not run."
-
-    if split_data_global is None: return default_shap_paths, default_lime_path, "Error: Data not prepared. Run 'Train & Compare' first."
-    if best_model_details_global is None: return default_shap_paths, default_lime_path, "Error: Best model not identified. Run 'Train & Compare' first."
-
-    # --- Get data and model ---
-    X_train, X_test, y_train, y_test, label_encoder = split_data_global
-    best_model_name = best_model_details_global['name']
-    best_model = best_model_details_global['model']
-
-    print(f"--- Generating explanations for the best model: {best_model_name} ---")
-    # ... (Setup output dir, define paths, WandB init logic as before) ...
-    output_dir = "./explainability_plots";
-    if os.path.exists(output_dir): shutil.rmtree(output_dir)
-    os.makedirs(output_dir)
-    shap_summary_path = os.path.join(output_dir, f"shap_summary_{best_model_name}.png")
-    lime_path = os.path.join(output_dir, f"lime_instance_{best_model_name}.png")
-    all_shap_paths = []; status_message = f"Explaining best model: {best_model_name}" # Initialize gallery list
-    # WandB Init...
-    run_name = f"Explain_{best_model_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
-    wandb_run_explain = None
-    if wandb.run is None or wandb.run.mode != "disabled":
-         try:
-             if wandb.run and wandb.run.id: wandb.finish()
-             wandb_run_explain = wandb.init(project="ai-data-analysis-gradio", name=run_name, config={...}, reinit=True)
-             print(f"WandB run '{run_name}' initialized for Explainability.")
-         except Exception as e: print(f"Error initializing Wandb run for Explainability: {e}"); wandb_run_explain = None
-    else: print("WandB disabled, skipping Explainability run logging."); wandb_run_explain = None
+            results.append(metrics)
+            trained_models[name] = model # Store the fitted model
+            print(f"  {name} - F1: {metrics['F1 Score (Weighted)']:.4f}, Time: {duration:.2f}s")
+
+            # Log individual model metrics to WandB
+            if wandb_run:
+                wandb.log({f"{name}_{k.lower().replace(' (weighted)','_w').replace(' ','_')}": v
+                           for k, v in metrics.items() if k != "Model"}, commit=False)
+
+        except Exception as e:
+            print(f"  ERROR training/evaluating {name}: {e}")
+            results.append({"Model": name, "Error": str(e)})
+            if wandb_run: wandb.log({f"{name}_error": str(e)}, commit=False)
+
+    # --- 5. Process Comparison Results ---
+    if not results:
+        if wandb_run: wandb.finish()
+        return pd.DataFrame({"Error": ["No models trained successfully."]}), None, None, "Error: Model training failed."
+
+    comparison_df = pd.DataFrame(results)
+    # Sort by F1, handle potential errors where F1 might be NaN
+    if "F1 Score (Weighted)" in comparison_df.columns:
+        comparison_df = comparison_df.sort_values(by="F1 Score (Weighted)", ascending=False, na_position='last').reset_index(drop=True)
+        best_model_row = comparison_df.iloc[0]
+        best_model_name = best_model_row['Model']
+        if pd.notna(best_model_row["F1 Score (Weighted)"]) and best_model_name in trained_models:
+             best_model = trained_models[best_model_name]
+             print(f"Best model determined: {best_model_name} (F1: {best_model_row['F1 Score (Weighted)']:.4f})")
+        else:
+             best_model = None # Best row had an error or NaN F1
+             best_model_name = "N/A (Error or No Valid Model)"
+             print("Warning: Could not determine a valid best model from results.")
+    else:
+         print("Warning: F1 Score column missing, cannot determine best model.")
+         best_model = None
+         best_model_name = "N/A (F1 Missing)"
 
 
-    try:
-        # --- SHAP Explanation (logic as before) ---
-        print("Calculating SHAP values...")
-        shap_values = None; explainer = None; X_test_for_shap = X_test
-        # ... (logic to determine explainer and calculate shap_values based on model type/pipeline) ...
-        # Simplified example logic:
-        if isinstance(best_model, Pipeline):
-             final_estimator = best_model.steps[-1][1]
-             if isinstance(final_estimator, (RandomForestClassifier, GradientBoostingClassifier)):
-                  print("Using TreeExplainer for Pipeline")
-                  # ... (transform data, init explainer, get shap_values) ...
-             elif isinstance(final_estimator, LogisticRegression):
-                  print("Using KernelExplainer for Pipeline")
-                  # ... (init explainer, get shap_values for subset) ...
-             else: print("SHAP not implemented for this pipeline step.")
-        elif isinstance(best_model, (RandomForestClassifier, GradientBoostingClassifier)):
-             print("Using TreeExplainer for standalone model")
-             explainer = shap.TreeExplainer(best_model); shap_values = explainer.shap_values(X_test_for_shap)
-        else: print("SHAP not implemented for this model type.")
-
-
-        # --- Generate SHAP Plots (logic as before) ---
-        if shap_values is not None:
-             # ... (Generate summary plot, calculate global importance, generate dependence plots) ...
-             # Important: Ensure generated paths are added to `all_shap_paths`
-             # Example:
-             # if summary plot saved: all_shap_paths.append(shap_summary_path)
-             # if dep plot saved: all_shap_paths.append(dep_path)
-             print("Generating SHAP plots...") # Placeholder print
-             # ... (SHAP plot generation, saving, and appending to all_shap_paths) ...
-             # Example: If summary plot is generated and saved:
-             if os.path.exists(shap_summary_path):
-                 all_shap_paths.append(shap_summary_path)
-             # Example: If dependence plots are generated and saved:
-             # for dep_path in shap_dep_paths:
-             #    if os.path.exists(dep_path):
-             #        all_shap_paths.append(dep_path)
-
-
-        # --- LIME Explanation (logic as before) ---
-        print("Generating LIME explanation...")
+    # Log comparison table to WandB
+    if wandb_run and not comparison_df.empty:
         try:
-            # ... (LIME explainer setup, explain_instance, plot saving logic) ...
-             print("LIME explanation generated.") # Placeholder print
-        except Exception as lime_e:
-            print(f"Error generating LIME explanation: {lime_e}")
-            lime_path = None # Indicate failure
-
-        # --- Finalize and Return ---
-        status_message = f"Explanations finished for {best_model_name}." # Update status
-        valid_lime_path = lime_path if lime_path and os.path.exists(lime_path) else None
-        valid_shap_paths = [p for p in all_shap_paths if p and os.path.exists(p)]
-        print(f"Returning {len(valid_shap_paths)} SHAP paths, LIME path: {valid_lime_path}")
-        # Ensure return matches type hint
-        return valid_shap_paths, valid_lime_path, status_message
+            wandb.log({"model_comparison": wandb.Table(dataframe=comparison_df)}, commit=False)
+            print("Logged comparison table to WandB.")
+        except Exception as e:
+            print(f"Error logging comparison table: {e}")
+
+    # --- 6. Explain Best Model (if found) ---
+    shap_plot_path = None
+    lime_plot_path = None
+    explain_status = f"Compared {len(trained_models)} models. Best: {best_model_name}."
+
+    if best_model:
+        print(f"Generating explanations for {best_model_name}...")
+        explain_dir = "./explain_plots"
+        if os.path.exists(explain_dir): shutil.rmtree(explain_dir)
+        os.makedirs(explain_dir)
+        shap_plot_path = os.path.join(explain_dir, f"shap_{best_model_name}.png")
+        lime_plot_path = os.path.join(explain_dir, f"lime_{best_model_name}.png")
 
-    except Exception as e:
-        print(f"An error occurred during explainability: {e}"); import traceback; traceback.print_exc()
-        status_message = f"Error during explanation: {e}"
-        if wandb_run_explain: wandb_run_explain.finish(exit_code=1)
-        # Ensure return matches type hint
-        return [], None, status_message
-    finally:
-        plt.close('all')
-        # Finish WandB run logic...
-        if wandb_run_explain and wandb.run and wandb.run.id == wandb_run_explain.id: wandb_run_explain.finish(); print(f"WandB run '{run_name}' finished.")
-        wandb_run_explain = None
-
-
-# --- Gradio Interface ---
-# (Keep Gradio UI layout and connections exactly as before)
-print("--- Setting up Gradio Interface ---")
-with gr.Blocks(theme=gr.themes.Soft(), title="AI Data Analysis & Model Comparison") as demo:
-    gr.Markdown( ... ) # Title markdown
-    # Row 1: File Upload ...
+        try:
+            # --- SHAP ---
+            explainer = None
+            shap_values = None
+            X_test_for_shap = X_test # Default
+
+            if isinstance(best_model, Pipeline):
+                final_estimator = best_model.steps[-1][1]
+                if isinstance(final_estimator, (RandomForestClassifier, LogisticRegression)): # Add others if needed
+                     # Use KernelExplainer for pipelines generally (safer)
+                     print("  Using SHAP KernelExplainer for Pipeline...")
+                     predict_proba_pipe = lambda x_np: best_model.predict_proba(pd.DataFrame(x_np, columns=X_test.columns))
+                     # Smaller background for KernelExplainer
+                     X_train_summary = shap.kmeans(X_train.values, min(50, X_train.shape[0]), random_state=42)
+                     explainer = shap.KernelExplainer(predict_proba_pipe, X_train_summary)
+                     X_test_sample = X_test.sample(min(50, X_test.shape[0]), random_state=42) # Sample for faster explanation
+                     shap_values = explainer.shap_values(X_test_sample.values, nsamples='auto')
+                     X_test_for_shap = X_test_sample # Use the same sample for plotting
+                     print("  SHAP values calculated (Kernel).")
+                else:
+                     print(f"  SHAP not configured for pipeline step: {type(final_estimator)}")
+            elif isinstance(best_model, RandomForestClassifier):
+                print("  Using SHAP TreeExplainer...")
+                explainer = shap.TreeExplainer(best_model)
+                shap_values = explainer.shap_values(X_test)
+                print("  SHAP values calculated (Tree).")
+            else:
+                print(f"  SHAP not configured for model type: {type(best_model)}")
+
+            if shap_values is not None:
+                plt.figure()
+                # Use class index 1 for binary, or average for multiclass summary
+                plot_values = shap_values
+                shap_title = f"SHAP Summary ({best_model_name})"
+                if isinstance(shap_values, list) and len(class_names) == 2:
+                    plot_values = shap_values[1] # Positive class
+                    shap_title += f" - Class: {class_names[1]}"
+
+                shap.summary_plot(plot_values, X_test_for_shap, plot_type="dot", show=False, class_names=class_names)
+                plt.title(shap_title)
+                plt.tight_layout()
+                plt.savefig(shap_plot_path, bbox_inches='tight')
+                plt.clf()
+                print(f"  SHAP plot saved: {shap_plot_path}")
+                if wandb_run: wandb.log({"shap_summary_best": wandb.Image(shap_plot_path)}, commit=False)
+            else:
+                shap_plot_path = None # No plot generated
+
+            # --- LIME ---
+            print("  Generating LIME explanation...")
+            # LIME needs predict_proba function
+            predict_fn_lime = None
+            if hasattr(best_model, 'predict_proba'):
+                predict_fn_lime = lambda x_np: best_model.predict_proba(pd.DataFrame(x_np, columns=X_train.columns))
+            else: # Fallback if no predict_proba
+                num_classes_lime = len(class_names)
+                predict_fn_lime = lambda x: np.ones((len(x), num_classes_lime)) / num_classes_lime
+
+            lime_explainer = lime.lime_tabular.LimeTabularExplainer(
+                training_data=X_train.values,
+                feature_names=X_train.columns.tolist(),
+                class_names=class_names,
+                mode='classification' # Assume classification
+            )
+            instance_idx = 0 # Explain the first test instance
+            lime_exp = lime_explainer.explain_instance(
+                data_row=X_test.iloc[instance_idx].values,
+                predict_fn=predict_fn_lime,
+                num_features=10
+            )
+            lime_fig = lime_exp.as_pyplot_figure()
+            lime_fig.suptitle(f"LIME Explanation ({best_model_name} - Instance {instance_idx})")
+            lime_fig.savefig(lime_plot_path, bbox_inches='tight')
+            plt.clf()
+            print(f"  LIME plot saved: {lime_plot_path}")
+            if wandb_run: wandb.log({"lime_explanation_best": wandb.Image(lime_plot_path)}, commit=False)
+
+            explain_status += f" Explanations generated for {best_model_name}."
+
+        except Exception as e:
+            print(f"  ERROR during explanation: {e}")
+            import traceback
+            traceback.print_exc()
+            explain_status += f" Explanation failed for {best_model_name}: {e}"
+            # Keep paths as None if error occurred
+            if not os.path.exists(shap_plot_path): shap_plot_path = None
+            if not os.path.exists(lime_plot_path): lime_plot_path = None
+
+    else:
+        explain_status += " No best model found to explain."
+
+
+    # --- 7. Finish WandB Run and Return ---
+    if wandb_run:
+        wandb.log({}, commit=True) # Ensure final commit
+        wandb.finish()
+        print(f"WandB Run '{run_name}' finished.")
+
+    # Filter out non-existent plot paths before returning
+    valid_shap_path = shap_plot_path if shap_plot_path and os.path.exists(shap_plot_path) else None
+    valid_lime_path = lime_plot_path if lime_plot_path and os.path.exists(lime_plot_path) else None
+
+    return comparison_df, valid_shap_path, valid_lime_path, explain_status
+
+
+# --- Gradio UI ---
+with gr.Blocks() as demo:
+    gr.Markdown("## 📊 AI Data Analysis, Model Comparison & Explainability")
+
+    # --- Row 1: Upload ---
     with gr.Row():
-        with gr.Column(scale=1): file_input = gr.File(...)
-        with gr.Column(scale=2): df_output = gr.DataFrame(...)
-    # Row 2: AI Agent ...
-    with gr.Accordion("🤖 Step 2 (Optional): Run AI Agent for Insights & Visuals", open=False):
+        with gr.Column(scale=1):
+            file_input = gr.File(label="1. Upload CSV or Excel", type="filepath", file_types=[".csv", ".xls", ".xlsx"])
+        with gr.Column(scale=2):
+            df_output = gr.DataFrame(label="Cleaned Data Preview", interactive=False)
+
+    # --- Row 2: AI Agent (Optional) ---
+    with gr.Accordion("🤖 Step 2 (Optional): Run AI Agent Insights", open=False):
          with gr.Row():
-             with gr.Column(scale=1): agent_notes = gr.Textbox(...); agent_btn = gr.Button(...)
-             with gr.Column(scale=2): insights_output = gr.HTML(...)
-         with gr.Row(): visual_output = gr.Gallery(...)
-    # Row 3: Model Training ...
-    with gr.Accordion("⚙️ Step 3: Train & Compare Models", open=True):
+             with gr.Column(scale=1):
+                agent_notes = gr.Textbox(label="Optional requests for Agent", placeholder="e.g., 'Focus on column X'")
+                agent_btn = gr.Button("Run AI Analysis", interactive=(model_api is not None))
+             with gr.Column(scale=2):
+                 insights_output = gr.HTML(label="AI Agent Report")
          with gr.Row():
-             with gr.Column(scale=1): tune_rf_checkbox = gr.Checkbox(...); tune_gb_checkbox = gr.Checkbox(...); optuna_trials_slider = gr.Slider(...); compare_btn = gr.Button(...)
-             with gr.Column(scale=2): comparison_output = gr.DataFrame(...)
-    # Row 4: Explainability ...
-    with gr.Accordion("💡 Step 4: Explain Best Model (SHAP & LIME)", open=False):
-         with gr.Row(): explain_btn = gr.Button(...); explain_status = gr.Textbox(...)
-         with gr.Row(): shap_gallery = gr.Gallery(...); lime_img = gr.Image(...)
+             visual_output = gr.Gallery(label="AI Agent Visualizations", height=350, object_fit="contain", columns=3, preview=True)
+
+    # --- Row 3: Compare & Explain ---
+    with gr.Accordion("⚙️💡 Step 3: Compare Models & Explain Best", open=True):
+        with gr.Row():
+            compare_explain_btn = gr.Button("Run Comparison & Explain Best Model", variant="primary")
+        with gr.Row():
+            comparison_output = gr.DataFrame(label="Model Comparison Results", interactive=False)
+        with gr.Row():
+            explain_status_output = gr.Textbox(label="Status", interactive=False)
+        with gr.Row():
+            # Only one SHAP plot expected now (summary)
+            shap_img_output = gr.Image(label="SHAP Summary (Best Model)", type="filepath", interactive=False)
+            lime_img_output = gr.Image(label="LIME Explanation (Best Model - Instance 0)", type="filepath", interactive=False)
+
 
     # --- Connect Components ---
-    print("Connecting Gradio components...")
     file_input.change(fn=upload_file, inputs=file_input, outputs=df_output)
+
+    # AI Agent connection (Keep as is)
     agent_btn.click(fn=analyze_data, inputs=[file_input, agent_notes], outputs=[insights_output, visual_output])
-    compare_btn.click(fn=train_and_compare_models, inputs=[tune_rf_checkbox, tune_gb_checkbox, optuna_trials_slider], outputs=[comparison_output])
-    explain_btn.click(fn=explainability, inputs=[], outputs=[shap_gallery, lime_img, explain_status])
-    print("Gradio components connected.")
-
-
-# --- Launch the App ---
-if __name__ == "__main__":
-    print("--- Cleaning up temporary directories/files ---")
-    # (Cleanup logic as before)
-    temp_dirs = ['./figures', './explainability_plots', './saved_models', './__pycache__']
-    for d in temp_dirs:
-        if os.path.exists(d):
-            try: shutil.rmtree(d); print(f"Cleaned up directory: {d}")
-            except Exception as e: print(f"Warning: Could not clean up directory {d}: {e}")
-
-    print("--- Launching Gradio App ---")
-    demo.launch(debug=False, share=False)
-    print("--- Gradio App Closed ---")
+
+    # NEW: Connection for combined comparison and explainability
+    compare_explain_btn.click(
+        fn=run_comparison_and_explainability,
+        inputs=[], # Takes data from global df_global
+        outputs=[comparison_output, shap_img_output, lime_img_output, explain_status_output]
+    )
+
+# --- Launch ---
+print("Launching Gradio App...")
+demo.launch(debug=True) # Use debug=True for more detailed errors during development