import gradio as gr
import pandas as pd
import matplotlib.pyplot as plt
import io
import base64
import google.generativeai as genai

def process_file(api_key, file, instructions):
    # Set up Gemini API
    genai.configure(api_key=api_key)
    model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')

    # Read the file
    if file.name.endswith('.csv'):
        df = pd.read_csv(file.name)
    else:
        df = pd.read_excel(file.name)

    # Analyze data and get visualization suggestions from Gemini
    data_description = df.describe().to_string()
    prompt = f"Given this data: {data_description}\n"
    if instructions:
        prompt += f"And these instructions: {instructions}\n"
    prompt += "Suggest 3 ways to visualize this data."

    response = model.generate_content(prompt)
    suggestions = response.text.split('\n')

    # Generate visualizations
    visualizations = []
    for i, suggestion in enumerate(suggestions[:3]):
        plt.figure(figsize=(10, 6))
        plt.title(f"Visualization {i+1}")
        plt.text(0.5, 0.5, suggestion, ha='center', va='center', wrap=True)
        buf = io.BytesIO()
        plt.savefig(buf, format='png')
        buf.seek(0)
        img_str = base64.b64encode(buf.getvalue()).decode()
        visualizations.append(f"data:image/png;base64,{img_str}")
        plt.close()

    return visualizations

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("Data Visualization with Gemini")
    api_key = gr.Textbox(label="Enter Gemini API Key")
    file = gr.File(label="Upload Excel or CSV file")
    instructions = gr.Textbox(label="Optional visualization instructions")
    submit = gr.Button("Generate Visualizations")
    outputs = [gr.Image(label=f"Visualization {i+1}") for i in range(3)]
    
    submit.click(
        fn=process_file,
        inputs=[api_key, file, instructions],
        outputs=outputs
    )

demo.launch(share=True)