# -*- coding: utf-8 -*-
"""
 @Time    : 2020/5/29 13:52
 @Author  : 杰森·家乐森
 @File    : ae_train.py
 @Software: PyCharm
"""

import time
import json
import numpy as np
import pandas as pd
import tensorflow as tf
from sklearn.metrics import r2_score
from tensorflow.keras import backend
from tensorflow.keras.models import Model
from sklearn.preprocessing import MinMaxScaler
from tensorflow.keras.layers import Dense, Input


def rmse(y_true, y_pred):
    return backend.sqrt(backend.mean(tf.keras.losses.mean_squared_error(y_true, y_pred), axis=-1))


def train(layers, data, test, epoch):
    """
    自编码训练函数
    :param data: 归一化数据
    :param layers: 网络结构，int类型数组
    :param epoch: 训练次数
    :return:
    """
    mms = MinMaxScaler()
    data = mms.fit_transform(data)
    test = mms.transform(test)
    mid = len(layers) // 2 - 1
    layer = [Input(shape=(layers[0],))]
    for i in range(len(layers) - 1):
        if i == mid:
            layer.append(Dense(layers[i + 1])(layer[i]))
        else:
            layer.append(Dense(layers[i + 1], activation="sigmoid")(layer[i]))
    autoencoder = Model(layer[0], layer[-1])
    autoencoder.compile(optimizer="adam", loss="mse",
                        metrics=['accuracy'])
    autoencoder.fit(data, data, epochs=epoch, batch_size=32, shuffle=True)
    test_data = autoencoder.predict(data, batch_size=400)
    test_out = autoencoder.predict(test, batch_size=400)
    with tf.Session():
        spe = rmse(data, test_data).eval()
        limit = 3 * spe
        mse = rmse(test, test_out).eval()
    # r2 = r2_score(data, test_data)
    r2 = r2_score(test, test_out)
    weight = autoencoder.get_weights()
    weights = []
    bias = []
    is_weight = True
    model = {
        "weights": weights,
        "bias": bias
    }
    for i in range(len(weight)):
        if is_weight:
            weights.append(weight[i].tolist())
            is_weight = False
        else:
            bias.append(weight[i].tolist())
            is_weight = True
    return model, float(limit), r2, float(mse)


if __name__ == '__main__':
    with open('data.json', 'r') as f:
        data = np.array(json.load(f)['data'])
    layer_arr = [
        [10, 1, 10],
        [10, 8, 1, 8, 10],
        [10, 8, 6, 1, 6, 8, 10]
    ]
    for i in range(len(layer_arr)):
        k = len(layer_arr[i]) // 2
        for j in range(5):
            layer_arr[i][k] = j + 1
            model, limit, r2, mse = train(layer_arr[i], data[:3000, :], data, 5000)
            model_config = {
                "layer": layer_arr[i],
                "r2": r2,
                "mse": mse,
                "model": model,
                "spe": limit
            }
            with open('model%d_%d.json' % (len(layer_arr[i]), j + 1), 'w') as f:
                f.write(json.dumps(model_config))
    print('OK')