import gradio as gr
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from sklearn.preprocessing import StandardScaler

from sklearn.linear_model import LogisticRegression
from sklearn.svm import SVC
from sklearn.neighbors import KNeighborsClassifier
from sklearn.tree import DecisionTreeClassifier

import mlflow
import mlflow.sklearn

sns.set(style="whitegrid")

# ----------------------------
# Load Iris Data
# ----------------------------
iris = load_iris()
df = pd.DataFrame(iris.data, columns=iris.feature_names)
df["target"] = iris.target

# Scale
scaler = StandardScaler()
x_scaled = scaler.fit_transform(iris.data)
y = iris.target

# Split
x_train, x_test, y_train, y_test = train_test_split(
    x_scaled, y, test_size=0.2, random_state=42
)

# ----------------------------
# MLflow Setup
# ----------------------------
mlflow.set_experiment("iris-Experiment")

models = {
    "Logistic Regression": LogisticRegression(),
    "KNN (k=3)": KNeighborsClassifier(n_neighbors=3),
    "Decision Tree": DecisionTreeClassifier(),
    "Linear SVM": SVC(kernel="linear"),
}

trained_models = {}

def train_and_log(model, name):
    with mlflow.start_run(run_name=name):
        model.fit(x_train, y_train)
        y_pred = model.predict(x_test)
        acc = accuracy_score(y_test, y_pred)

        mlflow.log_param("model_name", name)
        mlflow.log_metric("accuracy", acc)
        mlflow.sklearn.log_model(model, name)

        trained_models[name] = model

for name, model in models.items():
    train_and_log(model, name)

# ----------------------------
# Visualization Functions
# ----------------------------
def create_plot(func_name):
    plt.figure(figsize=(7,5))

    if func_name == "Pairplot":
        sns.pairplot(df, hue="target")
        plt.savefig("pairplot.png")
        return "pairplot.png"

    elif func_name == "Correlation Heatmap":
        plt.figure(figsize=(8,6))
        sns.heatmap(df.iloc[:, :-1].corr(), annot=True, cmap="viridis")
        plt.title("Feature Correlation Heatmap")
        plt.savefig("heatmap.png")
        return "heatmap.png"

    elif func_name == "Boxplot":
        plt.figure(figsize=(8,6))
        df_melt = df.melt(id_vars="target", var_name="feature", value_name="value")
        sns.boxplot(x="feature", y="value", hue="target", data=df_melt)
        plt.xticks(rotation=45)
        plt.title("Boxplot of Features")
        plt.tight_layout()
        plt.savefig("boxplot.png")
        return "boxplot.png"

    elif func_name == "Violin Plot":
        plt.figure(figsize=(8,6))
        sns.violinplot(x="feature", y="value", hue="target", data=df_melt, split=True)
        plt.xticks(rotation=45)
        plt.title("Violin Plot of Features")
        plt.tight_layout()
        plt.savefig("violin.png")
        return "violin.png"

    elif func_name == "Countplot":
        sns.countplot(x=df["target"])
        plt.title("Class Distribution")
        plt.savefig("countplot.png")
        return "countplot.png"

    elif func_name == "Scatterplot (Petal Length vs Width)":
        sns.scatterplot(
            x=df["petal length (cm)"],
            y=df["petal width (cm)"],
            hue=df["target"],
            palette="deep",
            s=70
        )
        plt.title("Petal Length vs Petal Width")
        plt.savefig("scatter.png")
        return "scatter.png"


# ----------------------------
# Prediction Function
# ----------------------------
def predict(model_name, sepal_length, sepal_width, petal_length, petal_width):
    model = trained_models[model_name]
    x = np.array([[sepal_length, sepal_width, petal_length, petal_width]])
    x_scaled = scaler.transform(x)
    pred = model.predict(x_scaled)[0]
    return f"Predicted Class: {iris.target_names[pred]}"

# ----------------------------
# Build Gradio UI
# ----------------------------
with gr.Blocks() as demo:
    gr.Markdown("## 🌸 Iris Classification with MLflow + Rich Visualizations")

    with gr.Row():
        model_choice = gr.Dropdown(list(trained_models.keys()), label="Select ML Model")

    with gr.Row():
        sepal_length = gr.Number(label="Sepal Length")
        sepal_width = gr.Number(label="Sepal Width")
        petal_length = gr.Number(label="Petal Length")
        petal_width = gr.Number(label="Petal Width")

    predict_btn = gr.Button("Predict")
    output = gr.Textbox(label="Prediction")

    predict_btn.click(
        predict,
        inputs=[model_choice, sepal_length, sepal_width, petal_length, petal_width],
        outputs=output
    )

    gr.Markdown("### 📊 Choose Visualization")

    graph_choice = gr.Dropdown(
        [
            "Pairplot",
            "Correlation Heatmap",
            "Boxplot",
            "Violin Plot",
            "Countplot",
            "Scatterplot (Petal Length vs Width)",
        ],
        label="Select Graph"
    )

    plot_btn = gr.Button("Show Graph")
    graph_output = gr.Image()

    plot_btn.click(create_plot, inputs=graph_choice, outputs=graph_output)

demo.launch()