import numpy as np
from sklearn.cluster import KMeans
from scipy.spatial import distance

class AlgorithmKMeans():
    # method: AlgorithmKNN::constructor
    #
    # arguments:
    #  win_input: GUI input display
    #  win_output: GUI output display
    #  win_log: GUI process log
    #
    # return: none
    #
    def __init__(self, win_input, win_output, win_log, n_clusters, init, n_init, maxiter):
        # create class data
        #
        AlgorithmKMeans.__CLASS_NAME__ = self.__class__.__name__

        # copy the inputs into class data
        #
        self.input_d = win_input
        self.output_d = win_output
        self.log_d = win_log
        self.n_clusters = n_clusters
        self.initial = init
        self.n_init = n_init
        self.maxiter = maxiter

    # method: AlgorithmKNN::initialize
    #
    # arguments: None
    #
    # return: True
    #
    # initialize variables for KNN
    #
    def initialize(self, data):
        # initialize variables
        #
        self.data = data
        self.classes = len(self.data)
        self.means = [d.mean(axis=0) for d in self.data]
        self.kmeans = KMeans(n_clusters=self.n_clusters,
                             init=self.initial, n_init=self.n_init,
                             max_iter=self.maxiter)
        self.X = np.empty((0, 0))
        return True


    def run_algo(self, data):
        self.initialize(data)
        self.train()
        return True


    def train(self):
        data = np.vstack((self.data))
        self.kmeans.fit(data)
        self.compute_cluster_mean(self.input_d)

    def compute_cluster_mean(self, ax):
        ax.canvas.axes.scatter(self.kmeans.cluster_centers_[:,0],
                               self.kmeans.cluster_centers_[:,1],
                               c='black', s=8)


    def predict(self, ax, X):
        X = np.concatenate(X, axis=0)
        X = np.reshape(X, (-1, 2))
        res = (ax.canvas.axes.get_xlim()[1] -
               ax.canvas.axes.get_ylim()[0]) / 100

        x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
        y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
        xx, yy = np.meshgrid(np.arange(x_min, x_max, res),
                             np.arange(y_min, y_max, res))
        Z = self.kmeans.predict(np.c_[xx.ravel(), yy.ravel()])
        Z = Z.reshape(xx.shape)

        return xx, yy, Z

    def prediction_classifier(self, data):
        centers = self.kmeans.cluster_centers_

        data = np.array([data])
        cluster = self.kmeans.predict(data)
        predicted_center = centers[cluster]
        prediction_list = []
        for means in self.means:
            dist = distance.euclidean(means, predicted_center)
            prediction_list.append(dist)
        prediction = prediction_list.index((min(prediction_list)))

        return prediction