import numpy as np
import math
import random
from scipy import misc

IMAGE = 'ball.jpg'

def pix_variance(x, y):
    a = x[0] - y[0]
    b = x[1] - y[1]
    c = x[2] - y[2]
    d = [a, b, c]
    return (d)

def cluster_points(X, mu):
    clusters  = {} #сюда с наименьшим расстоянием
    for x in X:
        bestmukey = min([(i[0], np.linalg.norm(pix_variance(x[0], mu[i[0]][0]))) \
                    for i in enumerate(mu)], key=lambda t:t[1])[0]
        try:
            clusters[bestmukey].append(x)
        except KeyError:
            clusters[bestmukey] = [x]
    #print("clusters =", clusters)
    print()
    print()
    return clusters


def reevaluate_centers(mu, clusters):
    #все короче перекрашиваем
    newmusplited = []
    keys = sorted(clusters.keys())
    for k in keys:
        newmu = []
        clusters_in_one = []
        for j in clusters[k]:
            clusters_in_one.append(np.array([j[0][0], j[0][1], j[0][2], j[1][0], j[1][1]]))
        newmu.append(np.mean(clusters_in_one, axis=0))
        newmusplited.append([[int(newmu[0][0]), int(newmu[0][1]), int(newmu[0][2])], [int(newmu[0][3]), int(newmu[0][4])]])
    return newmusplited

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


def find_centers(X, K):
    oldmu = random.sample(X, K) #можно задать и специально
    mu = random.sample(X, K)
    print("oldmu :=", oldmu)
    print("mu :=", mu)
    #oldmu = pixels[80*240 - 50]
    #mu = pixels[80*240 - 100]
    while not has_converged(mu, oldmu):
        oldmu = mu
        print(mu)
        print("*")
        clusters = cluster_points(X, mu)
        mu = reevaluate_centers(oldmu, clusters)
    return(mu, clusters)


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



def main():
    for countK in range(2, 3): #два цвета
        image = misc.imread(IMAGE)
        dim_x, dim_y, colors = image.shape
        print("Loaded image of shape {x}, {y}".format(x=dim_x, y=dim_y))
        pixels = []

        for i in range(dim_x):
            for j in range(dim_y):
                pixels.append([[image[i, j, 0], image[i, j, 1], image[i, j, 2]], [i, j]])

        res = find_centers(pixels, countK)
        print(res[0][0][0])

        for k in range(0, len(res[0])):
            for xy in res[1][k]:
                i = xy[1][0]
                j = xy[1][1]
                image[i, j, 0] = res[0][k][0][0]
                image[i, j, 1] = res[0][k][0][1]
                image[i, j, 2] = res[0][k][0][2]

        print("-------------------------")
        print("Done for " + str(countK))
        print("-------------------------")
        misc.imsave(('outBall' + str(countK) + '.png'), image)


if __name__ == '__main__':
    main()