import streamlit as st
from PIL import Image
from texify.inference import batch_inference
from texify.model.model import load_model
from texify.model.processor import load_processor
from texify.output import replace_katex_invalid
from sympy import symbols, simplify, diff, integrate, Eq, Derivative, Integral, solve, plot
from sympy.parsing.latex import parse_latex
from sympy.abc import x, y, z, t
import matplotlib.pyplot as plt
import re
import io

# ✅ Page Configuration
st.set_page_config(
    page_title="Texify Math Solver",
    layout="centered",
    page_icon="🧠"
)

# ✅ Cache model and processor
@st.cache_resource
def load_texify():
    model = load_model()
    processor = load_processor()
    return model, processor

model, processor = load_texify()

# ✅ SymPy Solver with Extended Support
def solve_with_sympy(latex_expr):
    steps = []
    try:
        latex_expr = latex_expr.strip()
        latex_expr = re.sub(r'^\$+|\$+$', '', latex_expr)

        if len(latex_expr) < 3:
            return ["⚠️ Expression too short to parse."], None

        parsed_expr = parse_latex(latex_expr)
        steps.append("🧩 **Parsed Expression:** $" + latex_expr + "$")

        # ✅ Integral Handling
        if isinstance(parsed_expr, Integral) or str(parsed_expr).startswith("Integral"):
            result = parsed_expr.doit()
            steps.append("∫ **Integral Result:** " + str(result))
            steps.append("🔎 **Approximate Value:** " + str(result.evalf()))
            return steps, None

        # ✅ Derivative Handling
        elif isinstance(parsed_expr, Derivative) or str(parsed_expr).startswith("Derivative") or "diff" in latex_expr:
            result = diff(parsed_expr, x)
            deriv_expr = Derivative(parsed_expr, x)
            steps.append("📘 **Derivative Setup:** " + str(deriv_expr) + " = " + str(result))
            steps.append("🧮 **Derivative:** " + str(result))
            return steps, None

        # ✅ Equation Solving
        elif "=" in latex_expr:
            lhs, rhs = latex_expr.split("=")
            eq = Eq(parse_latex(lhs), parse_latex(rhs))
            result = solve(eq)
            steps.append("🧮 **Equation Solved:** " + str(eq))
            steps.append("✅ **Solution:** " + str(result))
            return steps, None

        # ✅ Simplify + Plot Expression
        else:
            result = simplify(parsed_expr)
            steps.append("🧹 **Simplified:** " + str(result))
            steps.append("🔎 **Approximate Value:** " + str(result.evalf()))

            # ✅ Plotting
            fig, ax = plt.subplots()
            p = plot(parsed_expr, (x, -10, 10), show=False)
            for expr in p:
                expr.line_color = 'blue'
                expr.label = str(parsed_expr)
                expr.plot(ax)
            ax.set_title("📉 Graph of the Function")
            ax.set_xlabel("x")
            ax.set_ylabel("f(x)")
            ax.legend()
            ax.grid(True)

            return steps, fig

    except Exception as e:
        return ["❌ Error solving expression: " + str(e),
                "ℹ️ Tip: OCR may miss '^', brackets, or grouping. Please check manually."], None

# ✅ Streamlit Interface
st.title("🧠 Texify Math OCR + Solver")
st.markdown(
    """
    <div style='text-align: center; color: #4A6FA5;'>
        📚 Upload a math image. Texify will extract LaTeX, solve it, and show plots!
    </div>
    """,
    unsafe_allow_html=True
)

uploaded_file = st.file_uploader("📤 Upload Image (PNG, JPG, WEBP)", type=["png", "jpg", "jpeg", "webp"])

if uploaded_file:
    image = Image.open(uploaded_file).convert("RGB")
    st.image(image, caption="🖼️ Uploaded Image", use_column_width=True)

    with st.spinner("🔍 Running OCR with Texify..."):
        ocr_result = batch_inference([image], model, processor)[0]
        display_latex = replace_katex_invalid(ocr_result)

    st.markdown("---")
    st.markdown("### ✍️ **OCR Result**")
    try:
        st.latex(display_latex.replace("$$", "").strip())
    except:
        st.markdown(f"<span style='color:red;'>OCR output could not be rendered as LaTeX.</span>", unsafe_allow_html=True)

    st.code(ocr_result)

    st.markdown("### ✏️ **Edit LaTeX (if needed)**")
    corrected_latex = st.text_input("✏️ Edit LaTeX Expression", value=ocr_result.strip())
    cleaned_latex = re.sub(r'^\$+|\$+$', '', corrected_latex.strip())

    st.markdown("### 👀 Preview")
    st.latex(cleaned_latex)

    if st.button("🔎 Solve Now"):
        st.markdown("---")
        st.markdown("### 🧠 **Solution Steps**")
        steps, fig = solve_with_sympy(cleaned_latex)
        for step in steps:
            st.markdown(step, unsafe_allow_html=True)

        if fig:
            st.pyplot(fig)