import gradio as gr
from codeexecutor import get_majority_vote, type_check, postprocess_completion, draw_polynomial_plot
import re
iterations = 4
# Function to generate mock predictions (as the model isn't loaded)
def get_prediction(question):
    return "Solve the following mathematical problem: what is the sum of polynomial 2x+3 and 3x?\n### Solution: To solve the problem of summing the polynomials \\(2x + 3\\) and \\(3x\\), we can follow these steps:\n\n1. Define the polynomials.\n2. Sum the polynomials.\n3. Simplify the resulting polynomial expression.\n\nThe sum of the polynomials \\(2x + 3\\) and \\(3x\\) is \\(\\boxed{5x + 3}\\).\n"

# Function to parse the prediction to extract the answer and steps
def parse_prediction(prediction):
    lines = prediction.strip().split('\n')
    answer = None
    steps = []
    for line in lines:
        match = re.match(r'^\s*(?:Answer|answer)\s*[:=]\s*(.*)', line)
        if match:
            answer = match.group(1).strip()
        else:
            steps.append(line)
    if answer is None:
        answer = lines[-1].strip()
        steps = lines
    steps_text = '\n'.join(steps).strip()
    return answer, steps_text

# Function to extract boxed answers
def extract_boxed_answer(text):
    match = re.search(r'\\boxed\{(.*?)\}', text)
    if match:
        return match.group(1)
    return None

# Function to perform majority voting and get steps
def majority_vote_with_steps(question, num_iterations=10):
    all_predictions = []
    all_answers = []
    steps_list = []

    for _ in range(num_iterations):
        prediction = get_prediction(question)
        answer, success = postprocess_completion(prediction, return_status=True, last_code_block=True)

        if success:
            all_predictions.append(prediction)
            all_answers.append(answer)
            steps_list.append(prediction)
        else:
            answer, steps = parse_prediction(prediction)
            all_predictions.append(prediction)
            all_answers.append(answer)
            steps_list.append(steps)

    if success:
        majority_voted_ans = get_majority_vote(all_answers)
        expression = majority_voted_ans
        if type_check(expression) == "Polynomial":
            plotfile = draw_polynomial_plot(expression)
    else:
        plotfile = None

    # Find the steps corresponding to the majority voted answer
    for i, ans in enumerate(all_answers):
        if ans == majority_voted_ans:
            steps_solution = steps_list[i]
            answer = parse_prediction(steps_solution)
            break
    else:
        answer = majority_voted_ans
        steps_solution = "No steps found"

    return answer, steps_solution, plotfile

# Function to handle chat-like interaction
def chat_interface(history, question):
    final_answer, steps_solution, plotfile = majority_vote_with_steps(question, iterations)
    history.append(("User", question))
    history.append(("MathBot", f"Answer: {final_answer}\nSteps:\n{steps_solution}"))
    return history, plotfile

# Gradio app setup using Blocks for layout management
with gr.Blocks() as interface:
    with gr.Column():
        chat_history = gr.Chatbot(label="Chat with MathBot", elem_id="chat_history")
        math_question = gr.Textbox(label="Your Question", placeholder="Ask a math question...", elem_id="math_question")
        chatbot_output = gr.Chatbot(label="Chat History")
        polynomial_plot = gr.Image(label="Polynomial Plot")
    
    math_question.submit(chat_interface, inputs=[chat_history, math_question], outputs=[chatbot_output, polynomial_plot])


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