import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import silhouette_score

customers = pd.read_csv('mall_customers.csv')
print(f'Real customers in this dataset: {len(customers)}')
print(customers[['Age', 'Annual Income (k$)', 'Spending Score (1-100)']].describe())

plt.figure(figsize=(8, 6))
plt.scatter(customers['Annual Income (k$)'], customers['Spending Score (1-100)'], color='steelblue')
plt.title('Real Mall Customers: Income vs. Spending Score')
plt.xlabel('Annual Income (k$)')
plt.ylabel('Spending Score (1-100)')
plt.show()

X = customers[['Annual Income (k$)', 'Spending Score (1-100)']].values
rng = np.random.RandomState(42)
initial_idx = rng.choice(len(X), size=3, replace=False)
centroids = X[initial_idx].astype(float)
print(f'Real starting centroids, picked from 3 random real customers:')
print(centroids)

def assign_clusters(points, centers):
    distances = np.sqrt(((points[:, None, :] - centers[None, :, :]) ** 2).sum(axis=2))
    return distances.argmin(axis=1), distances

labels, distances = assign_clusters(X, centroids)
inertia_round_0 = sum(distances[i, labels[i]] ** 2 for i in range(len(X)))
print(f'Real total squared distance after the first assignment: {inertia_round_0:.2f}')

updated_centroids = np.array([X[labels == k].mean(axis=0) for k in range(3)])
print('Real updated centroids after averaging their assigned real customers:')
print(updated_centroids)
labels_2, distances_2 = assign_clusters(X, updated_centroids)
inertia_round_1 = sum(distances_2[i, labels_2[i]] ** 2 for i in range(len(X)))
print(f'Real total squared distance after one full real iteration: {inertia_round_1:.2f}')

inertias = []
for k in range(1, 11):
    model = KMeans(n_clusters=k, random_state=42, n_init=10).fit(X)
    inertias.append(model.inertia_)
    print(f'k={k}: real inertia={model.inertia_:.2f}')

plt.figure(figsize=(8, 5))
plt.plot(range(1, 11), inertias, marker='o', color='steelblue')
plt.title('Real Elbow Plot')
plt.xlabel('Number of Clusters (K)')
plt.ylabel('Inertia')
plt.show()

for k in range(2, 9):
    model = KMeans(n_clusters=k, random_state=42, n_init=10).fit(X)
    score = silhouette_score(X, model.labels_)
    print(f'k={k}: real silhouette score={score:.4f}')

final_model = KMeans(n_clusters=5, random_state=42, n_init=10).fit(X)
customers['cluster'] = final_model.labels_
print(customers['cluster'].value_counts().sort_index())

plt.figure(figsize=(8, 6))
plt.scatter(customers['Annual Income (k$)'], customers['Spending Score (1-100)'], c=customers['cluster'], cmap='viridis')
plt.scatter(final_model.cluster_centers_[:, 0], final_model.cluster_centers_[:, 1], c='red', marker='X', s=200, label='Centroids')
plt.title('Real K=5 Mall Customer Segments')
plt.xlabel('Annual Income (k$)')
plt.ylabel('Spending Score (1-100)')
plt.legend()
plt.savefig('kmeans_final_clusters.png', dpi=150)
plt.show()

cluster_summary = customers.groupby('cluster')[['Annual Income (k$)', 'Spending Score (1-100)']].mean()
cluster_summary['count'] = customers['cluster'].value_counts().sort_index()
print(cluster_summary)

for seed in [0, 1, 2, 3, 4]:
    single_run = KMeans(n_clusters=5, random_state=seed, n_init=1, init='random').fit(X)
    print(f'seed={seed}, single random init: real inertia={single_run.inertia_:.2f}')

X3 = customers[['Age', 'Annual Income (k$)', 'Spending Score (1-100)']].values
raw_model = KMeans(n_clusters=5, random_state=42, n_init=10).fit(X3)
raw_silhouette = silhouette_score(X3, raw_model.labels_)
print(f'Real unscaled 3-feature silhouette: {raw_silhouette:.4f}')
scaler = StandardScaler()
X3_scaled = scaler.fit_transform(X3)
scaled_model = KMeans(n_clusters=5, random_state=42, n_init=10).fit(X3_scaled)
scaled_silhouette = silhouette_score(X3_scaled, scaled_model.labels_)
print(f'Real scaled 3-feature silhouette: {scaled_silhouette:.4f}')

import joblib
joblib.dump(final_model, 'kmeans_model.joblib')
reloaded_kmeans = joblib.load('kmeans_model.joblib')
print(f'Real original model inertia: {final_model.inertia_:.4f}')
print(f'Real reloaded model inertia: {reloaded_kmeans.inertia_:.4f}')

customers[['CustomerID', 'Annual Income (k$)', 'Spending Score (1-100)', 'cluster']].to_csv('mall_customer_clusters.csv', index=False)
print('Real cluster assignments saved to mall_customer_clusters.csv')