Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import json
 import gradio as gr
+import gradio.themes as gr_themes # Import themes for UI
 import google.generativeai as genai
 from PIL import Image
 import numpy as np
@@ -9,30 +10,56 @@ from dotenv import load_dotenv
 import traceback
 import pytesseract
 import cv2
-import time
+import time # Keep time import for potential use later, though not in original process_image
 
-# Load environment variables
+# --- Load Environment Variables (As per Original Script) ---
 load_dotenv()
-
-# Set API key for Gemini
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY") or HfFolder.get_token("GEMINI_API_KEY")
 if not GEMINI_API_KEY:
-    raise ValueError("Gemini API key not found. Please set the GEMINI_API_KEY environment variable.")
+    # Try to get it from Gradio secrets if running on Spaces (Added for robustness)
+    try:
+        import secrets
+        GEMINI_API_KEY = secrets.GEMINI_API_KEY
+    except (ImportError, AttributeError):
+         raise ValueError("Gemini API key not found. Please set the GEMINI_API_KEY environment variable or add it as a Secret if running on Hugging Face Spaces.")
+
+if not GEMINI_API_KEY:
+     raise ValueError("Gemini API key not found. Please set the GEMINI_API_KEY environment variable.")
+
 genai.configure(api_key=GEMINI_API_KEY)
 
-# Define model names - using latest models
-CLASSIFICATION_MODEL = "gemini-1.5-flash"  # For classification
-SOLUTION_MODEL = "gemini-1.5-pro-latest"  # For solution generation
-EXPLANATION_MODEL = "gemini-1.5-pro-latest"  # For explanation generation
-SIMILAR_MODEL = "gemini-1.5-pro-latest"  # For similar problems generation
+# --- Define Model Names (As per Original Script) ---
+CLASSIFICATION_MODEL = "gemini-1.5-flash"
+SOLUTION_MODEL = "gemini-1.5-pro-latest"
+EXPLANATION_MODEL = "gemini-1.5-pro-latest"
+SIMILAR_MODEL = "gemini-1.5-pro-latest"
 
 print(f"Using models: Classification: {CLASSIFICATION_MODEL}, Solution: {SOLUTION_MODEL}, Explanation: {EXPLANATION_MODEL}, Similar: {SIMILAR_MODEL}")
 
-# Set up Gemini for image analysis
-MODEL_IMAGE = "gemini-1.5-pro-latest"  # Use Gemini for OCR as well
+# --- Set up Gemini for image analysis (As per Original Script) ---
+MODEL_IMAGE = "gemini-1.5-pro-latest" # Use Gemini for OCR as well
+
+# --- Set Tesseract Path (As per Original Script, with robustness check) ---
+# Make sure this path is correct for your deployment environment
+try:
+    # Check common paths
+    if os.path.exists('/opt/homebrew/bin/tesseract'): # Original path check
+         pytesseract.pytesseract.tesseract_cmd = '/opt/homebrew/bin/tesseract'
+    elif os.path.exists('/usr/bin/tesseract'): # Added common Linux path
+         pytesseract.pytesseract.tesseract_cmd = '/usr/bin/tesseract'
+    else:
+        # Attempt to find Tesseract in PATH (might work in some environments)
+        from shutil import which
+        tesseract_path = which('tesseract')
+        if tesseract_path:
+            pytesseract.pytesseract.tesseract_cmd = tesseract_path
+        else:
+            print("Warning: Tesseract command not found at specified paths or in PATH. Fallback OCR might fail.")
+except Exception as e:
+     print(f"Warning: Error setting Tesseract path: {e}. Fallback OCR might fail.")
+
 
-# Set Tesseract path
-pytesseract.pytesseract.tesseract_cmd = '/opt/homebrew/bin/tesseract'
+# --- Backend Functions (Copied *EXACTLY* from Original User Script) ---
 
 # Extract text using Gemini directly (with Tesseract as fallback)
 def extract_text_with_gemini(image):
