import loadData
import numpy as np
import tensorflow as tf
import pandas as pd

'''
class ConcatLabel：
这是一个给data加上label的类
使用示例：ConcatLabel(data, label,if_load)
作用是按是否采用加载将data文件分别添加上label
需要传入的参数有：
data:要添加label的data  是一个list类型或者np.array
label:添加的label为几   最好是int类型
if_load:是否采用按加载来分  true则采用，false则不采用   布尔类型
'''


class ConcatLabel(object):

    def __init__(self, data, label, if_load):
        self.data = data
        self.if_load = if_load
        self.label = label

    def __call__(self, *args, **kwargs):
        data = np.array(self.data)
        if not self.if_load:
            length, dim = data.shape
            label_list = [[self.label] for i in range(length)]
            label_list = np.array(label_list)
            dataWithLabel = tf.concat([data, label_list], axis=1)
            return dataWithLabel
        height, length, dim = data.shape
        data_new = tf.reshape(data, [height * length, dim])
        label_list = [[self.label] for i in range(height * length)]
        label_list = np.array(label_list)
        dataWithLabel = tf.concat([data_new, label_list], axis=1)
        dataWithLabel = tf.reshape(dataWithLabel, [height, length, dim + 1])
        return dataWithLabel


'''
class PieceAndBag：
这是一个给data分成train_data和test_data的类
使用示例：ConcatLabel(data, if_load)
作用是按是否采用加载将data文件分成(train_data,train_label和(test_data,test_label)
需要传入的参数有：
data:要添加label的data  是一个list类型或者np.array
if_load:是否采用按加载来分  true则采用，false则不采用   布尔类型
'''


class PieceAndBag(object):

    def __init__(self, data, if_load):
        self.data = data
        self.if_load = if_load

    def __call__(self, *args, **kwargs):
        data = self.data
        data = np.array(data)
        test_data = []
        train_data = []
        test_label = []
        train_label = []
        if not self.if_load:
            heigth, dim = data.shape
            for i in range(heigth):
                if i % 6 == 0:
                    test_data.append(data[i, :-1])
                    test_label.append(data[i, -1])
                else:
                    train_data.append(data[i, :-1])
                    train_label.append(data[i, -1])
            train_data = np.array(train_data)
            train_label = np.array(train_label)
            test_data = np.array(test_data)
            test_label = np.array(test_label)
            return (train_data, train_label), (test_data, test_label)
        height, length, dim = data.shape
        data = tf.reshape(data, [height * length, dim])
        for i in range(height * length):
            if i % 6 == 0:
                test_data.append(data[i, :-1])
                test_label.append(data[i, -1])
            else:
                train_data.append(data[i, :-1])
                train_label.append(data[i, -1])

        train_data = np.array(train_data)
        train_label = np.array(train_label)
        test_data = np.array(test_data)
        test_label = np.array(test_label)
        print(train_data.shape)
        print(train_label.shape)
        print(test_data.shape)
        print(test_label.shape)
        train_data = tf.reshape(np.array(train_data), [height, length, dim - 1])
        train_label = tf.reshape(np.array(train_label), [height, length, 1])
        test_data = tf.reshape(np.array(test_data), [height, length, dim - 1])
        test_label = tf.reshape(np.array(test_label), [height, length, 1])

        train_data = tf.transpose(train_data, [1, 2, 0])
        train_label = tf.transpose(train_label, [1, 2, 0])
        test_data = tf.transpose(test_data, [1, 2, 0])
        test_label = tf.transpose(test_label, [1, 2, 0])
        return (train_data, train_label), (test_data, test_label)


