import numpy as np
from scipy.cluster.vq import kmeans, whiten
from operator import itemgetter
from math import ceil
from mpi4py import MPI
import time

startTime = time.time()

comm = MPI.COMM_WORLD
rank = comm.Get_rank(); size = comm.Get_size()

np.random.seed(seed=rank) # XXX should use parallel RNG
obs = whiten(np.genfromtxt('alteredData.csv', dtype=float, delimiter=','))

K = 10
nstart = 100
n = int(ceil(float(nstart) / size))
centroids, distortion = kmeans(obs, K, n)
results = comm.gather((centroids, distortion), root=0)

if rank == 0:
    results.sort(key=itemgetter(1))
    result = results[0]
    print('Best distortion for %d tries: %f' % (nstart, result[1]))
    print('Result Array:', result)
    print('TIme it took:', time.time() - startTime, 'seconds')