import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist
from scipy.cluster.hierarchy import dendrogram, linkage
from scipy.cluster.vq import vq, whiten
from scipy.stats import zscore

def euclidean_distance(a, b):
    return np.sqrt(np.sum((a-b)**2))

def get_clusters(data, k):
    dist_matrix = cdist(data, data, metric=euclidean_distance)
    cluster_labels = np.zeros(shape=len(data))

    for i in range(k):
        closest_clusters = np.argmin(dist_matrix, axis=1)
        closest_clusters_data = data[closest_clusters == i]
        center = np.mean(closest_clusters_data, axis=0)

        for j, data_point in enumerate(data):
            if closest_clusters[j] == i:
                dist_matrix[j, :] = euclidean_distance(data_point, center)
            else:
                dist_matrix[j, :] = dist_matrix[j, i]

        cluster_labels[closest_clusters == i] = i

    return cluster_labels.astype(int)

# Synthetic dataset
np.random.seed(0)
n_points_per_cluster = 50

C1 = np.random.multivariate_normal([15, 15], [[5, 0], [0, 5]], n_points_per_cluster)
C2 = np.random.multivariate_normal([5, 25], [[5, 0], [0, 5]], n_points_per_cluster)
C3 = np.random.multivariate_normal([25, 15], [[5, 0], [0, 5]], n_points_per_cluster)

data = np.vstack((C1, C2, C3))
z = zscore(data)

k = 3
cluster_labels = get_clusters(z, k)

# Visualize clusters
fig, ax = plt.subplots()
for i in range(k):
    ax.scatter(data[cluster_labels == i, 0], data[cluster_labels == i, 1], label=f'Cluster {i+1}')
ax.legend()
plt.show()