import random
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


class K_means(object):  # 0继承

    def __init__(self, input_dir, k, column, outfilepath, col):  # path:csv文件路径  k：聚成的类别数
        # column：数据的列数（维数），和下面的 createDataSet（）方法 中的usecols=[1,2]中的列数是一样的

        self.list_data = self.createDataSet(input_dir, col)
        self.final_centroids, self.cluster = self.kmeans(self.list_data, k, input_dir, outfilepath)
        self.plot(self.final_centroids, self.cluster, k, column)

    def createDataSet(self, path, col):  # 将数据转成list形式，例如[[1,2,……]，[1,1,……],……]
        self.data_fd = open(path, encoding='utf-8')
        self.data_frame = pd.read_csv(self.data_fd, usecols=col, header=0)
        self.np_data = np.array(self.data_frame).astype(np.float64)
        self.list_data = self.np_data.tolist()
        return self.list_data

    # 得到每行数据到聚类中心的欧氏距离 ，并转成list形式，例如[[1 5]
    #                                                [2,6]
    #                                                [3,2]]  表示第一行数据到第一个中心的距离是1，到第二个中心的距离是5，第二行数据到第一个中心的距离是2，以此类推
    def calcDis(self, dataSet, centroids, k):  # dataSet：数据  centroids：聚类中心  k；类别数
        self.clalist = []
        for data in dataSet:
            diff = np.tile(data, (k, 1)) - centroids
            squaredDiff = diff ** 2
            squaredDist = np.sum(squaredDiff, axis=1)
            distance = squaredDist ** 0.5
            self.clalist.append(distance)
        self.clalist = np.array(self.clalist)
        return self.clalist

    # 得到新的聚类中心，新旧中心的变化量，如果变化量（是个array数组）有任何一个数不为0，就要重新执行这个方法，直到变化量为0
    def classify(self, dataSet, centroids, k):  # dataSet：数据   centroids：聚类中心  k：类别数
        self.calist = self.calcDis(dataSet, centroids, k)

        self.minDistIndices = np.argmin(self.clalist, axis=1)  # 本行代码按每行最小值分组， 如self.clalist= [[1 5]
        #                                                                                                      [2,6]
        #                                                                                                      [3,2]], 得到的是[ 0 0 1 ]

        self.newCentroids = pd.DataFrame(dataSet).groupby(self.minDistIndices).mean()  # 得到新的聚类中心
        self.newCentroids = self.newCentroids.values
        self.changed = self.newCentroids - centroids
        return self.changed, self.newCentroids

    # kmeans（）是这个大类的主方法
    def kmeans(self, dataSet, k, path, outfilepath):  # dataSet：数据  k：聚成的类别数  ，path：预处理的csv文件路径
        self.centroids = random.sample(dataSet, k)
        self.changed, self.newCentroids = self.classify(self.list_data, self.centroids, k)

        while np.any(self.changed != 0):
            self.changed, self.newCentroids = self.classify(dataSet, self.newCentroids, k)

        self.final_centroids = sorted(self.newCentroids.tolist())
        self.clalist = self.calcDis(dataSet, self.final_centroids, k)  # 再调用calcDis（）方法，得到欧拉距离
        self.minDistIndices = np.argmin(self.clalist, axis=1)
        self.category_list = self.minDistIndices.tolist()
        self.original_data = pd.read_csv(path, low_memory=False)
        self.original_df = pd.DataFrame(self.original_data)
        self.original_df['category'] = self.category_list
        outfilename = 'new_data.csv'
        self.original_df.to_csv(outfilepath +'/'+ outfilename, index=None)  # 生成新的csv文件，默认路径在此脚本同目录下
        self.cluster = []

        for i in range(k):
            self.cluster.append([])
        for i, j in enumerate(self.minDistIndices):  # enumerate()可同时遍历索引和遍历元素
            self.cluster[j].append(dataSet[i])

        print(f'最终质心为：{self.final_centroids}')
        for i in range(k):
            print(f'集群{i + 1}为：{self.cluster[i]}')

        return self.final_centroids, self.cluster

    # 画图的方法，只能画二维图
    def plot(self, centroids, cluster, k, column):
        if column > 2:
            return
        for i in range(len(centroids)):
            plt.scatter(centroids[i][0], centroids[i][1], color='k', s=50)
        for i in range(k):
            for j in range(len(cluster[i])):
                if i == 0:
                    plt.scatter(cluster[i][j][0], cluster[i][j][1], color='r', s=5)
                if i == 1:
                    plt.scatter(cluster[i][j][0], cluster[i][j][1], color='b', s=5)
                if i == 2:
                    plt.scatter(cluster[i][j][0], cluster[i][j][1], color='g', s=5)
        plt.show()
# if __name__ == '__main__':
#     input_dir = 'D:\PycharmProjects\pycode\data_process/df_finaldata.csv'  # 文件路径是csv所在路径
#     k = 2  # 聚成几类就改成几
#     column = 2  # 数据有几维，就改成几，column和此脚本第18行的usecols对应，usecols是取的csv文件的哪些列
#     outfilepath = 'D:/'
#     col = [1, 2]
#     object = K_means(input_dir, k, len(col), outfilepath, col)  # 实例化