import gradio as gr
import ctranslate2
from transformers import AutoTokenizer
from huggingface_hub import snapshot_download
from codeexecutor import get_majority_vote,type_check,postprocess_completion,draw_polynomial_plot


import re
import os
# Define the model and tokenizer loading
model_prompt = "Explain and solve the following mathematical problem step by step, showing all work: "
tokenizer = AutoTokenizer.from_pretrained("AI-MO/NuminaMath-7B-TIR")
model_path = snapshot_download(repo_id="Makima57/deepseek-math-Numina")
generator = ctranslate2.Generator(model_path, device="cpu", compute_type="int8")
iterations = 4
test=True

# Function to generate predictions using the model
def get_prediction(question):
    if test==True:
        text="Solve the following mathematical problem: what is  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\nLet's implement this in Python using the sympy library.\n\n```python\nimport sympy as sp\n\n# Define the variable\nx = sp.symbols('x')\n\n# Define the polynomials\npoly1 = 2*x + 3\npoly2 = 3*x\n\n# Sum the polynomials\nsum_poly = poly1 + poly2\n\n# Simplify the resulting polynomial\nsimplified_sum_poly = sp.simplify(sum_poly)\n\n# Print the simplified polynomial\nprint(simplified_sum_poly)\n```\n```output\n5*x + 3\n```\nThe sum of the polynomials \\(2x + 3\\) and \\(3x\\) is \\(\\boxed{5x + 3}\\).\n"  

        return text
    input_text = model_prompt + question
    input_tokens = tokenizer.tokenize(input_text)
    results = generator.generate_batch(
        [input_tokens],
        max_length=512,
        sampling_temperature=0.7,
        sampling_topk=40,
    )
    output_tokens = results[0].sequences[0]
    predicted_answer = tokenizer.convert_tokens_to_string(output_tokens)
    return predicted_answer

# 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:
        # Check for "Answer:" or "answer:"
        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:
        # If no "Answer:" found, assume last line is the answer
        answer = lines[-1].strip()
        steps = lines
    steps_text = '\n'.join(steps).strip()
    return answer, steps_text

# 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,sucess= postprocess_completion(prediction, return_status=True, last_code_block=True)
        print(answer,sucess)
        if sucess:
            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)

    majority_voted_ans = get_majority_vote(all_answers)
    if sucess:
            print(type_check(majority_voted_ans))
            if type_check(majority_voted_ans) == "Polynomial":
                plotfile = draw_polynomial_plot(majority_voted_ans) 
    else:
        if os.path.exists("polynomial_plot.png"):
            plotfile = "polynomial_plot.png"
        else:

            plotfile = None


    # Get the majority voted answer
    


    # 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

def gradio_interface(question, correct_answer):
    final_answer, steps_solution,plotfile = majority_vote_with_steps(question, iterations)
    print(plotfile)
    print(os.path.exists(plotfile),os.path.exists("polynomial_plot.png"))
    return question, final_answer, steps_solution, correct_answer,plotfile

# Custom CSS for enhanced design (unchanged)
custom_css = """
    body {
        background-color: #fafafa;
        font-family: 'Open Sans', sans-serif;
    }
    .gradio-container {
        background-color: #ffffff;
        border: 3px solid #007acc;
        border-radius: 15px;
        padding: 20px;
        box-shadow: 0 8px 20px rgba(0, 0, 0, 0.15);
        max-width: 800px;
        margin: 50px auto;
    }
    h1 {
        font-family: 'Poppins', sans-serif;
        color: #007acc;
        font-weight: bold;
        font-size: 32px;
        text-align: center;
        margin-bottom: 20px;
    }
    p {
        font-family: 'Roboto', sans-serif;
        font-size: 18px;
        color: #333;
        text-align: center;
        margin-bottom: 15px;
    }
    input, textarea {
        font-family: 'Montserrat', sans-serif;
        font-size: 16px;
        padding: 10px;
        border: 2px solid #007acc;
        border-radius: 10px;
        background-color: #f1f8ff;
        margin-bottom: 15px;
    }
    #math_question, #correct_answer {
        font-size: 20px;
        font-family: 'Poppins', sans-serif;
        font-weight: 500px;
        color: #007acc;
        margin-bottom: 5px;
        display: inline-block;
    }
    
    textarea {
        min-height: 150px;
    }
    .gr-button-primary {
        background-color: #007acc !important;
        color: white !important;
        border-radius: 10px !important;
        font-size: 18px !important;
        font-weight: bold !important;
        padding: 10px 20px !important;
        font-family: 'Montserrat', sans-serif !important;
        transition: background-color 0.3s ease !important;
    }
    .gr-button-primary:hover {
        background-color: #005f99 !important;
    }
    .gr-button-secondary {
        background-color: #f44336 !important;
        color: white !important;
        border-radius: 10px !important;
        font-size: 18px !important;
        font-weight: bold !important;
        padding: 10px 20px !important;
        font-family: 'Montserrat', sans-serif !important;
        transition: background-color 0.3s ease !important;
    }
    .gr-button-secondary:hover {
        background-color: #c62828 !important;
    }
    .gr-output {
        background-color: #e0f7fa;
        border: 2px solid #007acc;
        border-radius: 10px;
        padding: 15px;
        font-size: 16px;
        font-family: 'Roboto', sans-serif;
        font-weight: bold;
        color: #00796b;
    }
"""

# Define the directory path
flagging_dir = "./flagged_data" 

# Create the directory if it doesn't exist
if not os.path.exists(flagging_dir):
    os.makedirs(flagging_dir)

# Gradio app setup with flagging
interface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(label="🧠 Math Question", placeholder="Enter your math question here...", elem_id="math_question"),
    ],
    outputs=[
        gr.Textbox(label="Question", interactive=False),  # Non-editable
        gr.Textbox(label="Answer", interactive=False),  # Non-editable
        gr.Textbox(label="Solution", interactive=True),  # Editable textbox for correct solution
        gr.Image(label="Polynomial Plot", type="filepath")
    ],
    title="🔢 Math Question Solver",
    description="Enter a math question to get the model's majority-voted answer and steps to solve the problem.",
    css=custom_css,  # Apply custom CSS
    flagging_dir=flagging_dir,  # Directory to save flagged data
    allow_flagging="auto"  # Allow users to auto flag data
)

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