import numpy
import pylab
import vlfeat
from vlfeat.plotop.vl_plotframe import vl_plotframe
#from vlfeat.kmeans.vl_hikmeanshist import vl_hikmeanshist


def color_gen():
	colors = ['r', 'g', 'y']
	for i in range(1000):
		yield colors[i%len(colors)]
	 

def xplot(X, tree):
	x1 = []
	x2 = []

	centers = tree.get_centers()
	C = centers.shape[1]
	 
	for k in range(centers.shape[1]):
		x1.append([X[0], centers[0, k]])
		x2.append([X[1], centers[1, k]])
	
	if tree.sub.size() > 0:		
		i = 0
		for sub in tree.sub:
			[x1p, x2p] = xplot(tree.get_centers()[:, i], sub)
			x1 += x1p
			x2 += x2p
			i += 1
	return [x1, x2]


def plottree(tree, linewidth=2):
	""" PLOTTRE  Plot hierarchical K-means tree
	PLOTTREE(TREE) plots a tree generated by HIKMEASN(). 
	See also VL_HIKMEANS().
	"""
	
	x1 = []
	x2 = []
	i = 0
	
	for sub in tree.sub:	
		[x1p, x2p] = xplot(tree.get_centers()[:, i], sub)
	  	x1 += x1p
	  	x2 += x2p
	  	i += 1
	
	for i in range(len(x1)):
		pylab.plot(x1[i], x2[i], '+-b', linewidth=linewidth)
	
	
def vl_test_hikmeans():
	""" VL_TEST_HIKMEANS Test VL_HIKMEANS function
	"""
	K = 3
	nleaves = 100
	data = numpy.array(numpy.random.rand(2, 1000) * 255, 'uint8')
	datat = numpy.array(numpy.random.rand(2, 10000) * 255, 'uint8')
	
	[tree, A] = vlfeat.vl_hikmeans(data, K, nleaves, verb=1)
	AT = vlfeat.vl_hikmeanspush(tree, datat)
		
	pylab.figure()
	plottree(tree)
	pylab.xlim(0, 255)
	pylab.ylim(0, 255)
	print('hikmeans-tree') ;
	
	pylab.figure()
	gen = color_gen()
	for k in range(K*K):
		color = gen.next()
		sel = pylab.find(A[-1, :] == k)
		pylab.plot(data[0, sel], data[1, sel], '.', color=color)
		sel = pylab.find(AT[-1, :] == k)
		pylab.plot(datat[0, sel], datat[1, sel], '+', color=color)
	
	plottree(tree, linewidth=4)
	pylab.xlim(0, 255)
	pylab.ylim(0, 255)
	print('hikmeans-clusters') ;
	

def vl_test_hikmeans_io():	
	""" VL_TEST_HIKMEANS_IO Test VL_HIKMEANS function
	"""
	K = 3
	nleaves = 100
	data = numpy.array(numpy.random.rand(2, 1000) * 255, 'uint8')
	
	[tree, A] = vlfeat.vl_hikmeans(data, K, nleaves, verb=1)
	tree.save('/tmp/test.vlhkm')
	
	tree2 = vlfeat._vlfeat.VlHIKMTree(0, 0)
	tree2.load('/tmp/test.vlhkm')
	
	pylab.figure()
	plottree(tree)
	pylab.xlim(0, 255)
	pylab.ylim(0, 255)
	
	pylab.figure()
	plottree(tree2)
	pylab.xlim(0, 255)
	pylab.ylim(0, 255)
	print('hikmeans-tree-io')


if __name__ == '__main__':
	vl_test_hikmeans()
	vl_test_hikmeans_io()
	pylab.show()

	
	
	