import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
from datasets import load_dataset
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import LabelEncoder

st.title("🌸 AI Flower Species Classification System")
st.write("Built with **K-Nearest Neighbors (KNN)** and trained using a dataset fetched live from **Hugging Face**.")

@st.cache_data
def load_hf_data():
    dataset = load_dataset("scikit-learn/iris", split="train")
    return pd.DataFrame(dataset)

with st.spinner("Fetching dataset from Hugging Face..."):
    df = load_hf_data()

st.success("Dataset successfully loaded from Hugging Face!")

# Dynamically select features and target
feature_columns = df.columns[:4]
target_column = df.columns[-1]

X = df[feature_columns]
y_text = df[target_column]

# Convert text labels (Iris-setosa, etc.) into numbers (0, 1, 2) for model and graph coloring
encoder = LabelEncoder()
y = encoder.fit_transform(y_text)

# Scale and Train KNN Model
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

knn = KNeighborsClassifier(n_neighbors=3)
knn.fit(X_scaled, y)

# Sidebar UI Controls for User Input
st.sidebar.header("🎛️ Input Flower Measurements")
inputs = []
for col in feature_columns:
    val = st.sidebar.slider(f"{col}", float(X[col].min()), float(X[col].max()), float(X[col].mean()))
    inputs.append(val)

# Prediction Button
if st.button("Predict Species"):
    user_input = scaler.transform([inputs])
    pred = knn.predict(user_input)
    
    predicted_name = encoder.inverse_transform(pred)[0]
    
    st.subheader("✨ Result:")
    st.success(f"The predicted flower species is: **{predicted_name}**")

# Graph Visualization
st.subheader("📊 Dataset Visualization & Your Input")
fig, ax = plt.subplots(figsize=(8, 5))
scatter = ax.scatter(X.iloc[:, 0], X.iloc[:, 1], c=y, cmap='viridis', s=60, edgecolors='k', label='Dataset Clusters')
ax.scatter(inputs[0], inputs[1], color='red', marker='X', s=250, label='Your Custom Input')
ax.set_xlabel(feature_columns[0])
ax.set_ylabel(feature_columns[1])
plt.colorbar(scatter, label='Species Code')
ax.legend()
st.pyplot(fig)