app.py

import gradio as gr
from PIL import Image
from pix2tex.cli import LatexOCR
import sympy as sp
from sympy.parsing.latex import parse_latex
import re

model = LatexOCR()

def preprocess_handwritten_image(pil_img):
    return pil_img.convert('RGB')

# Tab 1 cleaning (KEEP UNCHANGED)
def clean_latex(latex):
    latex = latex.replace('\\ ', '')
    latex = latex.replace('\\\\', '\\')
    latex = re.sub(r'\\[ \t\n\r\f\v]*', '', latex)
    latex = re.sub(r'\\([+\-=])', r'\1', latex)
    replacements = {
        r'\\chi': 'x', r'chi': 'x',
        r'\\xi': 'x', r'xi': 'x',
        r'\\alpha': 'x', r'alpha': 'x',
        r'\\beta': 'b', r'beta': 'b',
        r'\\gamma': 'y', r'gamma': 'y',
        r'\\vartheta': '3', r'vartheta': '3',
        r'\\mathcalW': 'x', r'mathcalW': 'x',
        r'\\pi': 'pi', r'pi': 'pi',
        r'\\mathrm': '', r'mathrm': '',
    }
    for wrong, correct in replacements.items():
        latex = re.sub(wrong, correct, latex)
    latex = re.sub(r'\\(cal|mathcal)\s*\{?\s*[Xx]\s*\}?', 'x', latex)
    latex = re.sub(r'\\(cal|mathcal)\s*\{?\s*[Yy]\s*\}?', 'y', latex)
    latex = re.sub(r'\\(cal|mathcal)\s*\{?\s*[Zz]\s*\}?', 'z', latex)
    latex = latex.replace('cal x', 'x').replace('cal X', 'x')
    latex = latex.replace('mathcal x', 'x').replace('mathcal X', 'x')
    latex = latex.replace('{', '').replace('}', '')
    latex = latex.strip().rstrip(',.')
    latex = re.sub(r'(?<![a-zA-Z0-9])e(?![a-zA-Z0-9])', 'E', latex)
    latex = re.sub(r'(\d)([a-zA-Z])', r'\1*\2', latex)
    latex = re.sub(r'(\d+)\s*i', r'\1*I', latex)
    latex = re.sub(r'(?<![a-zA-Z0-9])i(?![a-zA-Z0-9])', 'I', latex)
    latex = re.sub(r'\(([^()]+?)\)\s*([a-zA-Z](\^\d+)?)', r'(\1)*\2', latex)
    latex = latex.replace(r'\cdot', '*')
    latex = latex.replace('−', '-')
    latex = re.sub(r'[^\w\s^=+*\-().]', '', latex)
    if '=' not in latex:
        latex += '=0'
    latex = latex.replace('pi', '3.1416')
    latex = latex.replace('e', '2.7183')
    latex = latex.replace('E', '2.7183')
    return latex

# ✅ Tab 2 specific cleaner
def clean_latex2(latex):
    latex = latex.replace('\\ ', '')
    latex = latex.replace('\\\\', '\\')
    latex = re.sub(r'\\[ \t\n\r\f\v]*', '', latex)
    latex = latex.replace(r'\times', 'x')
    latex = latex.replace(r'\cdot', '*')
    latex = latex.replace('−', '-').replace('–', '-')

    # Remove LaTeX wrappers like \mathcal{}, \mathbf{}, etc.
    latex = re.sub(r'\\(text|mathbf|mathrm|mathit|textbf|mathcal|cal)\s*\{([^{}]+)\}', r'\2', latex)

    # Remove remaining LaTeX commands (like \!\!\! and \,)
    latex = re.sub(r'\\[a-zA-Z]+', '', latex)

    # Remove nested or empty braces
    latex = re.sub(r'\{+\}+', '', latex)
    latex = re.sub(r'\{([^\{\}]*)\}', r'\1', latex)

    # Remove extra commas, semicolons, weird formatting characters
    latex = re.sub(r'[;,]', '\n', latex)
    latex = re.sub(r'[^\w\s=+\-*/().]', '', latex)

    # Replace multiple spaces with single
    latex = re.sub(r'\s+', ' ', latex)

    # Normalize equals
    latex = re.sub(r'(?<![=<>])=(?![=<>])', ' = ', latex)

    # Replace ambiguous variable artifacts
    replacements = {
        'chi': 'x', 'xi': 'x', 'alpha': 'x', 'beta': 'b',
        'gamma': 'y', 'vartheta': '3', 'mathcal': '',
        'cal': '', 'mathrm': ''
    }
    for wrong, right in replacements.items():
        latex = re.sub(wrong, right, latex)

    return latex.strip()

