import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import tensorflow as tf
num_puntos = 2000
conjunto_puntos = []
for i in range(num_puntos):
    if np.random.random() < 0.5:
        conjunto_puntos.append([np.random.normal(0.0,  0.9),
                                       np.random.normal(0.0, 0.9)])
    else:
        conjunto_puntos.append([np.random.normal(3.0, 0.5), 
                                       np.random.normal(1, 0.5)])
    
    
vectors = tf.constant(conjunto_puntos)
k = 5
centroides = tf.Variable(tf.slice(tf.random_shuffle(vectors),[0,0],[k,-1]))
expanded_vectors = tf.expand_dims(vectors, 0)
expanded_centroides = tf.expand_dims(centroides, 1)
diff = tf.subtract(expanded_vectors, expanded_centroides)
sqr = tf.square(diff)
distances = tf.reduce_sum(sqr, 2)
assignments = tf.argmin(distances, 0)

means = tf.concat([tf.reduce_mean(tf.gather(vectors,tf.reshape(tf.where(tf.equal(assignments,c)),[1,-1])),1) for c in range(k)],0)

update_centroides = tf.assign(centroides, means)
init_op = tf.global_variables_initializer()

sess = tf.Session()
sess.run(init_op)

num_steps=100

for step in range(num_steps):
    _, centroid_values, assignment_values,_ = sess.run([update_centroides, centroides, assignments, means])
    
    data={"x": [], "y": [], "cluster": []}
    
    for i in range(len(assignment_values)):
        data["x"].append(conjunto_puntos[i][0])
        data["y"].append(conjunto_puntos[i][1])
        data["cluster"].append(assignment_values[i])
        
    df=pd.DataFrame(data)
    sns.lmplot("x","y",data=df,fit_reg=False, size=6, hue="cluster", legend=False)
    plt.show()