"""
手写kmeans
"""
import numpy as np
from matplotlib import pyplot as plt
from sklearn.datasets import make_blobs
import typing


class KMeans:
    def __init__(self, k: int):
        self.k = k
        self._centers = None

    def fit(self, nda: np.ndarray, n_iters=10, callback: typing.Callable = None):
        n_features = nda.shape[1]
        centers = self.random_centers(self.k, n_features)

        for i in range(n_iters):
            labels = self.assign(centers, nda)
            centers = self.update(nda, labels, self.k)
            if callback:
                callback(nda, labels, i)

    def predict(self, nda: np.ndarray):
        return self.assign(self.centers, nda)

    @property
    def centers(self):
        if not self._centers:
            raise AttributeError("Call 'fit' before reference to centers.")
        return self._centers

    @staticmethod
    def random_centers(k, n_features):
        return np.random.random((k, n_features))

    @staticmethod
    def assign(centers, nda):
        n = nda.shape[0]
        labels = np.empty(n)
        for i, arr in enumerate(nda):
            labels[i] = KMeans.nearest_center(centers, arr)
        return labels

    @staticmethod
    def update(nda, labels, k):
        centers = np.empty(k)
        for i in range(k):
            center = KMeans.cal_center(nda, labels, i)
            centers[i] = center
        return centers

    @staticmethod
    def distance(arr1, arr2):
        return np.sum((arr1 - arr2) ** 2)

    @staticmethod
    def cal_center(nda, labels, i):
        return np.mean(nda[labels == i])

    @staticmethod
    def nearest_center(centers, nda):
        j = -1
        min_dis = np.PINF
        for i, center in enumerate(centers):
            dis = KMeans.distance(center, nda)
            if dis < min_dis:
                min_dis = dis
                j = i
        return j


def my_plot(nda, labels, i):
    if i % 10 == 0:
        plt.scatter(nda[:, 0], nda[:, 1], c=labels)
        plt.title("i = %s" % i)
        plt.savefig("%s.png" % i)


def main():
    X, y = make_blobs(n_samples=1000, n_features=2, centers=[[-1, -1], [0, 0], [1, 1], [2, 2]],
                      cluster_std=[0.4, 0.2, 0.2, 0.2],
                      random_state=9)

    kmeans = KMeans(4)
    kmeans.fit(X, n_iters=40, callback=my_plot)


if __name__ == '__main__':
    main()

#在console中输出图片
%matplotlib inline

#在单独窗口中输出图片
%matplotlib qt5