import gradio as gr
import google.generativeai as genai
from google.generativeai.types import HarmCategory, HarmBlockThreshold
import sympy as sp

# Set up the Gemini generative model using your API key
model = genai.GenerativeModel("gemini-1.5-flash", api_key="AIzaSyCjAahw7YvVFWSYHWRaGwZ9cA9emobfNok")

# Define the functions Gemini will call
def algebra_simplification(expression):
    import sympy as sp
    try:
        expr = sp.sympify(expression)
        simplified_expr = sp.simplify(expr)
        return str(simplified_expr)
    except Exception as e:
        return f"Error: {str(e)}"

def algebra_solve(equation):
    import sympy as sp
    try:
        eq = sp.sympify(equation)
        variables = list(eq.free_symbols)
        solutions = sp.solve(eq, *variables)
        return str(solutions)
    except Exception as e:
        return f"Error: {str(e)}"

# Chat model and function calling logic
def gemini_chat(math_problem):
    chat = model.start_chat(history=[])
    
    # Gemini model should figure out whether to call the appropriate function
    response = chat.send_message(math_problem, functions=[
        {
            "name": "algebra_simplification",
            "parameters": {
                "expression": math_problem
            },
            "fn": algebra_simplification
        },
        {
            "name": "algebra_solve",
            "parameters": {
                "equation": math_problem
            },
            "fn": algebra_solve
        }
    ], function_call="auto")  # Set function_call to 'auto'
    
    return response.text

# Gradio Interface
def process_math_problem(math_problem):
    # Use Gemini's auto function calling to process the math problem
    return gemini_chat(math_problem)

with gr.Blocks() as demo:
    gr.Markdown("# Gemini Algebra Solver")
    
    input_problem = gr.Textbox(label="Enter a Math Problem", placeholder="e.g., Simplify (x^2 - 4)/(x - 2) or Solve x^2 - 4 = 0")
    output_solution = gr.Textbox(label="Solution")
    
    solve_button = gr.Button("Solve")
    solve_button.click(fn=process_math_problem, inputs=input_problem, outputs=output_solution)

demo.launch()