import gradio as gr
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, confusion_matrix

# --- Step 1: Load dataset ---
data = load_iris()
X, y = data.data, data.target
feature_names = data.feature_names
target_names = data.target_names

# --- Step 2: Split dataset ---
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# --- Step 3: Train model ---
model = LogisticRegression(max_iter=200)
model.fit(X_train, y_train)

# --- Step 4: Evaluate model ---
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
conf_mat = confusion_matrix(y_test, y_pred)

# --- Step 5: Graph 1 - Confusion Matrix ---
def plot_confusion_matrix():
    fig, ax = plt.subplots(figsize=(4, 3))
    sns.heatmap(conf_mat, annot=True, fmt="d", cmap="Blues",
                xticklabels=target_names, yticklabels=target_names)
    plt.xlabel("Predicted")
    plt.ylabel("Actual")
    plt.title("Confusion Matrix")
    fig.canvas.draw()
    img = np.array(fig.canvas.renderer.buffer_rgba())
    plt.close(fig)
    return img

# --- Step 6: Graph 2 - Feature Importance ---
def plot_feature_importance():
    importance = np.abs(model.coef_).mean(axis=0)
    fig, ax = plt.subplots(figsize=(5, 3))
    sns.barplot(x=importance, y=feature_names, palette="viridis")
    plt.title("Feature Importance")
    fig.canvas.draw()
    img = np.array(fig.canvas.renderer.buffer_rgba())
    plt.close(fig)
    return img

# --- Step 7: Prediction Function ---
def predict_flower(sepal_length, sepal_width, petal_length, petal_width):
    features = [[sepal_length, sepal_width, petal_length, petal_width]]
    prediction = model.predict(features)[0]
    class_name = target_names[prediction]
    return f"🌸 Predicted Flower: {class_name}"

# --- Step 8: Gradio Interface ---
with gr.Blocks() as demo:
    gr.Markdown("## 🌼 Iris Flower Classifier with Graphs")
    gr.Markdown(f"**Model Accuracy:** {accuracy:.2f}")

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

    output_text = gr.Textbox(label="Prediction Result")

    predict_btn = gr.Button("🔍 Predict")
    predict_btn.click(
        predict_flower,
        inputs=[sepal_length, sepal_width, petal_length, petal_width],
        outputs=output_text
    )

    gr.Markdown("### 📊 Model Performance Graphs")
    gr.Image(value=plot_confusion_matrix(), label="Confusion Matrix")
    gr.Image(value=plot_feature_importance(), label="Feature Importance")

# --- Step 9: Launch the app ---
demo.launch()