import numpy as np
import random
import matplotlib.pyplot as plt

def init_board(N):
	random.seed(100)
	X = np.array([(random.uniform(-1, 1), random.uniform(-1, 1)) for i in range(N)])
	random.seed(None)
	return X

def cluster_points(X, mu):
    clusters  = {}
    for x in X:
        bestmukey = min([(i[0], np.linalg.norm(x-mu[i[0]])) \
                    for i in enumerate(mu)], key=lambda t:t[1])[0]
        try:
            clusters[bestmukey].append(x)
        except KeyError:
            clusters[bestmukey] = [x]
    return clusters

def reevaluate_centers(mu, clusters):
    newmu = []
    keys = sorted(clusters.keys())
    for k in keys:
        newmu.append(np.mean(clusters[k], axis = 0))
    return newmu

def has_converged(mu, oldmu):
    return (set([tuple(a) for a in mu]) == set([tuple(a) for a in oldmu]))

def find_centers(X, K):
    # Initialize to K random centers
    random.seed(None)
    oldmu = random.sample(X, K)
    mu = random.sample(X, K)
    while not has_converged(mu, oldmu):
        oldmu = mu
        # Assign all points in X to clusters
        clusters = cluster_points(X, mu)
        # Reevaluate centers
        mu = reevaluate_centers(oldmu, clusters)
    return(mu, clusters)

def change_coords(array):
    return list(map(list, zip(*array)))

def parse_output(data):
    clusters = data[1]
    points1 = change_coords(clusters[0])
    plt.plot(points1[0], points1[1], 'ro')
    points2 = change_coords(clusters[1])
    plt.plot(points2[0], points2[1], 'g^')
    points3 = change_coords(clusters[2])
    plt.plot(points3[0], points3[1], 'ys')
    centroids = change_coords(data[0])
    plt.plot(centroids[0], centroids[1], 'kx')
    plt.axis([-1.0, 1, -1.0, 1])
    plt.show()

data = init_board(15)
# data = np.array([[0.560, 0.256], [-0.560, 0.206], [0.012, -0.789], [0.011, 0.785], [0.126, 0.692], 
# 				[-0.211, 0.785], [-0.011, 0.585], [-0.911, -0.785], [-0.811, -0.678], [-0.7511, -0.9], 
# 				[-0.891, -0.555], [0.896, -0.675], [0.786, -0.775], [0.681, -0.599], [0.596, -0.888]])
print(data)
print(type(data))
out = find_centers(list(data), 3)
parse_output(out)
