import gradio as gr
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.metrics import (
    accuracy_score,
    mean_absolute_error,
    mean_squared_error,
    r2_score,
)
import numpy as np

# ======================
# Model imports
# ======================
from sklearn.linear_model import (
    LinearRegression,
    LogisticRegression,
    Perceptron,
)
from sklearn.neighbors import (
    KNeighborsClassifier,
    KNeighborsRegressor,
)
from sklearn.naive_bayes import GaussianNB
from sklearn.tree import (
    DecisionTreeClassifier,
    DecisionTreeRegressor,
)
from sklearn.svm import SVC, SVR
from sklearn.neural_network import (
    MLPClassifier,
    MLPRegressor,
)

# ======================
# Model registry
# ======================
REGRESSION_MODELS = {
    "Linear Regression": LinearRegression,
    "KNN Regressor": KNeighborsRegressor,
    "Decision Tree Regressor": DecisionTreeRegressor,
    "SVR": SVR,
    "MLP Regressor": MLPRegressor,
}

CLASSIFICATION_MODELS = {
    "Logistic Regression": LogisticRegression,
    "KNN Classifier": KNeighborsClassifier,
    "Naive Bayes": GaussianNB,
    "Perceptron": Perceptron,
    "Decision Tree Classifier": DecisionTreeClassifier,
    "SVM Classifier": SVC,
    "MLP Classifier": MLPClassifier,
}

# ======================
# UI Logic
# ======================
def update_models(task_type):
    if task_type == "Regression":
        return gr.update(choices=list(REGRESSION_MODELS.keys()), value=None)
    return gr.update(choices=list(CLASSIFICATION_MODELS.keys()), value=None)


def train_model(file, task_type, model_name):
    df = pd.read_csv(file.name)

    # Assumption: last column is target
    X = df.iloc[:, :-1]
    y = df.iloc[:, -1]

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42
    )

    if task_type == "Regression":
        model = REGRESSION_MODELS[model_name]()
        model.fit(X_train, y_train)

        preds = model.predict(X_test)

        mae = mean_absolute_error(y_test, preds)
        mse = mean_squared_error(y_test, preds)
        rmse = np.sqrt(mse)
        r2 = r2_score(y_test, preds)

        return (
            f"Model: {model_name}\n"
            f"Task: Regression\n"
            f"MAE: {mae:.4f}\n"
            f"MSE: {mse:.4f}\n"
            f"RMSE: {rmse:.4f}\n"
            f"R² Score: {r2:.4f}"
        )

    else:
        model = CLASSIFICATION_MODELS[model_name]()
        model.fit(X_train, y_train)

        preds = model.predict(X_test)
        acc = accuracy_score(y_test, preds)

        return (
            f"Model: {model_name}\n"
            f"Task: Classification\n"
            f"Accuracy: {acc:.4f}"
        )


# ======================
# Gradio App
# ======================
with gr.Blocks() as demo:
    gr.Markdown("## Supervised Learning Model Trainer")
    gr.Markdown(
        "Upload a CSV file. The **last column is treated as target**."
    )

    file_input = gr.File(label="Upload CSV", file_types=[".csv"])
    task_dropdown = gr.Dropdown(
        ["Regression", "Classification"],
        label="Task Type",
    )
    model_dropdown = gr.Dropdown(label="Model")
    output = gr.Textbox(label="Result", lines=5)

    train_btn = gr.Button("Generate")

    task_dropdown.change(
        update_models,
        inputs=task_dropdown,
        outputs=model_dropdown,
    )

    train_btn.click(
        train_model,
        inputs=[file_input, task_dropdown, model_dropdown],
        outputs=output,
    )

demo.launch()