image.py

# image.py
import gradio as gr
import sympy as sp
from pix2text import Pix2Text
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import re
import io
import logging

# Configure logging for debugging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Define symbolic variables for polynomial and linear system solving
x, y = sp.symbols('x y')

# Initialize Pix2Text model globally to avoid reloading
try:
    p2t_model = Pix2Text.from_config()
    logger.info("Pix2Text model loaded successfully")
except Exception as e:
    logger.error(f"Failed to load Pix2Text model: {e}")
    p2t_model = None

def clean_latex_expression(latex_str):
    """Clean and normalize LaTeX expression for better parsing"""
    if not latex_str:
        return ""
    
    # Remove common LaTeX artifacts
    latex_str = latex_str.strip()
    latex_str = re.sub(r'\\[a-zA-Z]+\{([^}]*)\}', r'\1', latex_str)  # Remove LaTeX commands
    latex_str = re.sub(r'\\\\', r'\\', latex_str)  # Fix double backslashes
    latex_str = re.sub(r'\s+', ' ', latex_str)  # Normalize whitespace
    
    return latex_str

def parse_equation_type(latex_str):
    """Determine if the equation is polynomial or linear system"""
    try:
        # Clean the latex string
        cleaned = clean_latex_expression(latex_str)
        
        # Check if it contains system indicators
        if '\\\\' in cleaned or '\n' in cleaned or 'system' in cleaned.lower():
            return 'linear_system'
        
        # Check for polynomial indicators
        if 'x^' in cleaned or '^' in cleaned:
            return 'polynomial'
        
        # Check for linear equation indicators
        if ('x' in cleaned and 'y' in cleaned) or ('=' in cleaned and any(op in cleaned for op in ['+', '-'])):
            # Could be either, try to parse as polynomial first
            try:
                expr = sp.sympify(cleaned.split('=')[0] if '=' in cleaned else cleaned)
                if x in expr.free_symbols:
                    degree = sp.degree(expr, x)
                    if degree > 1:
                        return 'polynomial'
                    elif degree == 1 and y in expr.free_symbols:
                        return 'linear_system'
                    else:
                        return 'polynomial'
            except:
                pass
        
        return 'polynomial'  # Default fallback
        
    except Exception as e:
        logger.error(f"Error determining equation type: {e}")
        return 'polynomial'

def extract_polynomial_coefficients(latex_str):
    """Extract polynomial coefficients from LaTeX string"""
    try:
        # Clean and parse the expression
        cleaned = clean_latex_expression(latex_str)
        if '=' in cleaned:
            cleaned = cleaned.split('=')[0].strip()
        
        # Try to parse with SymPy
        expr = sp.sympify(cleaned)
        
        if x not in expr.free_symbols:
            raise ValueError("No variable x found in expression")
        
        # Get polynomial degree
        degree = sp.degree(expr, x)
        
        # Extract coefficients
        poly = sp.Poly(expr, x)
        coeffs = [float(poly.coeff_monomial(x**i)) for i in range(degree, -1, -1)]
        
        return {
            "type": "polynomial",
            "degree": degree,
            "coeffs": " ".join(map(str, coeffs)),
            "latex": latex_str,
            "success": True
        }
        
    except Exception as e:
        logger.error(f"Error extracting polynomial coefficients: {e}")
        return {
            "type": "polynomial",
            "degree": 2,
            "coeffs": "1 0 0",  # Default quadratic
            "latex": latex_str,
            "success": False,
            "error": str(e)
        }

def extract_linear_system_coefficients(latex_str):
    """Extract linear system coefficients from LaTeX string"""
    try:
        # Clean and split equations
        cleaned = clean_latex_expression(latex_str)
        
        # Split by common separators
        equations = re.split(r'\\\\|\n|;', cleaned)
        if len(equations) < 2:
            # Try to detect two equations in one line
            equations = re.split(r'(?<=[0-9])\s*(?=[+-]?\s*[0-9]*[xy])', cleaned)
        
        if len(equations) < 2:
            raise ValueError("Could not find two equations in system")
        
        eq1_str = equations[0].strip()
        eq2_str = equations[1].strip()
        
        # Parse each equation
        def parse_linear_eq(eq_str):
            # Convert to standard form ax + by = c
            if '=' not in eq_str:
                raise ValueError("No equals sign found")
            
            left, right = eq_str.split('=')
            expr = sp.sympify(left) - sp.sympify(right)
            
            # Extract coefficients
            a = float(expr.coeff(x, 1)) if expr.coeff(x, 1) else 0
            b = float(expr.coeff(y, 1)) if expr.coeff(y, 1) else 0
            c = float(-expr.as_coefficients_dict()[1]) if 1 in expr.as_coefficients_dict() else 0
            
            return f"{a} {b} {c}"
        
        eq1_coeffs = parse_linear_eq(eq1_str)
        eq2_coeffs = parse_linear_eq(eq2_str)
        
        return {
            "type": "linear",
            "eq1_coeffs": eq1_coeffs,
            "eq2_coeffs": eq2_coeffs,
            "latex": latex_str,
            "success": True
        }
        
    except Exception as e:
        logger.error(f"Error extracting linear system coefficients: {e}")
        return {
            "type": "linear",
            "eq1_coeffs": "1 1 3",  # Default system
            "eq2_coeffs": "1 -1 1",
            "latex": latex_str,
            "success": False,
            "error": str(e)
        }