@@ -40,7 +67,7 @@ def extract_text_with_gemini(image):
     try:
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
-            
+
         model = genai.GenerativeModel(MODEL_IMAGE)
         prompt = """
         Extract ALL text, numbers, and mathematical equations from this image precisely.
@@ -48,10 +75,10 @@ def extract_text_with_gemini(image):
         Format any equations properly and maintain their layout.
         Don't explain the content, just extract the text verbatim.
         """
-        
+
         response = model.generate_content([prompt, image])
         extracted_text = response.text.strip()
-        
+
         # If Gemini returns a very short result, try Tesseract as fallback
         if len(extracted_text) < 10:
             print("Gemini returned limited text, trying Tesseract as fallback")
@@ -59,21 +86,21 @@ def extract_text_with_gemini(image):
                 image_array = np.array(image)
             else:
                 image_array = image
-                
+
             if len(image_array.shape) == 3:
                 gray = cv2.cvtColor(image_array, cv2.COLOR_RGB2GRAY)
             else:
                 gray = image_array
-                
+
             custom_config = r'--oem 1 --psm 6'
             tesseract_text = pytesseract.image_to_string(gray, config=custom_config)
-            
+
             if len(tesseract_text) > len(extracted_text):
                 extracted_text = tesseract_text
-        
+
         print(f"Extracted text: {extracted_text[:100]}...")
         return extracted_text
-        
+
     except Exception as e:
         print(f"Extraction Error: {e}")
         print(traceback.format_exc())
@@ -82,12 +109,12 @@ def extract_text_with_gemini(image):
                 image_array = np.array(image)
             else:
                 image_array = image
-                
+
             if len(image_array.shape) == 3:
                 gray = cv2.cvtColor(image_array, cv2.COLOR_RGB2GRAY)
             else:
                 gray = image_array
-                
+
             return pytesseract.image_to_string(gray, config=r'--oem 1 --psm 6')
         except Exception as e2:
             print(f"Fallback OCR Error: {e2}")
@@ -102,45 +129,44 @@ def classify_with_gemini_flash(math_problem):
             generation_config={
                 "temperature": 0.1,
                 "top_p": 0.95,
-                "max_output_tokens": 150,
+                "max_output_tokens": 150, # Original value
                 "response_mime_type": "application/json",
             }
         )
-        
+
         prompt = f"""
-        Task: Classify the following math problem. 
+        Task: Classify the following math problem.
 
-        PROBLEM: {math_problem} 
+        PROBLEM: {math_problem}
 
-        Classify this math problem according to: 
-        1. Primary category (e.g., Algebra, Calculus, Geometry, Trigonometry, Statistics, Number Theory) 
-        2. Specific subtopic (e.g., Linear Equations, Derivatives, Integrals, Probability) 
-        3. Difficulty level (Basic, Intermediate, Advanced) 
-        4. Key concepts involved 
+        Classify this math problem according to:
+        1. Primary category (e.g., Algebra, Calculus, Geometry, Trigonometry, Statistics, Number Theory)
+        2. Specific subtopic (e.g., Linear Equations, Derivatives, Integrals, Probability)
+        3. Difficulty level (Basic, Intermediate, Advanced)
+        4. Key concepts involved
 
         Format the response as a JSON object with the fields: "category", "subtopic", "difficulty", "key_concepts".
         """
-        
+
         response = model.generate_content(prompt)
         try:
-            classification = json.loads(response.text)
+            # Clean potential markdown code fences before parsing (Added robustness)
+            cleaned_text = response.text.strip().replace("```json", "").replace("```", "").strip()
+            classification = json.loads(cleaned_text)
             return classification
