import warnings
import numpy as np
np.random.seed(42)
warnings.filterwarnings('ignore')
# We suppress warnings so outputs are cleaner.

# Data setup (Mall Customers Dataset)
import pandas as pd
url = 'https://gist.githubusercontent.com/pravalliyaram/5c05f43d2351249927b8a3f3cc3e5ecf/raw/Mall_Customers.csv'
df = pd.read_csv(url)
print(df.shape)
print(df.head(3))

# Select only relevant numerical features for clustering
X = df[['Age', 'Annual Income (k$)', 'Spending Score (1-100)']]
print(X.head())

# Standardize our features (very important for clustering!)
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# Fit KMeans clustering with k=3
from sklearn.cluster import KMeans
kmeans = KMeans(n_clusters=3, random_state=42)
labels = kmeans.fit_predict(X_scaled)

# Scatter plot to view clusters
import matplotlib.pyplot as plt
plt.figure(figsize=(6,4))
plt.scatter(X_scaled[:, 1], X_scaled[:, 2], c=labels, cmap='viridis', alpha=0.7)
plt.xlabel('Annual Income (scaled)')
plt.ylabel('Spending Score (scaled)')
plt.title('Cluster assignments')
plt.show()

# Inertia: sum of squared distances to nearest center (lower is better)
print('Inertia:', kmeans.inertia_)

# Silhouette score measures how similar a point is to its cluster vs others (-1=bad, 1=good)
from sklearn.metrics import silhouette_score
score = silhouette_score(X_scaled, labels)
print('Silhouette Score:', round(score,2))

# Davies-Bouldin Index: lower is better (measures overlap)
from sklearn.metrics import davies_bouldin_score
db_score = davies_bouldin_score(X_scaled, labels)
print('Davies-Bouldin Index:', round(db_score,2))

# Calinski-Harabasz Score: Higher is better
from sklearn.metrics import calinski_harabasz_score
ch_score = calinski_harabasz_score(X_scaled, labels)
print('Calinski-Harabasz Score:', int(ch_score))

# Try different k, compare silhouette scores
scores = []
for k in range(2,7):
    km = KMeans(n_clusters=k, random_state=42)
    labs = km.fit_predict(X_scaled)
    score = silhouette_score(X_scaled, labs)
    scores.append(round(score,2))
print('Silhouette scores for k=2 to k=6:', scores)

# See what our clusters look like by average values
df['Cluster'] = labels
print(df.groupby('Cluster')[['Age', 'Annual Income (k$)', 'Spending Score (1-100)']].mean().round(1))

# User challenge: Try clustering with your own choice of k using input()
k = int(input('Choose a value for k (2-6): '))
custom_km = KMeans(n_clusters=k, random_state=42)
custom_labels = custom_km.fit_predict(X_scaled)
print('Silhouette Score:', round(silhouette_score(X_scaled, custom_labels), 2))

# Pro tip: Use elbow method to find best k visually
inertias = []
K = range(1, 8)
for k in K:
    km = KMeans(n_clusters=k, random_state=42)
    km.fit(X_scaled)
    inertias.append(km.inertia_)
plt.plot(K, inertias, marker='o')
plt.xlabel('Number of clusters k')
plt.ylabel('Inertia')
plt.title('Elbow Method for optimal k')
plt.show()