def extract_equation_from_image(image_file):
    """Extract equation from image using Pix2Text with improved error handling"""
    try:
        if p2t_model is None:
            return {
                "type": "error",
                "latex": "❌ Pix2Text model not loaded. Please check installation.",
                "success": False
            }
        
        if image_file is None:
            return {
                "type": "error",
                "latex": "❌ No image file provided.",
                "success": False
            }
        
        # Open and process the image
        if isinstance(image_file, str):
            image = Image.open(image_file)
        else:
            image = Image.open(image_file.name)
        
        # Convert to RGB if needed
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        logger.info(f"Processing image of size: {image.size}")
        
        # Extract text and formulas using Pix2Text
        result = p2t_model.recognize_text_formula(image)
        
        if not result or result.strip() == "":
            return {
                "type": "error",
                "latex": "❌ No text or formulas detected in the image.",
                "success": False
            }
        
        logger.info(f"Extracted text: {result}")
        
        # Determine equation type
        eq_type = parse_equation_type(result)
        
        if eq_type == 'polynomial':
            return extract_polynomial_coefficients(result)
        else:
            return extract_linear_system_coefficients(result)
            
    except Exception as e:
        logger.error(f"Error processing image: {e}")
        return {
            "type": "error",
            "latex": f"❌ Error processing image: {str(e)}",
            "success": False
        }

# Import solve functions from other modules
def solve_polynomial(degree, coeff_string, real_only):
    """Solve polynomial equation (imported from polynomial.py logic)"""
    try:
        coeffs = list(map(float, coeff_string.strip().split()))
        if len(coeffs) != degree + 1:
            return f"⚠️ Please enter exactly {degree + 1} coefficients.", None, None

        poly = sum([coeffs[i] * x**(degree - i) for i in range(degree + 1)])
        simplified = sp.simplify(poly)
        factored = sp.factor(simplified)
        roots = sp.solve(sp.Eq(simplified, 0), x)

        if real_only:
            roots = [r for r in roots if sp.im(r) == 0]

        roots_output = "$$\n" + "\\ ".join(
            [f"r_{{{i}}} = {sp.latex(sp.nsimplify(r, rational=True))}" for i, r in enumerate(roots, 1)]
        ) + "\n$$"

        steps_output = f"""
### 🧐 Polynomial Expression
$$ {sp.latex(poly)} = 0 $$
### ✏️ Simplified
$$ {sp.latex(simplified)} = 0 $$
### 🤩 Factored
$$ {sp.latex(factored)} = 0 $$
### 🥮 Roots {'(Only Real)' if real_only else '(All Roots)'}
{roots_output}
        """

        x_vals = np.linspace(-10, 10, 400)
        y_vals = np.polyval(coeffs, x_vals)

        fig, ax = plt.subplots(figsize=(6, 4))
        ax.plot(x_vals, y_vals, label="Polynomial", color="blue")
        ax.axhline(0, color='black', linewidth=0.5)
        ax.axvline(0, color='black', linewidth=0.5)
        ax.grid(True)
        ax.set_title("📈 Graph of the Polynomial")
        ax.set_xlabel("x")
        ax.set_ylabel("f(x)")
        ax.legend()

        return steps_output, fig, ""
    except Exception as e:
        return f"❌ Error: {e}", None, ""