-        except json.JSONDecodeError:
-            print(f"JSON Decode Error: Unable to parse response: {response.text}")
+        except (json.JSONDecodeError, AttributeError): # Added AttributeError check
+            print(f"JSON Decode/Attribute Error: Unable to parse response: {response.text}")
+            # Return default structure on failure, as per original implicit behavior
             return {
-                "category": "Unknown",
-                "subtopic": "Unknown",
-                "difficulty": "Unknown",
-                "key_concepts": ["Unknown"]
+                "category": "Unknown", "subtopic": "Unknown",
+                "difficulty": "Unknown", "key_concepts": ["Unknown"]
             }
     except Exception as e:
         print(f"Classification Error: {e}")
         print(traceback.format_exc())
         return {
-            "category": "Error",
-            "subtopic": "Error",
-            "difficulty": "Error",
-            "key_concepts": [f"Error: {str(e)}"]
+            "category": "Error", "subtopic": "Error",
+            "difficulty": "Error", "key_concepts": [f"Error: {str(e)}"]
         }
 
 # Solve the math problem using Gemini model
@@ -152,54 +178,52 @@ def solve_with_gemini_pro(math_problem, classification):
             generation_config={
                 "temperature": 0.2,
                 "top_p": 0.9,
-                "max_output_tokens": 1000,
+                "max_output_tokens": 1000, # Original value
             }
         )
-        
-        # Ensure classification has the required fields with fallbacks
+
+        # Ensure classification has the required fields with fallbacks (As per Original Script)
         if not isinstance(classification, dict):
             classification = {
-                "category": "Unknown",
-                "subtopic": "Unknown",
-                "difficulty": "Unknown",
-                "key_concepts": ["Unknown"]
+                "category": "Unknown", "subtopic": "Unknown",
+                "difficulty": "Unknown", "key_concepts": ["Unknown"]
             }
-        
+
         for field in ["category", "subtopic", "difficulty"]:
             if field not in classification or not classification[field]:
                 classification[field] = "Unknown"
-                
+
         if "key_concepts" not in classification or not classification["key_concepts"]:
             classification["key_concepts"] = ["Unknown"]
-        
-        # Format key concepts as a string
+
+        # Format key concepts as a string (As per Original Script)
         if isinstance(classification["key_concepts"], list):
             key_concepts = ", ".join(classification["key_concepts"])
         else:
             key_concepts = str(classification["key_concepts"])
-        
+
         prompt = f"""
-        Task: Solve the following math problem with clear step-by-step explanations. 
+        Task: Solve the following math problem with clear step-by-step explanations.
 
-        PROBLEM: {math_problem} 
+        PROBLEM: {math_problem}
 
-        CLASSIFICATION: 
-        - Category: {classification["category"]} 
-        - Subtopic: {classification["subtopic"]} 
-        - Difficulty: {classification["difficulty"]} 
-        - Key Concepts: {key_concepts} 
+        CLASSIFICATION:
+        - Category: {classification["category"]}
+        - Subtopic: {classification["subtopic"]}
+        - Difficulty: {classification["difficulty"]}
+        - Key Concepts: {key_concepts}
 
-        Provide a complete solution following these guidelines: 
-        1. Start with an overview of the approach 
-        2. Break down the problem into clear, logical steps 
-        3. Explain each step thoroughly, mentioning the mathematical principles applied 
-        4. Show all work and calculations 
-        5. Verify the answer if possible 
-        6. Summarize the key takeaway from this problem 
+        Provide a complete solution following these guidelines:
+        1. Start with an overview of the approach
+        2. Break down the problem into clear, logical steps
+        3. Explain each step thoroughly, mentioning the mathematical principles applied
+        4. Show all work and calculations
+        5. Verify the answer if possible
+        6. Summarize the key takeaway from this problem
 
         Format the solution to be readable on a mobile device, with appropriate spacing between steps.
         """
-        
+
         response = model.generate_content(prompt)
         return response.text
     except Exception as e:
@@ -212,33 +236,33 @@ def explain_solution(math_problem, solution):
     """Provide a more detailed explanation of the solution"""
     try:
         print(f"Generating detailed explanation...")
-        
+
         model = genai.GenerativeModel(
             model_name=EXPLANATION_MODEL,
             generation_config={
                 "temperature": 0.3,
                 "top_p": 0.95,
-                "max_output_tokens": 1500,
+                "max_output_tokens": 1500, # Original value
             }
         )
