import gradio as gr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from io import StringIO
import sys
import contextlib
from sklearn import datasets, metrics, model_selection, preprocessing
import warnings
warnings.filterwarnings('ignore')

# Capture output helper
@contextlib.contextmanager
def capture_output():
    new_out, new_err = StringIO(), StringIO()
    old_out, old_err = sys.stdout, sys.stderr
    try:
        sys.stdout, sys.stderr = new_out, new_err
        yield sys.stdout, sys.stderr
    finally:
        sys.stdout, sys.stderr = old_out, old_err

def execute_code(code: str):
    """
    Execute the provided Python code and return the output
    """
    # Initialize output components
    output_text = ""
    output_plot = None
    output_df = None
    
    # Create a namespace for code execution
    namespace = {
        'np': np,
        'pd': pd,
        'plt': plt,
        'sns': sns,
        'datasets': datasets,
        'metrics': metrics,
        'model_selection': model_selection,
        'preprocessing': preprocessing
    }
    
    try:
        # Capture print outputs
        with capture_output() as (out, err):
            # Execute the code
            exec(code, namespace)
            
            # Capture print statements
            output_text = out.getvalue()
            if err.getvalue():
                output_text += "\nErrors:\n" + err.getvalue()
        
        # Check if there's a plot and convert to image
        if plt.gcf().axes:
            # Save plot to a temporary file
            import tempfile
            import os
            
            # Create a temporary file with .png extension
            temp_dir = tempfile.gettempdir()
            temp_file = os.path.join(temp_dir, f"plot_{os.getpid()}.png")
            
            # Save the current figure
            plt.savefig(temp_file, bbox_inches='tight', dpi=300)
            plt.close()
            
            # Set the output_plot to the file path
            output_plot = temp_file
        
        # Check for DataFrame output
        for var in namespace:
            if isinstance(namespace[var], pd.DataFrame):
                output_df = namespace[var]
                break
                
    except Exception as e:
        output_text = f"Error: {str(e)}"
    
    return output_text, output_plot, output_df

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("""
    # Python Data Science Code Executor
    Execute Python code with access to common data science libraries:
    - NumPy (np)
    - Pandas (pd) 
    - Matplotlib (plt)
    - Seaborn (sns)
    - Scikit-learn (datasets, metrics, model_selection, preprocessing)
    """)
    
    with gr.Row():
        with gr.Column():
            code_input = gr.Code(
                label="Python Code",
                language="python",
                value="""# Example: Load and visualize iris dataset
from sklearn.datasets import load_iris
iris = load_iris()
df = pd.DataFrame(iris.data, columns=iris.feature_names)
df['target'] = iris.target

# Create a scatter plot
plt.figure(figsize=(10, 6))
plt.scatter(df['sepal length (cm)'], df['sepal width (cm)'], c=df['target'])
plt.xlabel('Sepal Length (cm)')
plt.ylabel('Sepal Width (cm)')
plt.title('Iris Dataset - Sepal Length vs Width')

# Print first few rows
print("First 5 rows of the dataset:")
print(df.head())
"""
            )
            run_button = gr.Button("Execute Code", variant="primary")
        
        with gr.Column():
            output_text = gr.Textbox(label="Output", lines=5)
            output_plot = gr.Image(label="Plot Output")
            output_df = gr.Dataframe(label="DataFrame Output")
    
    # Handle code execution
    run_button.click(
        fn=execute_code,
        inputs=[code_input],
        outputs=[output_text, output_plot, output_df]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()