import gradio as gr
import pandas as pd
import numpy as np
import plotly.graph_objects as go
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.neural_network import MLPRegressor

css = '''
.gradio-container{max-width: 700px !important}
h1{text-align:center}
'''
def analyze_data(csv_file):

    df = pd.read_csv(csv_file.name)
    
    if len(df.columns) < 3:
        return "Error: CSV file must have at least 3 columns for analysis."
    

    x, y, z = df.iloc[:, 0], df.iloc[:, 1], df.iloc[:, 2]
    

    original_fig = go.Figure(data=[go.Scatter3d(
        x=x, y=y, z=z, mode='markers',
        marker=dict(size=5, color=z, colorscale='Viridis', opacity=0.8)
    )])
    original_fig.update_layout(title="Original Data: 3,000 points", scene=dict(
        xaxis_title=df.columns[0], yaxis_title=df.columns[1], zaxis_title=df.columns[2]))
    

    X = df.iloc[:, :2]
    y = df.iloc[:, 2]
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)
    
    xx, yy = np.meshgrid(np.linspace(X.iloc[:, 0].min(), X.iloc[:, 0].max(), 100),
                         np.linspace(X.iloc[:, 1].min(), X.iloc[:, 1].max(), 100))
    grid = np.c_[xx.ravel(), yy.ravel()]
    grid_scaled = scaler.transform(grid)
    
    models = [
        ("Good model", MLPRegressor(hidden_layer_sizes=(10, 5), max_iter=1000, random_state=42)),
        ("Underfitting", MLPRegressor(hidden_layer_sizes=(2,), max_iter=10, random_state=42)),
        ("Overfitting", MLPRegressor(hidden_layer_sizes=(100, 100), max_iter=1000, random_state=42)),
        ("4-layer NN", MLPRegressor(hidden_layer_sizes=(50, 30, 20), max_iter=100000, random_state=42))
    ]
    
    figs = []
    for name, model in models:
        model.fit(X_train_scaled, y_train)
        z_pred = model.predict(grid_scaled).reshape(xx.shape)
        
        fig = go.Figure(data=[go.Surface(x=xx, y=yy, z=z_pred, colorscale='Viridis')])
        fig.update_layout(title=f"{name}", scene=dict(
            xaxis_title=df.columns[0], yaxis_title=df.columns[1], zaxis_title=df.columns[2]))
        figs.append(fig)
    
    return original_fig, figs[0], figs[1], figs[2], figs[3]

iface = gr.Interface(
    fn=analyze_data,
    inputs=gr.File(label="Upload CSV file"),
    outputs=[
        gr.Plot(label="Original Data"),
        gr.Plot(label="1. Good model"),
        gr.Plot(label="2. Underfitting"),
        gr.Plot(label="3. Overfitting"),
        gr.Plot(label="4. 4-layer Neural Network")
    ],
    
    title="PLOTS REG 3D",
    description="Upload a CSV file with at least 3 columns to create various 3D graph visualizations and model predictions. Higher the size of CSV, Higher the Time Taken",
    css=css,
    theme="bethecloud/storj_theme",
)

iface.launch()