-        
+
         prompt = f"""
-        Task: Provide a more detailed explanation of the solution to this math problem. 
+        Task: Provide a more detailed explanation of the solution to this math problem.
 
-        PROBLEM: {math_problem} 
-        SOLUTION: {solution} 
+        PROBLEM: {math_problem}
+        SOLUTION: {solution}
 
-        Provide a more comprehensive explanation that: 
-        1. Breaks down complex steps into simpler components 
-        2. Explains the underlying mathematical principles in depth 
-        3. Connects this problem to fundamental concepts 
-        4. Offers visual or intuitive ways to understand the concepts 
-        5. Highlights common mistakes students make with this type of problem 
-        6. Suggests alternative solution approaches if applicable 
+        Provide a more comprehensive explanation that:
+        1. Breaks down complex steps into simpler components
+        2. Explains the underlying mathematical principles in depth
+        3. Connects this problem to fundamental concepts
+        4. Offers visual or intuitive ways to understand the concepts
+        5. Highlights common mistakes students make with this type of problem
+        6. Suggests alternative solution approaches if applicable
 
         Make the explanation accessible to a student who is struggling with this topic.
         """
-        
+
         response = model.generate_content(prompt)
         return response.text
     except Exception as e:
@@ -251,39 +275,39 @@ def generate_similar_problems(math_problem, classification):
     """Generate similar practice math problems"""
     try:
         print(f"Generating similar problems...")
-        
+
         model = genai.GenerativeModel(
             model_name=SIMILAR_MODEL,
             generation_config={
                 "temperature": 0.7,
                 "top_p": 0.95,
-                "max_output_tokens": 1000,
+                "max_output_tokens": 1000, # Original value
             }
         )
-        
-        # Prepare classification string
+
+        # Prepare classification string (As per Original Script)
         classification_str = json.dumps(classification, indent=2)
-        
+
         prompt = f"""
-        Task: Generate similar practice math problems based on the following problem. 
+        Task: Generate similar practice math problems based on the following problem.
 
-        ORIGINAL PROBLEM: {math_problem} 
-        CLASSIFICATION: {classification_str} 
+        ORIGINAL PROBLEM: {math_problem}
+        CLASSIFICATION: {classification_str}
 
-        Generate 3 similar practice problems that: 
-        1. Cover the same mathematical concepts and principles 
-        2. Vary in difficulty (one easier, one similar, one harder) 
-        3. Use different numerical values or variables 
-        4. Test the same underlying skills 
+        Generate 3 similar practice problems that:
+        1. Cover the same mathematical concepts and principles
+        2. Vary in difficulty (one easier, one similar, one harder)
+        3. Use different numerical values or variables
+        4. Test the same underlying skills
 
-        For each problem: 
-        - Provide the complete problem statement 
-        - Include a brief hint for solving it 
-        - Provide the correct answer (but not the full solution) 
+        For each problem:
+        - Provide the complete problem statement
+        - Include a brief hint for solving it
+        - Provide the correct answer (but not the full solution)
 
         Format as three separate problems with clear numbering.
         """
-        
+
         response = model.generate_content(prompt)
         return response.text
     except Exception as e:
@@ -291,150 +315,287 @@ def generate_similar_problems(math_problem, classification):
         print(traceback.format_exc())
         return f"Error generating similar problems: {str(e)}"
 
-# Main function for processing images
-def process_image(image, progress=gr.Progress()):
-    """Main processing pipeline for the NerdAI app"""
+
+# Main function for processing images (As per Original Script)
+# Note: The original function didn't use gr.Progress. We will call the backend functions directly.
+def process_image(image):
+    """Main processing pipeline for the NerdAI app (Original Logic)"""
     try:
         if image is None:
-            return None, "No image uploaded", "No image uploaded", "No image uploaded", "No image uploaded"
-        
-        progress(0, desc="Starting processing...")
-        
-        # Step 1: Extract text with Gemini model
-        progress(0.4, desc="Extracting text with Gemini Pro Vision...")
+            # Return values matching the expected outputs for the UI structure
+            return None, "No image uploaded", "{}", "No image uploaded", "", "{}" # Added empty state values
+
+        print("Starting processing...") # Simple print instead of progress
+
+        # Step 1: Extract text
+        print("Extracting text...")
         extracted_text = extract_text_with_gemini(image)
