import streamlit as st
import pandas as pd
import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt

# App Title
st.title("Unsupervised Learning: K-Means Clustering")

# Sidebar Section: Tab for downloading a sample dataset
st.sidebar.subheader("Sample Dataset")
st.sidebar.write("Download a sample dataset to test the app. This sample contains two numerical features for demonstration purposes.")
sample_data = {
    "Feature1": [1.0, 1.5, 3.0, 5.0, 3.5, 4.5, 3.5],
    "Feature2": [1.0, 2.0, 4.0, 7.0, 5.0, 5.0, 4.5],
}
sample_df = pd.DataFrame(sample_data)
sample_csv = sample_df.to_csv(index=False)
st.sidebar.download_button(
    label="Download Sample CSV",
    data=sample_csv,
    file_name="sample_data.csv",
    mime="text/csv"
)

# Main Section: Upload dataset
st.header("Step 1: Upload Your Dataset")
st.write("Upload a CSV file containing your data. Ensure it includes numerical features for clustering.")
uploaded_file = st.file_uploader("Upload your dataset (CSV format)", type="csv")
if uploaded_file:
    data = pd.read_csv(uploaded_file)
    st.write("Preview of the uploaded data:")
    st.dataframe(data)

    # Step 2: Select features for clustering
    st.subheader("Step 2: Feature Selection")
    st.write("Select the numerical features you want to use for clustering.")
    selected_features = st.multiselect(
        "Select features for clustering:", data.columns.tolist()
    )

    if selected_features:
        X = data[selected_features]

        # Step 3: Configure clustering parameters
        st.subheader("Step 3: Clustering Configuration")
        st.write("Choose the number of clusters you want to create using the slider below.")
        n_clusters = st.slider("Select the number of clusters:", min_value=2, max_value=10, value=3)

        # Apply K-Means Clustering
        model = KMeans(n_clusters=n_clusters, random_state=42)
        cluster_labels = model.fit_predict(X)

        # Step 4: Add cluster labels to the dataset
        data['Cluster'] = cluster_labels
        st.write("Clustered Data:")
        st.dataframe(data)

        # Step 5: Visualize the clusters
        st.subheader("Step 5: Cluster Visualization")
        st.write("Visualize the clustering results. Select at least two features for plotting.")
        if len(selected_features) >= 2:
            fig, ax = plt.subplots()
            scatter = ax.scatter(
                X[selected_features[0]],
                X[selected_features[1]],
                c=cluster_labels,
                cmap="viridis",
                s=50
            )
            ax.set_xlabel(selected_features[0])
            ax.set_ylabel(selected_features[1])
            ax.set_title("K-Means Clustering")
            legend = ax.legend(*scatter.legend_elements(), title="Clusters")
            ax.add_artist(legend)
            st.pyplot(fig)
        else:
            st.warning("Select at least 2 features for visualization.")

        # Step 6: Download the clustered data
        st.subheader("Step 6: Download Clustered Data")
        st.write("Download the dataset with the cluster labels added.")
        csv = data.to_csv(index=False)
        st.download_button(
            label="Download CSV",
            data=csv,
            file_name="clustered_data.csv",
            mime="text/csv"
        )
    else:
        st.warning("Please select features for clustering.")
else:
    st.info("Awaiting file upload. Use the sample dataset in the sidebar if you don’t have a file.")