import numpy as np
from sklearn.cluster import MeanShift
from sklearn.datasets.samples_generator import make_blobs
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import style

style.use('ggplot')

dummy_centers = [[1, 1, 1], [5, 5, 5], [3, 10, 10]]
X, _ = make_blobs(n_samples=100, centers=dummy_centers, cluster_std=1.5)

mean_shift = MeanShift()
mean_shift.fit(X)

labels = mean_shift.labels_
cluster_centers = mean_shift.cluster_centers_
number_clusters_ = len(np.unique(labels))

colors = 10 * ['r', 'g', 'b', 'c', 'k', 'y', 'm']

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')

for i in range(len(X)):
    ax.scatter(X[i][0], X[i][1], X[i][2], color=colors[labels[i]], marker='o')

ax.scatter(cluster_centers[:, 0], cluster_centers[:, 1], cluster_centers[:, 2], marker='x', color='k',
           s=150, linewidths=5, zorder=10)

plt.show()

print(cluster_centers)
print('Number of estimated clusters: ', number_clusters_)