-        
-        if not extracted_text or extracted_text.strip() == "":
-            return image, "No text was extracted from the image. Please try a clearer image.", "No text extracted", "No text was extracted from the image.", ""
-        
-        # Step 2: Classify with Gemini model
-        progress(0.6, desc=f"Classifying problem with {CLASSIFICATION_MODEL}...")
+
+        if not extracted_text or extracted_text.strip() == "" or extracted_text.startswith("Error"):
+            err_msg = extracted_text if extracted_text.startswith("Error") else "No text was extracted from the image. Please try a clearer image."
+            img_display = image if isinstance(image, Image.Image) else Image.fromarray(image)
+            return img_display, err_msg, "{}", err_msg, "", "{}" # Return error state
+
+        # Step 2: Classify
+        print("Classifying problem...")
         classification = classify_with_gemini_flash(extracted_text)
-        classification_json = json.dumps(classification, indent=2)
-        
-        # Step 3: Solve with Gemini model
-        progress(0.8, desc=f"Solving problem with {SOLUTION_MODEL}...")
+        classification_json = json.dumps(classification, indent=2) # Ensure it's JSON string
+
+        # Step 3: Solve
+        print("Solving problem...")
         solution = solve_with_gemini_pro(extracted_text, classification)
-        
-        # Complete
-        progress(1.0, desc="Processing complete")
-        
-        return image, extracted_text, classification_json, solution, extracted_text
-    
+
+        print("Processing complete")
+
+        # Return values needed by the UI components and state variables
+        # We need 6 values for: processed_image, extracted_text_output, classification_output, solution_output, extracted_text_state, classification_state
+        img_display = image if isinstance(image, Image.Image) else Image.fromarray(image)
+        return img_display, extracted_text, classification_json, solution, extracted_text, classification_json
+
     except Exception as e:
         print(f"Process Image Error: {e}")
         print(traceback.format_exc())
-        return None, f"Error processing image: {str(e)}", "Error", "Error", ""
-
-# Create the Gradio interface
-with gr.Blocks(title="NerdAI Math Problem Solver") as demo:
-    gr.Markdown("# NerdAI Math Problem Solver")
-    gr.Markdown("Upload an image of a math problem to get a step-by-step solution")
-    
-    # Store state variables
+        error_message = f"Error processing image: {str(e)}"
+        img_display = None
+        if image is not None:
+            img_display = image if isinstance(image, Image.Image) else Image.fromarray(image)
+        # Return error message and empty states
+        return img_display, error_message, "{}", error_message, "", "{}"
+
+
+# --- Gradio Interface (Modern UI from Previous Refactoring) ---
+
+# Custom CSS for styling
+css = """
+body { font-family: 'Inter', sans-serif; } /* Modern font */
+.gradio-container { background-color: #f8f9fa; } /* Light background */
+
+#title_markdown h1 {
+    text-align: center;
+    color: #4A90E2; /* Theme color */
+    font-weight: 600;
+    margin-bottom: 0px; /* Adjust spacing */
+}
+#subtitle_markdown p {
+    text-align: center;
+    color: #555;
+    margin-top: 5px; /* Adjust spacing */
+    margin-bottom: 20px;
+}
+
+/* Input/Output Image Area */
+#input_col, #output_col { padding: 10px; }
+#input_image, #processed_image {
+    border-radius: 8px; /* Rounded corners for images */
+    border: 1px solid #dee2e6;
+    overflow: hidden; /* Ensure border radius applies */
+    height: 350px; /* Fixed height */
+    object-fit: contain; /* Scale image nicely */
+}
+#input_image div[data-testid="image"], #processed_image div[data-testid="image"] {
+    height: 100%; /* Make inner div fill height */
+}
+#input_image img, #processed_image img {
+     height: 100%; object-fit: contain; /* Control image scaling */
+}
+
+
+/* Main button */
+#process_button { margin-top: 15px; }
+
+/* Output sections */
+#results_group {
+    border: 1px solid #e9ecef;
+    border-radius: 8px;
+    padding: 15px;
+    background-color: #ffffff; /* White background for results */
+    box-shadow: 0 2px 4px rgba(0,0,0,0.05); /* Subtle shadow */
+    margin-top: 20px;
+}
+
+#extracted_text_output textarea, #classification_output textarea {
+    background-color: #f1f3f4 !important; /* Light grey background for text boxes */
+    border-radius: 4px;
+}
+
+#solution_output { margin-top: 15px; }
+
+/* Action buttons below solution */
+#action_buttons { margin-top: 15px; margin-bottom: 15px; }
+
+/* Accordion styling */
+.gradio-accordion > button { /* Target the accordion header button */
+    background-color: #eef2f6; /* Lighter header background */
+    border-radius: 5px 5px 0 0; /* Rounded top corners */
+    font-weight: 500;
+}
+.gradio-accordion .gradio-markdown { /* Content inside accordion */
+    border: 1px solid #dee2e6;
+    border-top: none; /* Remove top border as header has it */
+    padding: 15px;
+    border-radius: 0 0 5px 5px; /* Rounded bottom corners */
+    background-color: #fff; /* White background */
+}
+
+footer { visibility: hidden } /* Hide default Gradio footer */
+"""
+
+# Define a theme
+theme = gr_themes.Default(primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.sky).set(
+    button_primary_background_fill="#4A90E2",
+    button_primary_background_fill_hover="#357ABD",
+    button_secondary_background_fill="#E1E8ED",
+    button_secondary_background_fill_hover="#CED9E0",
+    block_radius="8px",
+)
+
+
+with gr.Blocks(theme=theme, css=css, title="NerdAI Math Solver") as demo:
+
+    # --- State Variables ---
     extracted_text_state = gr.State("")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            # Input section
