# 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.datasets import make_blobs
from sklearn.metrics import adjusted_rand_score
from sklearn.mixture import GaussianMixture
from sklearn.model_selection import train_test_split

X, truth = make_blobs(n_samples=400, centers=3, cluster_std=0.8, random_state=44)
X = X @ np.array([[1.8, 0.7], [0.0, 0.5]])
train, test = train_test_split(np.arange(len(X)), random_state=44)
candidates = []
for k in [2, 3, 4]:
    model = GaussianMixture(n_components=k, covariance_type="full",
        reg_covar=1e-5, n_init=3, max_iter=150, random_state=44).fit(X[train])
    assert model.converged_
    candidates.append((model.bic(X[train]), model))
model = min(candidates, key=lambda item: item[0])[1]
labels = model.predict(X[test])
responsibility = model.predict_proba(X[test])
assert np.allclose(responsibility.sum(axis=1), 1)
assert np.isfinite(model.score_samples(X[test])).all()
print("Selected components:", model.n_components)
print("Held-out mean log density:", model.score(X[test]))
print("Synthetic ARI:", adjusted_rand_score(truth[test], labels))
samples, component = model.sample(5)
assert samples.shape == (5, 2) and component.shape == (5,)