from pylab import array, scatter, figure, show, mean, argmin, append
from random import random, seed
from sys import argv, exit

def init(X, k):
    """Simply use k random datapoints as initial means for the clusters."""
    return X[[int(X.shape[0] * random()) for i in range(k)]]

def init_pp(X, k):
    """Use the k-means++ algorithm to find initial means."""
    # Choose first mean at random
    N = X.shape[0]
    indices = [int(N * random())]
    m = [X[indices[0]]]
    # Initial distances
    D = [((x - m[0])**2).sum() for x in X]
    while len(m) < k:
        # Find new best mean
        best_sum = new = -1
        for i in range(N):
            if i not in indices:
                Dsum = sum([min(D[j], ((X[j] - X[i])**2).sum()) for j in range(N)])
                if best_sum == -1 or best_sum < Dsum:
                    best_sum, new = Dsum, i
        # Add new mean and update distances
        m.append(X[new])
        indices.append(new)
        D = [min(D[i], ((X[i] - X[new])**2).sum()) for i in range(N)]
    return array(m)

# Parse parameters (600 is a nice example seed)
pp = False
if len(argv) > 2:
    if argv[2] == 'pp':
        print 'Using k-means++'
        initial_means = init_pp
        pp = True
    else:
        seed(int(argv[2]))
    if len(argv) == 4:
        seed(int(argv[3]))
elif len(argv) < 2:
    print 'Usage: python %s K [ "pp" ] [ SEED ]' % argv[0]
    exit()
if not pp:
    print 'Using normal k-means'
    initial_means = init
k = int(argv[1])
if not 1 <= k <= 6:
    print 'K must be a value from 1-6 (we only defined six colors).'
    exit()

# Generate dataset, add a multiplication of k so that clusters are formed
n, N = 2, 200
X = array([[100 * random() + 70 for j in range(n)] for i in \
        range(int(N / k + N % k))])
for c in range(k - 1):
    d = (c + 2) * 70
    X = append(X, [[100 * random() + d for j in range(n)] for i in \
        range(int(N / k))], 0)

# Divide in clusters by applying k-means
M = initial_means(X, k)
Mp = M - 1
steps = 0

while (Mp - M).any():
    Mp = M
    clusters = [[] for i in range(k)]
    # Assignment step
    for x in X:
        clusters[argmin([((x - m)**2).sum() for m in M])].append(x)
    # Update step
    M = array([mean(c, 0) for c in clusters])
    steps += 1

print 'Completed in %d steps' % steps

# Plot clusters
figure(1)
colors = [[1,0,0], [0,1,0], [0,0,1], [0,0,0], [1,1,1], [.5,.5,.5]]
for i in range(k):
    c = array(clusters[i])
    scatter(c[:,0], c[:,1], c=colors[i])
show()
