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()
    columns_info = "\n".join([f"{col}: {df[col].dtype}" for col in df.columns])
    prompt = f"""
    Given this dataset:
    Columns and types:
    {columns_info}
    
    Data summary:
    {data_description}
    
    User instructions: {instructions if instructions else 'No specific instructions provided.'}
    
    Suggest 3 ways to visualize this data. For each visualization:
    1. Describe the visualization type and what it will show.
    2. Provide Python code using matplotlib to create the visualization.
    3. Explain why this visualization is useful for understanding the data.
    
    Format your response as:
    Visualization 1:
    Description: ...
    Code: ...
    Explanation: ...
    
    Visualization 2:
    ...
    
    Visualization 3:
    ...
    """

    response = model.generate_content(prompt)
    suggestions = response.text.split("Visualization")

    visualizations = []
    for i, suggestion in enumerate(suggestions[1:4], 1):  # Process only the first 3 visualizations
        parts = suggestion.split("Code:")
        description = parts[0].strip()
        code = parts[1].split("Explanation:")[0].strip()
        
        # Execute the code
        plt.figure(figsize=(10, 6))
        exec(code)
        plt.title(f"Visualization {i}")
        
        # Save the plot to a BytesIO object
        buf = io.BytesIO()
        plt.savefig(buf, format='png')
        buf.seek(0)
        img_str = base64.b64encode(buf.getvalue()).decode()
        plt.close()
        
        visualizations.append((f"data:image/png;base64,{img_str}", description, code))

    return visualizations

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Data Visualization with Gemini")
    api_key = gr.Textbox(label="Enter Gemini API Key", type="password")
    file = gr.File(label="Upload Excel or CSV file")
    instructions = gr.Textbox(label="Optional visualization instructions")
    submit = gr.Button("Generate Visualizations")
    
    with gr.Row():
        output1 = gr.Image(label="Visualization 1")
        output2 = gr.Image(label="Visualization 2")
        output3 = gr.Image(label="Visualization 3")
    
    with gr.Row():
        desc1 = gr.Textbox(label="Description 1")
        desc2 = gr.Textbox(label="Description 2")
        desc3 = gr.Textbox(label="Description 3")
    
    with gr.Row():
        code1 = gr.Code(language="python", label="Code 1")
        code2 = gr.Code(language="python", label="Code 2")
        code3 = gr.Code(language="python", label="Code 3")

    submit.click(
        fn=process_file,
        inputs=[api_key, file, instructions],
        outputs=[
            output1, desc1, code1,
            output2, desc2, code2,
            output3, desc3, code3
        ],
        show_progress=True,
    )

demo.launch()