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
import os

# Optional: import training function
from train import train_model

# Trigger data download only once
if not os.path.exists("dataset/train"):
    print("🚀 Running data preparation scripts...")
    os.system("python download_data.py")
    os.system("python generate_csv.py")


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


# Load Pix2Tex model
model = LatexOCR()


# Clean LaTeX output
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


# Solver logic
def solve_polynomial(image):
    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"
        output += "---\n"
        output += f"## 🧹 Cleaned LaTeX Used\n```\n{cleaned_latex}\n```\n"
        output += "---\n"
        output += f"## 🧠 Parsed Expression\n\n$$ {sp.latex(expr)} $$\n"
        output += "---\n"

        if isinstance(expr, sp.Equality):
            lhs = expr.lhs - expr.rhs
            output += "## ✏️ Step 1: Standard Form of the Polynomial\n"
            output += f"$$ {sp.latex(lhs)} = 0 $$\n"
            output += "---\n"

            output += "## 🧩 Step 2: Factor the Polynomial\n"
            factored = sp.factor(lhs)
            output += f"$$ {sp.latex(factored)} = 0 $$\n"
            output += "---\n"

            output += "## ✅ Step 3: Solve for 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}}}:\\quad {var} &= {sp.latex(val)} \\\\\n"
                output += "\\end{aligned}\n$$\n"
        else:
            simplified = sp.simplify(expr)
            output += "## ➕ Simplified Expression\n"
            output += f"$$ {sp.latex(simplified)} $$"

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


# Trigger training
def run_training():
    try:
        train_model("train.yaml")  # path to your training config
        return "✅ Training completed successfully."
    except Exception as e:
        return f"❌ Training failed: {str(e)}"


# Gradio UI with training button
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Polynomial Solver from Handwritten Image")
    gr.Markdown("Upload a handwritten polynomial image. The app will extract and solve it step-by-step.")

    with gr.Row():
        image_input = gr.Image(type="pil", label="📷 Upload Polynomial Image")
        solution_output = gr.Markdown(label="📋 Step-by-step Solution")

    image_input.change(fn=solve_polynomial, inputs=image_input, outputs=solution_output)

    gr.Markdown("----")

    gr.Markdown("## 🛠 Optional: Fine-tune Pix2Tex (CPU)")
    train_btn = gr.Button("🎓 Start CPU Training")
    train_output = gr.Textbox(label="Training Status")

    train_btn.click(fn=run_training, outputs=train_output)

demo.launch()