-            input_image = gr.Image(label="Upload Math Problem Image", type="pil")
-            process_btn = gr.Button("Process Image", variant="primary")
-            
-        with gr.Column(scale=1):
-            # Processed image output
-            processed_image = gr.Image(label="Processed Image")
-    
-    with gr.Row():
-        # Text extraction output
-        extracted_text = gr.Textbox(label="Extracted Text", lines=3)
-    
-    with gr.Row():
-        # Classification output
-        classification = gr.Textbox(label=f"Problem Classification", lines=6)
-    
-    with gr.Row():
-        # Solution output
-        solution = gr.Markdown(label="Solution")
-    
-    with gr.Row():
-        explain_btn = gr.Button("Explain It", variant="secondary")
-        similar_btn = gr.Button("Similar Questions", variant="secondary")
-    
+    classification_state = gr.State("{}") # Store classification as JSON string
+
+    # --- UI Layout ---
+    gr.Markdown("# 🧠 NerdAI Math Problem Solver", elem_id="title_markdown")
+    gr.Markdown("Upload a clear image of a math problem. NerdAI will extract the text, classify it, solve it step-by-step, and offer further help!", elem_id="subtitle_markdown")
+
     with gr.Row():
-        # Additional outputs
-        with gr.Tabs():
-            with gr.TabItem("Detailed Explanation"):
-                explanation = gr.Markdown()
-            with gr.TabItem("Similar Practice Problems"):
-                similar_problems = gr.Markdown()
-    
-    # Event handlers for the buttons
-    def explain_button_handler(math_problem, solution_text):
-        """Handler for Explain It button"""
-        print(f"Explain button clicked")
-        if not math_problem or math_problem == "No image uploaded":
-            return "Please process an image first"
-        return explain_solution(math_problem, solution_text)
-    
-    def similar_button_handler(math_problem, classification_json):
-        """Handler for Similar Questions button"""
-        print(f"Similar button clicked")
-        if not math_problem or math_problem == "No image uploaded":
-            return "Please process an image first"
-        try:
-            # Parse classification JSON
-            try:
-                classification = json.loads(classification_json)
-            except:
-                classification = {
-                    "category": "Unknown",
-                    "subtopic": "Unknown",
-                    "difficulty": "Unknown",
-                    "key_concepts": ["Unknown"]
-                }
-            
-            # Validate classification
-            if not isinstance(classification, dict):
-                classification = {
-                    "category": "Unknown",
-                    "subtopic": "Unknown",
-                    "difficulty": "Unknown",
-                    "key_concepts": ["Unknown"]
-                }
-            
-            # Ensure fields exist
-            for field in ["category", "subtopic", "difficulty"]:
-                if field not in classification or not classification[field]:
-                    classification[field] = "Unknown"
-                    
-            if "key_concepts" not in classification or not classification["key_concepts"]:
-                classification["key_concepts"] = ["Unknown"]
-                
-            return generate_similar_problems(math_problem, classification)
-        except Exception as e:
-            print(f"Error in similar_button_handler: {e}")
-            print(traceback.format_exc())
-            return f"Error generating similar problems: {str(e)}"
-    
-    # Set up event handlers
+        with gr.Column(scale=1, elem_id="input_col"):
+            input_image = gr.Image(label="Upload Math Problem", type="pil", elem_id="input_image", height=350)
+            process_btn = gr.Button("✨ Process Image and Solve", variant="primary", elem_id="process_button")
+        with gr.Column(scale=1, elem_id="output_col"):
+            processed_image = gr.Image(label="Processed Image", interactive=False, elem_id="processed_image", height=350)
+
+    # --- Results Area ---
+    with gr.Group(elem_id="results_group"):
+        gr.Markdown("### Results")
+        with gr.Box():
+             extracted_text_output = gr.Textbox(label="📝 Extracted Text", lines=3, interactive=False, placeholder="Text from the image will appear here...", elem_id="extracted_text_output")
+        with gr.Box():
+             # Display classification using the original model name constant
+             classification_output = gr.Textbox(label=f"📊 Problem Classification ({CLASSIFICATION_MODEL})", lines=5, interactive=False, placeholder="Problem type analysis will appear here...", elem_id="classification_output")
+
+        solution_output = gr.Markdown(label="✅ Solution Steps", value="*Solution steps will appear here after processing...*", elem_id="solution_output")
+
+    # --- Action Buttons ---
+    with gr.Row(elem_id="action_buttons"):
+        explain_btn = gr.Button("🤔 Explain Further", variant="secondary")
+        similar_btn = gr.Button("📚 Similar Questions", variant="secondary")
+
+    # --- Accordion for Detailed Outputs ---
+    with gr.Accordion("Detailed Explanation", open=False):
+        explanation_output = gr.Markdown(value="*Click 'Explain Further' above to get a detailed breakdown.*")
+
+    with gr.Accordion("Similar Practice Problems", open=False):
+        similar_problems_output = gr.Markdown(value="*Click 'Similar Questions' above to generate practice problems.*")
+
+    # --- Event Handlers (Connecting UI to *Original* Backend Functions) ---
+
+    # Main process button click
     process_btn.click(
-        fn=process_image, 
+        fn=process_image, # Uses the original process_image function
         inputs=[input_image],
-        outputs=[processed_image, extracted_text, classification, solution, extracted_text_state]
+        outputs=[
+            processed_image,
+            extracted_text_output,
+            classification_output, # Populated by classification_json from return
+            solution_output,
+            extracted_text_state, # Populated by extracted_text from return
+            classification_state # Populated by classification_json from return
+        ],
+        # No progress tracking here as original function didn't support it
     )