if __name__ == '__main__':
    data0 = loadData.DataDeal1(9, 'E:\data\DDS_data\平行齿轮箱齿轮表面磨损故障恒速\DATA')
    data = np.array(data0(9, 'E:\data\DDS_data\平行齿轮箱齿轮表面磨损故障恒速\DATA'))
    dataWithLabel1 = ConcatLabel(data, 1, True)
    dataWithLabel1 = np.array(dataWithLabel1(data, 1, True))

    # 导入第二类故障并打标签
    data0 = loadData.DataDeal1(9, 'E:\data\DDS_data\平行齿轮箱齿轮齿根裂纹故障恒速\DATA')
    data = np.array(data0(9, 'E:\data\DDS_data\平行齿轮箱齿轮齿根裂纹故障恒速\DATA'))
    dataWithLabel2 = ConcatLabel(data, 2, True)
    dataWithLabel2 = np.array(dataWithLabel2(data, 2, True))

    # 导入第三类故障并打标签
    data0 = loadData.DataDeal1(9, 'E:\data\DDS_data\平行齿轮箱齿轮断齿故障恒速\DATA')
    data = np.array(data0(9, 'E:\data\DDS_data\平行齿轮箱齿轮断齿故障恒速\DATA'))
    dataWithLabel3 = ConcatLabel(data, 3, True)
    dataWithLabel3 = np.array(dataWithLabel3(data, 3, True))

    # 导入第四类故障并打标签
    data0 = loadData.DataDeal1(9, 'E:\data\DDS_data\平行齿轮箱齿轮偏心故障恒速\DATA')
    data = np.array(data0(9, 'E:\data\DDS_data\平行齿轮箱齿轮偏心故障恒速\DATA'))
    dataWithLabel4 = ConcatLabel(data, 4, True)
    dataWithLabel4 = np.array(dataWithLabel4(data, 4, True))

    # 导入第五类故障并打标签
    data0 = loadData.DataDeal1(9, 'E:\data\DDS_data\平行齿轮箱齿轮缺齿故障恒速\DATA')
    data = np.array(data0(9, 'E:\data\DDS_data\平行齿轮箱齿轮缺齿故障恒速\DATA'))
    dataWithLabel5 = ConcatLabel(data, 5, True)
    dataWithLabel5 = np.array(dataWithLabel5(data, 5, True))

    # 导入第六类故障并打标签
    data0 = loadData.DataDeal1(9, 'E:\data\DDS_data\平行齿轮箱轴承复合故障恒速\DATA')
    data = np.array(data0(9, 'E:\data\DDS_data\平行齿轮箱轴承复合故障恒速\DATA'))
    dataWithLabel6 = ConcatLabel(data, 6, True)
    dataWithLabel6 = np.array(dataWithLabel6(data, 6, True))

    # 导入第七类故障并打标签
    data0 = loadData.DataDeal1(9, 'E:\data\DDS_data\平行齿轮箱轴承滚动体故障恒速\DATA')
    data = np.array(data0(9, 'E:\data\DDS_data\平行齿轮箱轴承滚动体故障恒速\DATA'))
    dataWithLabel7 = ConcatLabel(data, 7, True)
    dataWithLabel7 = np.array(dataWithLabel7(data, 7, True))

    # 导入第八类故障并打标签
    data0 = loadData.DataDeal1(9, 'E:\data\DDS_data\平行齿轮箱轴承内圈故障恒速\DATA')
    data = np.array(data0(9, 'E:\data\DDS_data\平行齿轮箱轴承内圈故障恒速\DATA'))
    dataWithLabel8 = ConcatLabel(data, 8, True)
    dataWithLabel8 = np.array(dataWithLabel8(data, 8, True))

    # 导入第九类故障并打标签
    data0 = loadData.DataDeal1(9, 'E:\data\DDS_data\平行齿轮箱轴承外圈故障恒速\DATA')
    data = np.array(data0(9, 'E:\data\DDS_data\平行齿轮箱轴承外圈故障恒速\DATA'))
    dataWithLabel9 = ConcatLabel(data, 9, True)
    dataWithLabel9 = np.array(dataWithLabel9(data, 9, True))

    data_all = tf.concat(
        [dataWithLabel1, dataWithLabel2, dataWithLabel3, dataWithLabel4, dataWithLabel5, dataWithLabel6, dataWithLabel7,
         dataWithLabel8, dataWithLabel9], axis=1)

    data_all = np.array(data_all)
    print(data_all.shape)
    data_all =tf.transpose(data,[1,2,0])
    print(data_all.shape)
    data_train = tf.random.shuffle(data_all)
    print(data_train)
    data_all = tf.transpose(data, [2, 0, 1])
    data_new = PieceAndBag(data_train, True)
    (train_data, train_label), (test_data, test_label) = data_new(data_train, True)
    train_data = np.array(train_data)
    train_label = np.array(train_label)
    test_data = np.array(test_data)
    test_label = np.array(test_label)

    print(train_data.shape)
    print(train_label.shape)
    print(test_data.shape)
    print(test_label.shape)