import gradio as gr
from huggingface_hub import InferenceClient

# Initialize the Hugging Face Inference Client
client = InferenceClient()

# Function to optimize code or simplify mathematical logic
def optimize_or_simplify(input_text, task_type):
    """
    Optimizes a given code snippet or simplifies a mathematical expression based on the selected task.

    Args:
        input_text (str): The user-provided code or math expression.
        task_type (str): The type of task ('Code Optimization' or 'Math Simplification').

    Returns:
        str: The optimized code or simplified mathematical result.
    """
    if task_type == "Code Optimization":
        prompt = (
            f"Optimize the following code for performance and readability. "
            f"Provide detailed suggestions and refactor the code:\n\n{input_text}"
        )
    else:  # Math Simplification
        prompt = (
            f"Simplify the following mathematical expression or algorithm, "
            f"ensuring accuracy and efficiency:\n\n{input_text}"
        )
    
    # Call the Hugging Face model
    response = client.chat.completions.create(
        model="Qwen/Qwen2.5-Coder-32B-Instruct",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.5,
        max_tokens=512
    )
    return response["choices"][0]["message"]["content"]

# Create Gradio interface
with gr.Blocks() as app:
    gr.Markdown("## Code Optimization and Math Simplification Assistant")
    gr.Markdown(
        "Refine your code for better performance or simplify complex mathematical expressions effortlessly. "
        "Choose a task and input your text to get started!"
    )

    with gr.Row():
        # Input section
        with gr.Column():
            task_type = gr.Radio(
                choices=["Code Optimization", "Math Simplification"],
                label="Select Task",
                value="Code Optimization"
            )
            input_text = gr.Textbox(
                lines=10,
                label="Input Text",
                placeholder="Enter your code or mathematical expression here"
            )
            process_button = gr.Button("Process")

        # Output section
        with gr.Column():
            gr.Markdown("### Output")
            output_result = gr.Textbox(lines=15, interactive=False)

    # Link button to function
    process_button.click(
        fn=optimize_or_simplify,
        inputs=[input_text, task_type],
        outputs=output_result
    )

# Launch the app
app.launch()