import random
import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr
import spaces  # Import the spaces module for ZeroGPU compatibility

# Ensure the environment has access to a CUDA-capable GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Load model and tokenizer directly to GPU if available
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained("microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True)

print("Loading model...")
model = AutoModelForCausalLM.from_pretrained("microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True, device_map="auto").to(device)

# Define templates for problems
templates = {
    "algebra": {
        "easy": ["Solve for x: {a}x + {b} = {c}", "Find the value of x: {a}x - {b} = {c}"],
        "medium": ["Solve for x: {a}x^2 + {b}x + {c} = 0", "Find the roots of: {a}x^2 - {b}x = {c}"],
        "hard": ["Solve for x: {a}x^3 + {b}x^2 + {c}x + {d} = 0", "Find the value of x in the equation: {a}x^3 - {b}x^2 + {c}x = {d}"]
    },
    "calculus": {
        "easy": ["Differentiate the function: f(x) = {a}x^2 + {b}x + {c}", "Find the derivative of: f(x) = {a}x^3 - {b}x + {c}"],
        "medium": ["Integrate the function: f(x) = {a}x^2 + {b}x + {c}", "Find the integral of: f(x) = {a}x^3 - {b}x + {c}"],
        "hard": ["Solve the differential equation: {a}dy/dx + {b}y = {c}", "Find the solution to the differential equation: {a}d^2y/dx^2 - {b}dy/dx + {c}y = 0"]
    },
    "geometry": {
        "easy": ["Find the area of a rectangle with length {a} and width {b}", "Calculate the perimeter of a rectangle with length {a} and width {b}"],
        "medium": ["Find the area of a triangle with base {a} and height {b}", "Calculate the circumference of a circle with radius {a}"],
        "hard": ["Find the volume of a cylinder with radius {a} and height {b}", "Calculate the surface area of a sphere with radius {a}"]
    },
    "trigonometry": {
        "easy": ["Find sin({a})", "Find cos({a})"],
        "medium": ["Calculate tan({a})", "Find the value of sin({a}) + cos({a})"],
        "hard": ["Solve for θ in the equation sin(θ) = {a}", "Find the angle θ for which tan(θ) = {a}"]
    }
}

def generate_synthetic_math_problems(num_problems, selected_templates):
    problems = []
    for _ in range(num_problems):
        # Randomly choose an area of mathematics from the selected templates
        area = random.choice(selected_templates)
        
        # Randomly choose a difficulty level
        difficulty = random.choice(list(templates[area].keys()))
        
        # Randomly choose a template
        template = random.choice(templates[area][difficulty])
        
        # Randomly generate parameters
        a = random.randint(1, 10)
        b = random.randint(1, 10)
        c = random.randint(1, 10)
        d = random.randint(1, 10)
        
        # Generate the problem using the template and parameters
        problem = template.format(a=a, b=b, c=c, d=d)
        problems.append(problem)
    
    return problems

@spaces.GPU(duration=60)
def solve_problem(problem, max_length):
    print(f"Solving problem: {problem}")
    with torch.no_grad():
        # Encode the problem
        inputs = tokenizer(problem, return_tensors="pt").to(device)
        
        # Generate a response from the model
        outputs = model.generate(inputs["input_ids"], max_length=max_length, num_return_sequences=1, do_sample=True)
        
        # Decode the response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Strip the answer to only the math (assuming answer is preceded by "The answer is ")
        if "The answer is " in response:
            answer = response.split("The answer is ")[-1].strip()
        else:
            answer = response.strip()
    
    print(f"Problem: {problem}, Answer: {answer}")
    return answer

@spaces.GPU(duration=120)
def generate_and_solve_problems(num_problems, max_length, selected_templates):
    problems = generate_synthetic_math_problems(num_problems, selected_templates)
    solved_problems = []

    for problem in problems:
        answer = solve_problem(problem, max_length)
        solved_problems.append({
            "problem": problem,
            "answer": answer
        })

    return solved_problems

def gradio_interface(num_problems, max_length, algebra, calculus, geometry, trigonometry):
    print(f"Generating and solving {num_problems} problems with max length {max_length}...")
    selected_templates = []
    if algebra:
        selected_templates.append("algebra")
    if calculus:
        selected_templates.append("calculus")
    if geometry:
        selected_templates.append("geometry")
    if trigonometry:
        selected_templates.append("trigonometry")
    
    if not selected_templates:
        return "Please select at least one math area."
    
    solved_problems = generate_and_solve_problems(num_problems, max_length, selected_templates)
    return json.dumps(solved_problems, indent=4)

# Create a Gradio interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Number(label="Number of Problems", value=10, precision=0),
        gr.Slider(label="Max Output Length", minimum=10, maximum=200, value=50),
        gr.Checkbox(label="Algebra", value=True),
        gr.Checkbox(label="Calculus", value=True),
        gr.Checkbox(label="Geometry", value=True),
        gr.Checkbox(label="Trigonometry", value=True)
    ],
    outputs=gr.Textbox(label="Generated and Solved Problems"),
    title="Synthetic Math Problem Generator and Solver",
    description="Generate and solve synthetic math problems using a HuggingFace model."
)

iface.launch()