import os
import gradio as gr
from PIL import Image
import sympy as sp
from sympy.parsing.latex import parse_latex
import re

try:
    from pix2tex.cli import LatexOCR
    if os.path.exists("trained_model"):
        model = LatexOCR(weights="trained_model")
    else:
        model = None
except Exception as e:
    model = None

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

def clean_latex(latex):
    latex = re.sub(r'\\(cal|mathcal)\s*X', 'x', latex)
    latex = latex.replace('{', '').replace('}', '')
    latex = latex.strip().rstrip(',.')
    latex = re.sub(r'(\d+)\s*\\pi', r'(\1*3.1416)', latex)
    latex = latex.replace(r'\pi', '3.1416')
    latex = re.sub(r'(\d+)\s*e', r'(\1*2.7183)', latex)
    latex = re.sub(r'(?<![a-zA-Z0-9])e(?![a-zA-Z0-9])', '2.7183', 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*([xX](\^\d+)?)', r'(\1)*\2', latex)
    if '=' not in latex:
        latex += '=0'
    return latex

def solve_polynomial(image):
    if model is None:
        return "❌ No trained model found. Please run training first to create `trained_model/`."

    try:
        img = preprocess_handwritten_image(image)
        latex_result = model(img)
        cleaned_latex = clean_latex(latex_result)
        expr = parse_latex(cleaned_latex)

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

        if isinstance(expr, sp.Equality):
            lhs = expr.lhs - expr.rhs
            output += "## ✏️ Step 1: Standard Form\n"
            output += f"$$ {sp.latex(lhs)} = 0 $$\n---\n"
            output += "## 🧩 Step 2: Factorized\n"
            output += f"$$ {sp.latex(sp.factor(lhs))} = 0 $$\n---\n"
            output += "## ✅ Step 3: Roots\n"
            roots = sp.solve(sp.Eq(lhs, 0), dict=True)
            if roots:
                output += "$$\n\\begin{aligned}\n"
                for i, sol in enumerate(roots, 1):
                    for var, val in sol.items():
                        output += f"\\text{{Root {i}}}: {var} &= {sp.latex(val)}\\\\\n"
                output += "\\end{aligned}\n$$"
        else:
            output += "## ➕ Simplified\n"
            output += f"$$ {sp.latex(sp.simplify(expr))} $$"

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

demo = gr.Interface(
    fn=solve_polynomial,
    inputs=gr.Image(type="pil", label="📷 Upload Image"),
    outputs=gr.Markdown(label="📋 Solution"),
    title="🧠 Handwritten Polynomial Solver",
    description="Upload a handwritten polynomial. The app extracts, solves, and explains it step-by-step.",
    allow_flagging="never"
)

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