import streamlit as st
from sympy import sympify, Eq, solve
import pytesseract
from PIL import Image
import easyocr
import numpy as np
from transformers import pipeline

# Initialize OCR reader and NLP models
easyocr_reader = easyocr.Reader(['en'])
fill_blanks_model = pipeline("fill-mask", model="bert-base-uncased")
qa_model = pipeline("question-answering")

# Streamlit app title
st.title("Quiz and Numerical Problem Solver")
st.markdown("**Input a question or problem as text or upload an image, and get solutions!**")

# Sidebar options
option = st.sidebar.selectbox("Select Input Type:", ("Text Input", "Image Upload"))

if option == "Text Input":
    # Text input
    user_input = st.text_area("Enter your question or problem:")
    if st.button("Solve"):
        if user_input.strip():
            try:
                st.subheader("Solution:")

                # Handle blanks (fill-in-the-blank questions)
                if "___" in user_input or "[MASK]" in user_input:
                    st.write("Detected fill-in-the-blank question:")
                    results = fill_blanks_model(user_input)
                    for result in results:
                        st.write(f"{result['sequence']} (Confidence: {result['score']:.2f})")

                # Handle MCQs
                elif "?" in user_input and any(option in user_input.lower() for option in ["a.", "b.", "c.", "d."]):
                    st.write("Detected multiple-choice question:")
                    question, *options = user_input.split("\n")
                    options = [opt.strip() for opt in options if opt.strip()]
                    answer = qa_model(question=question, context=" ".join(options))
                    st.write(f"Answer: {answer['answer']} (Confidence: {answer['score']:.2f})")

                # Handle mathematical expressions
                else:
                    try:
                        expr = sympify(user_input)  # Parse the input into a symbolic expression
                        if isinstance(expr, Eq):  # If it's an equation, solve it
                            solution = solve(expr)
                            st.write("Solutions:", solution)
                        else:  # Otherwise, solve the expression
                            solution = solve(expr)
                            st.write("Solution:", solution)
                    except Exception as e:
                        st.error(f"Error processing input: {str(e)}")

            except Exception as e:
                st.error("Error processing input. Please ensure it's a valid mathematical, quiz, or problem question.")
        else:
            st.error("Please input a valid text to solve.")

elif option == "Image Upload":
    # Image upload
    uploaded_image = st.file_uploader("Upload an image containing the problem:", type=["png", "jpg", "jpeg"])
    if uploaded_image:
        image = Image.open(uploaded_image)
        st.image(image, caption="Uploaded Image", use_column_width=True)

        if st.button("Extract & Solve"):
            # Extract text using OCR
            with st.spinner("Extracting text from image..."):
                try:
                    extracted_text = easyocr_reader.readtext(np.array(image), detail=0)
                    st.subheader("Extracted Text:")
                    full_text = "\n".join(extracted_text)
                    st.text(full_text)

                    # Attempt to process the extracted text
                    st.subheader("Solution:")

                    # Handle blanks, MCQs, or math dynamically
                    if "___" in full_text or "[MASK]" in full_text:
                        st.write("Detected fill-in-the-blank question:")
                        results = fill_blanks_model(full_text)
                        for result in results:
                            st.write(f"{result['sequence']} (Confidence: {result['score']:.2f})")
                    elif "?" in full_text and any(option in full_text.lower() for option in ["a.", "b.", "c.", "d."]):
                        st.write("Detected multiple-choice question:")
                        question, *options = full_text.split("\n")
                        options = [opt.strip() for opt in options if opt.strip()]
                        answer = qa_model(question=question, context=" ".join(options))
                        st.write(f"Answer: {answer['answer']} (Confidence: {answer['score']:.2f})")
                    else:
                        try:
                            expr = sympify(full_text)  # Parse the extracted text
                            if isinstance(expr, Eq):  # If it's an equation, solve it
                                solution = solve(expr)
                                st.write("Solutions:", solution)
                            else:  # Otherwise, solve the expression
                                solution = solve(expr)
                                st.write("Solution:", solution)
                        except Exception as e:
                            st.error(f"Error processing extracted text: {str(e)}")

                except Exception as e:
                    st.error("Error solving the problem from the extracted text. Ensure the image contains a valid problem.")

# Additional Notes
st.markdown("---")
st.markdown(
    "This app uses OCR for text extraction, symbolic computation for solving problems, and NLP for fill-in-the-blank and MCQ questions."
)