import pandas as pd
import matplotlib.pyplot as plt

from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans

# 1. LOAD DATASET

df = pd.read_csv("Mall_Customers.csv")

print(df.head())
print("Columns:", df.columns)

# 2. SELECT FEATURES

features = ["Annual Income (k$)", "Spending Score (1-100)"]

print("Missing Values:\n", df[features].isnull().sum())

df = df.dropna(subset=features).copy()
X = df[features]

# 3. STANDARDIZATION

scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

# ELBOW METHOD

inertia = []

for k in range(1, 11):
    model = KMeans(n_clusters=k, random_state=42, n_init=10)
    model.fit(X_scaled)
    inertia.append(model.inertia_)

plt.plot(range(1, 11), inertia, marker="o")
plt.xlabel("Number of Clusters (K)")
plt.ylabel("Inertia")
plt.title("Elbow Method")
plt.show()

k = int(input("Enter suitable K from elbow graph: "))

# 4. TRAIN K-MEANS

model = KMeans(n_clusters=k, random_state=42, n_init=10)
df["Cluster"] = model.fit_predict(X_scaled)

# 5 & 7. VISUALIZE CLUSTERS AND CENTROIDS

centroids = scaler.inverse_transform(model.cluster_centers_)

plt.scatter(X.iloc[:, 0], X.iloc[:, 1],
            c=df["Cluster"], cmap="viridis")

plt.scatter(centroids[:, 0], centroids[:, 1],
            color="red", marker="X", s=150, label="Centroids")

plt.xlabel("Annual Income")

plt.ylabel("Spending Score")
plt.title("Mall Customer Clusters")
plt.legend()
plt.show()

# 6. NUMBER OF CUSTOMERS IN EACH CLUSTER

print("\nCustomers per Cluster:")
print(df["Cluster"].value_counts().sort_index())


# ENHANCED: Income distribution; original elbow and centroid plots retained
X["Annual Income (k$)"].hist(bins=15)
plt.xlabel("Annual income (k$)")
plt.title("Customer Income Distribution")
plt.show()
