import numpy as np


def kmeans_plusplus_init(X, k, rng):
    n = X.shape[0]
    centroids = [X[rng.integers(n)]]
    for _ in range(1, k):
        dists = np.min(
            [np.sum((X - c) ** 2, axis=1) for c in centroids], axis=0
        )
        probs = dists / dists.sum()
        idx = rng.choice(n, p=probs)
        centroids.append(X[idx])
    return np.array(centroids)


def kmeans(X, k, max_iters=100, seed=0):
    rng = np.random.default_rng(seed)
    centroids = kmeans_plusplus_init(X, k, rng)
    for _ in range(max_iters):
        dists = np.linalg.norm(
            X[:, None, :] - centroids[None, :, :], axis=2
        )
        assignments = np.argmin(dists, axis=1)
        new_centroids = np.array([
            X[assignments == j].mean(axis=0)
            if np.any(assignments == j) else centroids[j]
            for j in range(k)
        ])
        if np.allclose(new_centroids, centroids):
            break
        centroids = new_centroids
    return centroids, assignments


def inertia(X, centroids, assignments):
    return sum(
        np.sum((X[assignments == j] - centroids[j]) ** 2)
        for j in range(len(centroids))
    )


def silhouette_score(X, assignments):
    n = len(X)
    k = len(np.unique(assignments))
    scores = np.zeros(n)
    dist_matrix = np.linalg.norm(
        X[:, None, :] - X[None, :, :], axis=2
    )
    for i in range(n):
        own = assignments[i]
        same_cluster = (assignments == own) & (np.arange(n) != i)
        a_i = (
            dist_matrix[i, same_cluster].mean()
            if same_cluster.sum() > 0 else 0.0
        )
        b_i = min(
            dist_matrix[i, assignments == other].mean()
            for other in range(k) if other != own
            and np.any(assignments == other)
        )
        scores[i] = (b_i - a_i) / max(a_i, b_i) if max(a_i, b_i) > 0 else 0.0
    return scores.mean()