def solve_linear_system_from_coeffs(eq1_str, eq2_str):
    """Solve linear system (imported from linear.py logic)"""
    try:
        coeffs1 = list(map(float, eq1_str.strip().split()))
        coeffs2 = list(map(float, eq2_str.strip().split()))

        if len(coeffs1) != 3 or len(coeffs2) != 3:
            return "⚠️ Please enter exactly 3 coefficients for each equation.", None, None, None

        a1, b1, c1 = coeffs1
        a2, b2, c2 = coeffs2

        eq1 = sp.Eq(a1 * x + b1 * y, c1)
        eq2 = sp.Eq(a2 * x + b2 * y, c2)

        sol = sp.solve([eq1, eq2], (x, y), dict=True)
        if not sol:
            return "❌ No unique solution.", None, None, None

        solution = sol[0]
        eq_latex = f"$$ {sp.latex(eq1)} \\ {sp.latex(eq2)} $$"

        steps = rf"""
### 📌 Step-by-step Solution
1. **Original Equations:**
   $$ {sp.latex(eq1)} $$
   $$ {sp.latex(eq2)} $$
2. **Standard Form:** Already provided.
3. **Solve using SymPy `solve`:** Internally applies substitution/elimination.
4. **Solve for `x` and `y`:**
   $$ x = {sp.latex(solution[x])}, \quad y = {sp.latex(solution[y])} $$
5. **Verification:** Substitute back into both equations."""

        x_vals = np.linspace(-10, 10, 400)
        f1 = sp.solve(eq1, y)
        f2 = sp.solve(eq2, y)

        fig, ax = plt.subplots()
        if f1:
            f1_func = sp.lambdify(x, f1[0], modules='numpy')
            ax.plot(x_vals, f1_func(x_vals), label=sp.latex(eq1))
        if f2:
            f2_func = sp.lambdify(x, f2[0], modules='numpy')
            ax.plot(x_vals, f2_func(x_vals), label=sp.latex(eq2))

        ax.plot(solution[x], solution[y], 'ro', label=f"Solution ({solution[x]}, {solution[y]})")
        ax.axhline(0, color='black', linewidth=0.5)
        ax.axvline(0, color='black', linewidth=0.5)
        ax.legend()
        ax.set_title("📊 Graph of the Equations")
        ax.grid(True)

        return eq_latex, steps, fig, ""
    except Exception as e:
        return f"❌ Error: {e}", None, None, None

def solve_extracted_equation(eq_data, real_only):
    """Route to appropriate solver based on equation type"""
    if eq_data["type"] == "polynomial":
        return solve_polynomial(eq_data["degree"], eq_data["coeffs"], real_only)
    elif eq_data["type"] == "linear":
        return solve_linear_system_from_coeffs(eq_data["eq1_coeffs"], eq_data["eq2_coeffs"])
    else:
        return "❌ Unknown equation type", None, ""