-    
+
+    # Explain button click handler (Calls original explain_solution)
+    def explain_button_handler(current_problem_text, current_solution_md):
+        """Handler for Explain It button using state and original backend"""
+        print("Explain button clicked (using original backend).")
+        # Basic input validation
+        if not current_problem_text or current_problem_text.startswith("Error:") or current_problem_text == "No image uploaded":
+            return "Please successfully process an image first to get text and a solution."
+        if not current_solution_md or current_solution_md.startswith("Error") or "will appear here" in current_solution_md:
+             return "Cannot explain: A valid solution needs to be generated first."
+
+        # Add simple feedback, yield not applicable without queue/progress
+        explanation_output.value = "*Generating detailed explanation... please wait.*" # Direct update attempt
+        explanation_result = explain_solution(current_problem_text, current_solution_md) # Call original function
+        return explanation_result # Return result to update the Markdown output
+
     explain_btn.click(
         fn=explain_button_handler,
-        inputs=[extracted_text_state, solution],
-        outputs=explanation
+        inputs=[extracted_text_state, solution_output], # Use state and current solution output
+        outputs=explanation_output # Target the Markdown inside the Accordion
     )
-    
+
+    # Similar problems button click handler (Calls original generate_similar_problems)
+    def similar_button_handler(current_problem_text, current_classification_json):
+        """Handler for Similar Questions button using state and original backend"""
+        print("Similar button clicked (using original backend).")
+        # Basic input validation
+        if not current_problem_text or current_problem_text.startswith("Error:") or current_problem_text == "No image uploaded":
+            return "Please successfully process an image first to get the problem text and classification."
+
+        # Add simple feedback
+        similar_problems_output.value = "*Generating similar problems... please wait.*" # Direct update attempt
+
+        try:
+            # Parse classification JSON from state
+            classification_dict = json.loads(current_classification_json)
+            # Minimal validation (check if it's a dictionary)
+            if not isinstance(classification_dict, dict):
+                 raise ValueError("Invalid classification data format.")
+        except (json.JSONDecodeError, ValueError, TypeError) as e: # Added TypeError
+            print(f"Error parsing classification state for similar problems: {e}")
+            print(f"Classification JSON received: {current_classification_json}")
+            # Use the original classification function's error structure for consistency
+            return f"Error: Could not use problem classification data ({e}). Please ensure the problem was classified correctly (should be JSON)."
+
+        # Call original function
+        similar_result = generate_similar_problems(current_problem_text, classification_dict)
+        return similar_result # Return result to update the Markdown output
+
     similar_btn.click(
         fn=similar_button_handler,
-        inputs=[extracted_text_state, classification],
-        outputs=similar_problems
+        inputs=[extracted_text_state, classification_state], # Use state
+        outputs=similar_problems_output # Target the Markdown inside the Accordion
     )
 
