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(df.shape)

# 2. FEATURES, MISSING VALUES AND DISTRIBUTIONS

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

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

df[features].hist()
plt.show()

df = df.dropna(subset=features)

X = df[features]

# 3. STANDARDIZATION

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

# 4. 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()

# 5. TRAIN K-MEANS

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

# 6. PLOT CLUSTERS AND CENTROIDS

centroids = scaler.inverse_transform(kmeans.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)

plt.xlabel("Annual Income")
plt.ylabel("Spending Score")
plt.title("Customer Segments")
plt.show()

# 7. DESCRIBE CUSTOMER SEGMENTS

print("\nCluster Summary:")
print(df.groupby("Cluster")[features].mean().round(2))

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

for i, (income, spending) in enumerate(centroids):
    income_type = "High" if income >= X.iloc[:, 0].mean() else "Low"
    spending_type = "High" if spending >= X.iloc[:, 1].mean() else "Low"
    print("Cluster", i, ":", income_type, "Income,", spending_type, "Spending")
