import numpy as np

# ---------------------------------------------------------------
# Part A: Competitive learning — reproduces the worked example
# ---------------------------------------------------------------
neurons = np.array([[0.2, 0.6], [0.8, 0.4], [0.5, 0.9]])  # N1, N2, N3
x = np.array([0.1, 0.5])
eta = 0.5

dists = np.linalg.norm(neurons - x, axis=1)
winner = np.argmin(dists)
print("Distances:", np.round(dists, 3), " Winner: N", winner + 1)

neurons[winner] = neurons[winner] + eta * (x - neurons[winner])
print("Updated winner weights:", np.round(neurons[winner], 3))
print("Other neurons unchanged:", np.round(neurons, 3))

# ---------------------------------------------------------------
# Part C: K-means from scratch — small synthetic dataset, k=2
# ---------------------------------------------------------------
def kmeans(X, centroids, n_iter=3):
    centroids = centroids.copy()
    for it in range(n_iter):
        dists = np.linalg.norm(X[:, None, :] - centroids[None, :, :], axis=2)
        assign = np.argmin(dists, axis=1)
        new_centroids = np.array([
            X[assign == k].mean(axis=0) if np.any(assign == k) else centroids[k]
            for k in range(len(centroids))
        ])
        print(f"Iter {it+1}: assignments={assign.tolist()}  "
              f"centroids={np.round(new_centroids, 3).tolist()}")
        if np.allclose(new_centroids, centroids):
            print("Converged.")
            centroids = new_centroids
            break
        centroids = new_centroids
    return centroids, assign

X = np.array([[1, 1], [1.5, 2], [3, 1], [5, 7], [3.5, 5], [4.5, 5]])
init_centroids = np.array([[1, 1], [5, 7]])
final_centroids, final_assign = kmeans(X, init_centroids, n_iter=3)
print("Final centroids:", np.round(final_centroids, 3))

J = sum(np.linalg.norm(X[i] - final_centroids[final_assign[i]]) ** 2 for i in range(len(X)))
print("Within-cluster sum of squares J =", round(J, 3))