# Source: content/notes/ml/unsupervised-learning.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import numpy as np
from sklearn.cluster import KMeans, MiniBatchKMeans
from sklearn.datasets import make_blobs
from sklearn.metrics import adjusted_rand_score, silhouette_score
from sklearn.model_selection import train_test_split
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

X, truth = make_blobs(n_samples=450, centers=3, cluster_std=0.75, random_state=11)
train, test = train_test_split(np.arange(len(X)), random_state=11)
model = make_pipeline(StandardScaler(), KMeans(n_clusters=3, n_init=10, random_state=11))
model.fit(X[train])
labels = model.predict(X[test])
scaled_test = model.named_steps["standardscaler"].transform(X[test])
centers = model.named_steps["kmeans"].cluster_centers_
distances = ((scaled_test[:, None, :] - centers[None, :, :]) ** 2).sum(axis=2)
assert np.array_equal(labels, distances.argmin(axis=1))
assert len(np.unique(labels)) == 3
print("Held-out silhouette:", silhouette_score(scaled_test, labels))
print("Synthetic external ARI:", adjusted_rand_score(truth[test], labels))
mini = make_pipeline(StandardScaler(), MiniBatchKMeans(
    n_clusters=3, n_init=5, batch_size=64, random_state=11))
mini.fit(X[train])
mini_labels = mini.predict(X[test])
assert mini_labels.shape == labels.shape
print("MiniBatch external ARI:", adjusted_rand_score(truth[test], mini_labels))