import gradio as gr
import ctranslate2
from transformers import AutoTokenizer
from huggingface_hub import snapshot_download
from codeexecutor import get_majority_vote
import re

# 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 = 10

# Function to generate predictions using the model
def get_prediction(question):
    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('
')
    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[:-1]
    steps_text = '
'.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, steps = parse_prediction(prediction)
        all_predictions.append(prediction)
        all_answers.append(answer)
        steps_list.append(steps)

    # Get the majority voted answer
    majority_voted_ans = get_majority_vote(all_answers)

    # 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]
            break
    else:
        steps_solution = "No steps found"

    return majority_voted_ans, steps_solution

# Gradio interface for user input and output
def gradio_interface(question, correct_answer):
    final_answer, steps_solution = majority_vote_with_steps(question, iterations)
    return {
        "Question": question,
        "Majority-Voted Answer": final_answer,
        "Steps to Solve": steps_solution,
        "Correct Solution": correct_answer
    }

# Custom CSS for enhanced design (unchanged)


# Gradio app setup
interface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(label="🧠 Math Question", placeholder="Enter your math question here...", elem_id="math_question"),
        gr.Textbox(label="✅ Correct Answer", placeholder="Enter the correct answer here...", elem_id="correct_answer"),
    ],
    outputs=[
        gr.JSON(label="📊 Results"),  # Display the results in a JSON format
    ],
    title="🔢 Math Question Solver",
    description="Enter a math question to get the model's majority-voted answer and steps to solve the problem.",
  
)

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