def image_tab():
    """Create the Image Upload Solver tab with improved functionality"""
    with gr.Tab("📷 Image Upload Solver"):
        gr.Markdown("## 📷 Solve Equations from Image")
        
        with gr.Row():
            # File input for uploading images
            image_input = gr.File(
                label="Upload Question Image",
                file_types=[".pdf", ".png", ".jpg", ".jpeg"],
                file_count="single"
            )
            # Button to trigger image processing
            image_upload_btn = gr.Button("📤 Process Image")
        
        gr.Markdown("**Supported Formats:** .pdf, .png, .jpg, .jpeg")

        with gr.Row():
            # Checkbox to toggle real roots only for polynomials
            real_image_checkbox = gr.Checkbox(label="Show Only Real Roots (for Polynomials)", value=False)
            # Button to preview the extracted equation
            preview_image_btn = gr.Button("🔍 Preview Equation")

        # Markdown component to display the extracted equation
        image_equation_display = gr.Markdown()
        
        with gr.Row():
            # Button to confirm and display the solution (initially hidden)
            confirm_image_btn = gr.Button("✅ Display Solution", visible=False)
            # Button to edit the equation manually (initially hidden)
            edit_image_btn = gr.Button("✏️ Make Changes Manually", visible=False)

        # Textbox for manual LaTeX editing (initially hidden)
        edit_latex_input = gr.Textbox(label="Edit LaTeX Equation", visible=False, lines=3)
        # Button to save manual changes (initially hidden)
        save_edit_btn = gr.Button("💾 Save Changes", visible=False)

        # Markdown component to display solution steps
        image_steps_md = gr.Markdown()
        # Plot component to display the graph
        image_plot_output = gr.Plot()
        # Textbox to display errors
        image_error_box = gr.Textbox(label="Status", visible=True, interactive=False)

        # State to store the extracted equation data
        extracted_eq_state = gr.State()

        def handle_image_upload(image_file):
            """Handle image upload and initial processing"""
            if image_file is None:
                return "⚠️ Please upload an image.", None, "", None, None
            
            try:
                eq_data = extract_equation_from_image(image_file)
                if eq_data["success"]:
                    status = "✅ Image processed successfully. Click 'Preview Equation' to see the extracted equation."
                else:
                    status = f"⚠️ Processing completed with issues: {eq_data.get('error', 'Unknown error')}"
                
                return status, eq_data, "", None, None
            except Exception as e:
                return f"❌ Error processing image: {str(e)}", None, "", None, None

        # Event handler for image upload button
        image_upload_btn.click(
            fn=handle_image_upload,
            inputs=[image_input],
            outputs=[image_error_box, extracted_eq_state, image_equation_display, 
                    image_steps_md, image_plot_output]
        )

        def preview_image_equation(eq_data, real_only):
            """Preview the extracted equation"""
            if eq_data is None:
                return ("⚠️ No equation data available. Please upload and process an image first.", 
                       gr.update(visible=False), gr.update(visible=False), "", None)
            
            if eq_data["type"] == "error":
                return (eq_data["latex"], gr.update(visible=False), gr.update(visible=False), "", None)
            
            # Create preview display
            if eq_data["type"] == "polynomial":
                eq_type_display = "Polynomial Equation"
                details = f"Degree: {eq_data['degree']}, Coefficients: {eq_data['coeffs']}"
            else:
                eq_type_display = "Linear System"
                details = f"Equation 1: {eq_data['eq1_coeffs']}, Equation 2: {eq_data['eq2_coeffs']}"
            
            preview_text = f"""
### ✅ Confirm {eq_type_display}

**Extracted LaTeX:** {eq_data['latex']}

**Parsed Details:** {details}

**Status:** {'✅ Successfully parsed' if eq_data.get('success', True) else '⚠️ Parsing had issues but proceeding with defaults'}
            """
            
            return (preview_text, gr.update(visible=True), gr.update(visible=True), "", None)

        # Event handler for preview button
        preview_image_btn.click(
            fn=preview_image_equation,
            inputs=[extracted_eq_state, real_image_checkbox],
            outputs=[image_equation_display, confirm_image_btn, edit_image_btn, 
                    image_steps_md, image_plot_output]
        )

        def confirm_image_solution(eq_data, real_only):
            """Confirm and solve the extracted equation"""
            if eq_data is None or eq_data["type"] == "error":
                return "⚠️ No valid equation to solve.", None, "No equation available"
            
            try:
                if eq_data["type"] == "polynomial":
                    steps, plot, error = solve_polynomial(eq_data["degree"], eq_data["coeffs"], real_only)
                    return steps, plot, error if error else "✅ Solution completed"
                elif eq_data["type"] == "linear":
                    eq_latex, steps, plot, error = solve_linear_system_from_coeffs(
                        eq_data["eq1_coeffs"], eq_data["eq2_coeffs"])
                    return steps, plot, error if error else "✅ Solution completed"
                else:
                    return "❌ Unknown equation type", None, "Unknown equation type"
            except Exception as e:
                return f"❌ Error solving equation: {str(e)}", None, str(e)

        # Event handler for confirm button
        confirm_image_btn.click(
            fn=confirm_image_solution,
            inputs=[extracted_eq_state, real_image_checkbox],
            outputs=[image_steps_md, image_plot_output, image_error_box]
        )

        def enable_manual_edit(eq_data):
            """Enable manual editing of the equation"""
            if eq_data is None:
                latex_value = "No equation to edit. Please upload an image first."
            elif eq_data["type"] == "error":
                latex_value = "Error in extraction. Please enter your equation manually."
            else:
                latex_value = eq_data.get("latex", "")
            
            return (gr.update(visible=True, value=latex_value), 
                   gr.update(visible=True), 
                   gr.update(visible=False), 
                   gr.update(visible=False))

        # Event handler for edit button
        edit_image_btn.click(
            fn=enable_manual_edit,
            inputs=[extracted_eq_state],
            outputs=[edit_latex_input, save_edit_btn, confirm_image_btn, edit_image_btn]
        )

        def save_manual_changes(latex_input, real_only):
            """Save manual changes and solve"""
            try:
                if not latex_input or latex_input.strip() == "":
                    return "⚠️ Please enter a valid equation.", None, "Empty input"
                
                # Determine equation type from manual input
                eq_type = parse_equation_type(latex_input)
                
                if eq_type == 'polynomial':
                    eq_data = extract_polynomial_coefficients(latex_input)
                    steps, plot, error = solve_polynomial(eq_data["degree"], eq_data["coeffs"], real_only)
                else:
                    eq_data = extract_linear_system_coefficients(latex_input)
                    eq_latex, steps, plot, error = solve_linear_system_from_coeffs(
                        eq_data["eq1_coeffs"], eq_data["eq2_coeffs"])
                
                return steps, plot, error if error else "✅ Manual equation solved"
                
            except Exception as e:
                return f"❌ Error parsing manual input: {str(e)}", None, str(e)

        # Event handler for save button
        save_edit_btn.click(
            fn=save_manual_changes,
            inputs=[edit_latex_input, real_image_checkbox],
            outputs=[image_steps_md, image_plot_output, image_error_box]
        )

    return (image_input, image_upload_btn, real_image_checkbox, preview_image_btn,
            image_equation_display, confirm_image_btn, edit_image_btn, edit_latex_input,
            save_edit_btn, image_steps_md, image_plot_output, image_error_box, extracted_eq_state)