# 变分自编码器
from keras.models import Model
from keras.layers import Dense, Input, Lambda
from keras.losses import mse, binary_crossentropy
from keras.utils import plot_model
import keras.backend as K

import pandas as pd
from random import shuffle
from sklearn.preprocessing import StandardScaler
import numpy as np
import tensorflow.keras as keras
import datetime
import matplotlib.pyplot as plt
from keras.layers import Input, Dense, Lambda
from matplotlib.pyplot import MultipleLocator

from keras import losses

data_set=pd.read_feather('data_set_sp.feather')

#打乱并切分训练集测试集
def shuffle_data(dataset_faults):
    sn_fau=list(set(dataset_faults['sn']))
    shuffle(sn_fau)
    newtrain=dataset_faults[dataset_faults['sn'].isin(sn_fau[:int(0.8*len(sn_fau))])]
    newtest=dataset_faults[dataset_faults['sn'].isin(sn_fau[int(0.8*len(sn_fau)):])]
    newtrain.reset_index(drop=True,inplace=True)
    newtest.reset_index(drop=True,inplace=True)
    return newtrain,newtest

#训练集数据标准化
def scaler_train(train):
    Xtrain=train.drop(['time','sn','split'],axis=1)
    Xsc_colnames=list(Xtrain.columns)
    scaler=StandardScaler()
    scaler.fit(Xtrain)  #保存train_sc的均值和标准差
    Xsc=scaler.transform(np.array(Xtrain))
    Xsc=pd.DataFrame(Xsc)
    Xsc.columns=Xsc_colnames
    Xsc['split']=train['split'].values
    return Xsc,scaler
#测试集数据标准化
def scaler_test_train(test,scaler):
    Xtest=test.drop(['time','sn','split'],axis=1)
    Xsc_colnames=list(Xtest.columns)
    Xtsc=scaler.transform(np.array(Xtest))
    Xtsc=pd.DataFrame(Xtsc)
    Xtsc.columns=Xsc_colnames
    Xtsc['split']=test['split'].values
    return Xtsc

train,test=shuffle_data(data_set)
Xsc,scaler = scaler_train(train)
Xtsc = scaler_test_train(test,scaler)

#时间窗口划分
def create_dataset(data_set,time_steps):   #X为dataframe，y为serie
    a=[]
    aa=np.empty(shape=[0,3])
    #index=pd.DataFrame()
    List_n_split=data_set['split'].drop_duplicates()
    for k in List_n_split:
        dataset=data_set[data_set['split']==k]
        #datatrain=data_train[data_train['split']==k]
        if len(dataset)>time_steps:
            dataset=dataset.reset_index(drop=True)
            dataset=dataset.drop(['split'],axis=1)
            dataX= []
            #index_step=[]
            for i in list(range(0,len(dataset)-60,60))[:-1]:      
                v1 = dataset.iloc[i:(i+time_steps)].values
                dataX.append(v1)
                #index_step.append(i)
            #dataset3=dataset.iloc[:len(dataset)-time_steps]
            #newdatatrain=datatrain[:len(dataset3)]
            dataX2=np.array(dataX)
            a.append(dataX2)             
            #index=index.append(newdatatrain)
    if len(a)>0:
        aa=np.vstack(a)
    #index.reset_index(drop=True,inplace=True)
    return aa

aa = create_dataset(Xsc,time_steps=120)
bb = create_dataset(Xtsc,time_steps=120)

def sampling(args):
    # 再参数化采样
    z_mean, z_log_var = args[0],args[1]
    batch = K.shape(z_mean)[0]
    dim = K.int_shape(z_mean)[1]
    #epsilon = K.random_normal(shape=(batch, dim))
    epsilon = K.random_normal(shape=(batch, dim,args[2]))
    return z_mean + K.exp(0.5 * z_log_var) * epsilon

def get_loss(args):
    """
    自定义损失函数
    x:原始样本
    xr: 重构样本
    m: 编码器隐变量z均值
    v: 编码器隐变量z方差的对数
    """
    x, xr, m, v = args
    dim = K.int_shape(x)[-1]
    print(dim)
    re_loss = dim * binary_crossentropy(x, xr)  # 重构正确性度量
    kl_loss = 1 + v - K.square(m) - K.exp(v)  # d维向量，隐变量z维度为d
    kl_loss = - 0.5 * K.sum(kl_loss, axis=-1)  # kl散度，罚项
    print(re_loss)
    print(kl_loss)
    #vae_loss = K.mean(re_loss + kl_loss)
    vae_loss =re_loss + kl_loss
    print(vae_loss)
    return vae_loss

# def get_loss(x, x_decoded_mean):
# 	  #my tips:logloss
#     xent_loss = input_dim * losses.binary_crossentropy(x, x_decoded_mean)
#     #my tips:see paper's appendix B
#     kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
#     return xent_loss + kl_loss


# 模型训练参数
batch_size = 128
epochs = 10
# 模型结构参数
input_dim = aa.shape[2] # 输入图像为28*28
latent_dim = 10  # 潜在因子z的维度
# 构建编码器（推断网络）
#inputs = Input(shape=(input_dim,))
inputs = Input(shape=(aa.shape[1],aa.shape[2]))
encode_h = Dense(20, activation='relu')(inputs)
z_mean = Dense(latent_dim, activation='relu')(encode_h)
z_log_var = Dense(latent_dim, activation='relu')(encode_h)  # log(sigma^2)
z_sample = Lambda(sampling, output_shape=([latent_dim]))([z_mean, z_log_var,latent_dim])  # 采样
print(z_sample)
encoder = Model(inputs, [z_mean, z_log_var, z_sample])  # 编码器输出结果为3层
# 构建解码器（生成网络）
inputs_decoder = Input(shape=(aa.shape[1],latent_dim))
print(inputs_decoder)
decode_h = Dense(20, activation='relu')(inputs_decoder)
x_mean_decoded = Dense(input_dim, activation='sigmoid')(decode_h)
decoder = Model(inputs_decoder, x_mean_decoded)
# 构建VAE模型
x_decoded = decoder(encoder(inputs)[2])
print(x_decoded)
outputs = Lambda(get_loss)([inputs, x_decoded, z_mean, z_log_var])  # 模型直接输出损失函数值
#outputs = Lambda(get_loss)(inputs_decoder, x_mean_decoded)  # 模型直接输出损失函数值

#decoder.compile(optimizer='rmsprop', loss=get_loss)


vae_model = Model(inputs, outputs)
vae_model.compile(optimizer='adam', loss=lambda y_true, y_pred: y_pred)
vae_model.fit(aa, aa,shuffle=True,epochs=epochs,verbose=2,batch_size=batch_size)

encoded_imgs = encoder.predict(bb)[2]
decoded_imgs = decoder.predict(encoded_imgs)