import gradio as gr
import matplotlib.pyplot as plt
import sys
import io
import numpy as np
from PIL import Image

def run_code(code):
    try:
        old_stdout = sys.stdout
        sys.stdout = io.StringIO()

        local_env = {'plt': plt, 'np': np}
        exec(code, {}, local_env)

        output = sys.stdout.getvalue()
        sys.stdout = old_stdout
        
        if 'plt' in code:
            buf = io.BytesIO()
            plt.draw()
            plt.savefig(buf, format='png')
            buf.seek(0)
            img = Image.open(buf)
            img_array = np.array(img)
            plt.close()

            if img_array.sum() > 0:
                return img_array
            else:
                return create_image("Empty plot")

        return create_image("No valid matplotlib plot found")

    except Exception as e:
        return create_image(f"Error: {str(e)}")

def create_image(message):
    fig, ax = plt.subplots()
    ax.text(0.5, 0.5, message, fontsize=20, ha='center', va='center')
    ax.axis('off')
    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)
    img = Image.open(buf)
    plt.close()
    return np.array(img)

interface = gr.Interface(
    fn=run_code, 
    inputs=gr.Code(language="python", label="Enter Python Code"), 
    outputs=gr.Image(type="numpy", label="Output"), 
    title="Matplotlib Plot Generator",
    description="Enter Python code that uses matplotlib to generate plots."
)

interface.launch()