import numpy as np
import matplotlib.pyplot as plt

mean1 = [2, 2]
cov1 = [[0.5, 0], [0, 0.5]]
x1, y1 = np.random.multivariate_normal(mean1, cov1, 200).T

mean2 = [-2, -2]
cov2 = [[0.5, 0], [0, 0.5]]
x2, y2 = np.random.multivariate_normal(mean2, cov2, 200).T

X = np.concatenate((np.array([x1, y1]).T, np.array([x2, y2]).T))

k = 2
np.random.shuffle(X)
initial_centers = X[:k]

def calculate_distance(X, centers):
    distances = []
    for x in X:
        distances.append(np.linalg.norm(x - centers, axis=1))
    return np.array(distances)

distances = calculate_distance(X, initial_centers)

def assign_points(X, centers):
    distances = calculate_distance(X, centers)
    return np.argmin(distances, axis=1)

labels = assign_points(X, initial_centers)

def update_centers(X, labels, k):
    centers = []
    for i in range(k):
        indices = np.where(labels == i)
        centers.append(np.mean(X[indices], axis=0))
    return np.array(centers)

new_centers = update_centers(X, labels, k)

def kmeans(X, k, max_iter=100):
    np.random.shuffle(X)
    centers = X[:k]
    for _ in range(max_iter):
        labels = assign_points(X, centers)
        new_centers = update_centers(X, labels, k)
        if np.all(centers == new_centers):
            break
        centers = new_centers
    return labels, centers

labels, centers = kmeans(X, k)

plt.scatter(X[:, 0], X[:, 1], c=labels)
plt.scatter(centers[:, 0], centers[:, 1], c='red', marker='x')
plt.show()