def solve_polynomial(image):
    try:
        img = preprocess_handwritten_image(image)
        latex_result = model(img)
        if not latex_result or len(latex_result.strip()) < 2:
            return "❌ Could not extract valid LaTeX from image.", "", ""
        cleaned_latex = clean_latex(latex_result)
        try:
            expr = parse_latex(cleaned_latex)
            expr = expr.subs(sp.pi, sp.Float(3.1416)).subs(sp.E, sp.Float(2.7183))
            if not isinstance(expr, sp.Equality):
                raise ValueError("Expression is not an equation.")
            lhs = expr.lhs - expr.rhs
            if not lhs.is_polynomial():
                raise ValueError("Not a polynomial")
        except:
            return f"❌ Could not parse expression:\n\n```latex\n{cleaned_latex}\n```", cleaned_latex, ""

        output = f"## 📄 Extracted LaTeX\n```latex\n{latex_result}\n```\n"
        output += f"---\n## 🧹 Cleaned LaTeX\n```latex\n{cleaned_latex}\n```\n"
        output += f"---\n## 🧠 Parsed Expression\n$$ {sp.latex(expr)} $$\n"

        lhs = expr.lhs - expr.rhs
        factor = sp.factor(lhs)
        output += f"---\n## ✏️ Standard Form\n$$ {sp.latex(lhs)} = 0 $$\n"
        output += f"---\n## 🧩 Factorized\n$$ {sp.latex(factor)} = 0 $$\n"

        x = next(iter(lhs.free_symbols))
        roots = sp.solve(lhs, x)
        output += "## ✅ Roots\n"
        output += "$$\n\\begin{aligned}\n"
        for i, r in enumerate(roots, 1):
            root_val = sp.N(r, 6)
            output += f"\\text{{Root {i}}}:\\quad x &\\approx {sp.latex(root_val)} \\\\\n"
        output += "\\end{aligned}\n$$\n"

        return output, cleaned_latex, ""
    except Exception as e:
        return f"❌ Error: {str(e)}", "", ""

def wrapped_solver(img):
    result, cleaned, prompt = solve_polynomial(img)
    return result, cleaned, prompt

# ✅ FIXED: moved outside of gr.Blocks()
def solve_system_of_equations(image):
    try:
        img = preprocess_handwritten_image(image)
        raw_latex = model(img)
        if not raw_latex or len(raw_latex.strip()) < 3:
            return "❌ Could not extract valid LaTeX."

        # Handle array environment
        if r"\begin{array" in raw_latex:
            lines = re.findall(r'\\begin{array}{[^}]+}(.+?)\\end{array}', raw_latex, re.DOTALL)
            if lines:
                raw_latex = lines[0]
            raw_latex = raw_latex.replace(r'\\', '\n')
            raw_latex = re.sub(r'&', '', raw_latex)

        raw_latex = raw_latex.replace("\n", " ")
        equations = re.split(r'[;,]', raw_latex)

        cleaned_equations = []
        for eq in equations:
            cleaned = clean_latex2(eq)
            cleaned = re.sub(r'(?<![=<>])=(?![=<>])', ' = ', cleaned)
            if '=' in cleaned:
                try:
                    parsed = parse_latex(cleaned)
                    if isinstance(parsed, sp.Equality):
                        cleaned_equations.append(cleaned)
                except Exception:
                    continue

        sympy_eqs = []
        symbols = set()
        for eq in cleaned_equations:
            parsed = parse_latex(eq)
            if isinstance(parsed, sp.Equality):
                sympy_eqs.append(parsed)
                symbols.update(parsed.free_symbols)

        if len(sympy_eqs) < 2:
            return f"❌ Not enough valid equations detected.\n\n```latex\n{raw_latex}\n```"

        sol = sp.solve(sympy_eqs, list(symbols), dict=True)
        output = f"## 🧾 Detected Equations:\n```latex\n{raw_latex}\n```\n"
        output += "---\n## ✏️ Parsed & Cleaned:\n"
        for eq in sympy_eqs:
            output += f"$$ {sp.latex(eq)} $$\n"
        output += "---\n## ✅ Solution:\n"
        if sol:
            output += "$$" + sp.latex(sol[0]) + "$$"
        else:
            output += "❌ No solution found or system may be inconsistent."
        return output
    except Exception as e:
        return f"❌ Error: {str(e)}"

# === UI ===
with gr.Blocks() as demo:
    with gr.Tab("🖼️ Parse from Image"):
        image_input = gr.Image(type="pil", label="📷 Upload Image of Polynomial")
        hidden_latex = gr.Textbox(visible=False)
        explanation_prompt = gr.Textbox(visible=False)
        output_box = gr.Markdown(label="📋 Step-by-step Solution")
        submit_btn = gr.Button("🔍 Solve")
        submit_btn.click(fn=wrapped_solver, inputs=[image_input], outputs=[output_box, hidden_latex, explanation_prompt])

    with gr.Tab("📐 Solve Simultaneous Equations"):
        sim_image_input = gr.Image(type="pil", label="📷 Upload Image with Simultaneous Equations")
        sim_output = gr.Markdown(label="📋 Solved System Output")
        sim_button = gr.Button("📌 Solve System")
        sim_button.click(fn=solve_system_of_equations, inputs=[sim_image_input], outputs=[sim_output])

if __name__ == "__main__":
    demo.launch()