# Lesson 6: clustering - group shoppers without any labels
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import KMeans

df = pd.read_csv("shoppers.csv")
X = StandardScaler().fit_transform(df[["yearly_spend", "visits"]])   # put both on the same scale

print("k  inertia (lower = tighter groups)")
for k in range(2, 7):
    print(k, round(KMeans(n_clusters=k, n_init=10, random_state=0).fit(X).inertia_, 1))

km = KMeans(n_clusters=3, n_init=10, random_state=0).fit(X)
df["group"] = km.labels_
summary = df.groupby("group").agg(shoppers=("shopper_id", "count"), avg_spend=("yearly_spend", "mean"), avg_visits=("visits", "mean")).round(0)
print("\n" + summary.sort_values("avg_spend").to_string())
df.to_csv("shoppers_grouped.csv", index=False)
print("\nSaved shoppers_grouped.csv")
