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 cleaner (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 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('–', '-')
    latex = re.sub(r'\\(text|mathbf|mathrm|mathit|textbf|mathcal|cal)\s*\{([^{}]+)\}', r'\2', latex)
    latex = re.sub(r'\\[a-zA-Z]+', '', latex)
    latex = re.sub(r'\{+\}+', '', latex)
    latex = re.sub(r'\{([^\{\}]*)\}', r'\1', latex)
    latex = re.sub(r'[;,]', '\n', latex)
    latex = re.sub(r'[^\w\s=+\-*/().]', '', latex)
    latex = re.sub(r'\s+', ' ', latex)
    latex = re.sub(r'(?<![=<>])=(?![=<>])', ' = ', latex)
    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()

# --- Polynomial Solver (unchanged) ---
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)
        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")
        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"
        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)}", "", ""

# --- NEW Multi-image System Solver ---
def solve_system_multi(img1, img2, img3):
    images = [img1, img2, img3]
    cleaned_equations = []
    raw_outputs = []
    for img in images:
        if img is None:
            continue
        try:
            img = preprocess_handwritten_image(img)
            latex = model(img)
            raw_outputs.append(latex)
            cleaned = clean_latex2(latex)
            parsed = parse_latex(cleaned)
            if isinstance(parsed, sp.Equality):
                cleaned_equations.append(parsed)
        except:
            continue

    if len(cleaned_equations) < 2:
        return f"❌ Could not parse enough valid equations."

    symbols = set()
    for eq in cleaned_equations:
        symbols.update(eq.free_symbols)

    sol = sp.solve(cleaned_equations, list(symbols), dict=True)
    output = "## 📄 Raw LaTeX:\n"
    for latex in raw_outputs:
        output += f"- ```latex\n{latex}\n```\n"
    output += "---\n## ✏️ Parsed Equations:\n"
    for eq in cleaned_equations:
        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

# === 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=solve_polynomial, inputs=[image_input], outputs=[output_box, hidden_latex, explanation_prompt])

    with gr.Tab("📐 Solve System (Multi-Image)"):
        img1 = gr.Image(type="pil", label="📷 Equation 1")
        img2 = gr.Image(type="pil", label="📷 Equation 2")
        img3 = gr.Image(type="pil", label="📷 Equation 3 (Optional)")
        sys_output = gr.Markdown(label="📋 Solved System Output")
        sys_btn = gr.Button("📌 Solve System")
        sys_btn.click(fn=solve_system_multi, inputs=[img1, img2, img3], outputs=[sys_output])

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