from sklearn.cluster import KMeans
import numpy as np
import matplotlib.pyplot as plt
from img_util import *
from sklearn.metrics.cluster import adjusted_rand_score
import copy
from sklearn.mixture import GaussianMixture
from itertools import product



def run_mspp_seg(img, width, height, count, data, target):
	new = copy.deepcopy(data)
	current = copy.deepcopy(data)

	t = 0
	h = 30
	threshold = 1
	v = [-1, 0, 1]

	offsets = list(product(v, repeat=3))

	while (t == 0) or np.linalg.norm(new - current, axis=1).sum() >= threshold: 
		current = copy.deepcopy(new)
		maxs = current.max(axis=0)
		shapes = np.ceil((maxs)/h).astype(int) + 1

		C = np.zeros(shapes)
		S = np.zeros(shapes.tolist() + [current.shape[1]]) 
		
		cell_idx = np.floor(current/h).astype(int)
		
		for pt, idx in zip(current, cell_idx):
			C[tuple(idx)] += 1
			S[tuple(idx)] += pt


		sum_vec = np.zeros_like(current)
		count = np.zeros(len(data))
		idx = np.floor(current/h).astype(int)


		for offset in offsets:
			neighbor_idx = idx + offset
			neighbor_idx[neighbor_idx<0] = 0
			
			neighbor_idx[:,0][neighbor_idx[:,0]>=shapes[0]] = shapes[0]-1
			neighbor_idx[:,1][neighbor_idx[:,1]>=shapes[1]] = shapes[1]-1
			neighbor_idx[:,2][neighbor_idx[:,2]>=shapes[2]] = shapes[2]-1
			
			temp = list(map(tuple, neighbor_idx.T))
			sum_vec += S[temp[0],temp[1],temp[2]]
			count += C[temp[0],temp[1],temp[2]]
		count[count == 0] = 1
		new = sum_vec / count.reshape(-1,1)

		t+=1
		if t>=10:
			break

	print("Termination at value:", np.linalg.norm(new - current, axis=1).sum(), ' at iter: ', t)



	new_img = new
	new_img = np.reshape(new_img,(height,width,3)).astype(int)
	new_img[new_img>255] = 255
	new_img[new_img<0] = 0

	#show_img(img)
	#show_img(new_img)

	h = 30
	cells = np.floor(new/h).astype(int)
	uniques, preds = np.unique(cells,axis=0,return_inverse=True)
	print("Num cells: ",len(uniques))


	return np.reshape(preds, (height,width)), preds

