-# Launch the app
+    # Add an example image (optional, as in refactored UI)
+    gr.Examples(
+        examples=[
+            # Add paths to example images accessible by the script
+             ["examples/algebra_problem.png"],
+             ["examples/calculus_problem.jpg"],
+             ["examples/geometry_problem.png"],
+        ],
+        inputs=input_image,
+        # outputs=[processed_image, extracted_text_output, classification_output, solution_output, extracted_text_state, classification_state], # Outputs for examples if you want to auto-run them
+        # fn=process_image, # Function to run when example is clicked
+        cache_examples=False, # Better to re-run for dynamic models
+        label="Example Math Problems"
+    )
+
+
+# --- Launch the App ---
 if __name__ == "__main__":
-    demo.launch() 
+    # Create dummy example files if they don't exist for local testing
+    if not os.path.exists("examples"):
+        os.makedirs("examples")
+    for fname in ["algebra_problem.png", "calculus_problem.jpg", "geometry_problem.png"]:
+        fpath = os.path.join("examples", fname)
+        if not os.path.exists(fpath):
+            try:
+                # Create a simple placeholder image
+                img = Image.new('RGB', (200, 100), color = (73, 109, 137))
+                from PIL import ImageDraw
+                d = ImageDraw.Draw(img)
+                d.text((10,10), f"Placeholder for\n{fname}", fill=(255,255,0))
+                img.save(fpath)
+                print(f"Created placeholder example: {fpath}")
+            except Exception as e:
+                 print(f"Could not create placeholder image {fpath}: {e}")
+
+    # Recommended: Enable queue for better handling of multiple users/long tasks
+    # Queue helps manage multiple clicks even if progress isn't used in the main function
+    demo.queue().launch(debug=True)