# Let us silence warnings for a cleaner notebook
import warnings
warnings.filterwarnings('ignore')

# Let us import some useful tools for our lesson
import pandas as pd
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
%matplotlib inline

# Load the mall customers data
data_url = "https://raw.githubusercontent.com/rithwik00/Mall-Customer-Dataset/main/Mall_Customers.csv"
df = pd.read_csv(data_url)
df.columns = df.columns.str.strip()

# Show the shape and a quick peek at the data
print('Shape of data:', df.shape)
df.head()

# Let us check if our data has missing values
df.isnull().sum()

# Select the columns we will use for clustering
features = df[["Annual Income (k$)", "Spending Score (1-100)"]]
features.head()

# Plot the customers by income and spending
plt.figure(figsize=(8,6))
plt.scatter(features["Annual Income (k$)"], features["Spending Score (1-100)"])
plt.xlabel("Annual Income (k$)")
plt.ylabel("Spending Score (1-100)")
plt.title("Customers by Income and Spending Score")
plt.show()

# Try clustering with 3 groups first
kmeans = KMeans(n_clusters=3, random_state=42)
kmeans.fit(features)
labels = kmeans.labels_
print(labels[:10])

# Change the number of clusters using input()
k = int(input("How many clusters would you like to try? "))
kmeans_alt = KMeans(n_clusters=k, random_state=42)
kmeans_alt.fit(features)
df["Cluster"] = kmeans_alt.labels_
df[["Annual Income (k$)", "Spending Score (1-100)", "Cluster"]].head()

# Plot the clusters with different colors
plt.figure(figsize=(8,6))
for i in range(k):
    plt.scatter(
        features[kmeans_alt.labels_ == i]["Annual Income (k$)"],
        features[kmeans_alt.labels_ == i]["Spending Score (1-100)"],
        label=f"Cluster {i+1}",
        alpha=0.7
    )
plt.xlabel("Annual Income (k$)")
plt.ylabel("Spending Score (1-100)")
plt.title(f"KMeans with {k} Clusters")
plt.legend()
plt.show()

# Check where each cluster center is
centers = kmeans_alt.cluster_centers_
print("Cluster Centers:")
print(centers)

# What happens if we pick way too many clusters?
k_too_many = 12
kmeans_many = KMeans(n_clusters=k_too_many, random_state=42)
kmeans_many.fit(features)
plt.figure(figsize=(8,6))
for i in range(k_too_many):
    plt.scatter(
        features[kmeans_many.labels_ == i]["Annual Income (k$)"],
        features[kmeans_many.labels_ == i]["Spending Score (1-100)"],
        label=f"Cluster {i+1}",
        alpha=0.7
    )
plt.xlabel("Annual Income (k$)")
plt.ylabel("Spending Score (1-100)")
plt.title("KMeans with Too Many Clusters")
plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')
plt.show()

# How to choose the best number of clusters? The Elbow method.
inertias = []
for k_try in range(1, 11):
    model = KMeans(n_clusters=k_try, random_state=42)
    model.fit(features)
    inertias.append(model.inertia_)
plt.figure(figsize=(6,4))
plt.plot(range(1, 11), inertias, marker='o')
plt.xlabel("Number of Clusters")
plt.ylabel("Inertia (sum of distances)")
plt.title("Elbow Method for Choosing k")
plt.show()

# Add clusters to our table, and look at each group
best_k = 5
final_kmeans = KMeans(n_clusters=best_k, random_state=42)
df["Cluster"] = final_kmeans.fit_predict(features)
df.groupby("Cluster").mean(numeric_only=True)

# Add some new data and predict which cluster it would join
new_customer = [[80, 50]]
cluster_for_new = final_kmeans.predict(new_customer)
print(f"Cluster for income 80k$ and spending score 50 is: {cluster_for_new[0]}")

# Remove the Cluster column if we are finished analyzing
df = df.drop(columns=["Cluster"])
df.head()

# How does KMeans handle out-of-range values?
wild_customer = [[1000, 0]]
wild_group = final_kmeans.predict(wild_customer)
print(f"Extreme customer goes to cluster: {wild_group[0]}")

# Extra tip: Resetting your notebook
df = pd.read_csv(data_url)
df.columns = df.columns.str.strip()
features = df[["Annual Income (k$)", "Spending Score (1-100)"]]

# Challenge: Try clustering by age and income instead
features_challenge = df[["Age", "Annual Income (k$)"]]
challenge_kmeans = KMeans(n_clusters=4, random_state=42)
challenge_kmeans.fit(features_challenge)
labels_challenge = challenge_kmeans.labels_
plt.figure(figsize=(8,6))
for i in range(4):
    plt.scatter(
        features_challenge[labels_challenge == i]["Age"],
        features_challenge[labels_challenge == i]["Annual Income (k$)"],
        label=f"Cluster {i+1}",
        alpha=0.7
    )
plt.xlabel("Age")
plt.ylabel("Annual Income (k$)")
plt.title("Clusters by Age and Income")
plt.legend()
plt.show()