# -*- coding: utf-8 -*-
import numpy as np
import tensorflow as tf
tf.compat.v1.disable_v2_behavior()
import math
import matplotlib.pyplot as plt
import pandas as pd
from scipy.io import savemat,loadmat
from sklearn.metrics import r2_score
from sklearn.preprocessing import MinMaxScaler,StandardScaler
tf.compat.v1.reset_default_graph()
tf.compat.v1.set_random_seed(0)
np.random.seed(0)
# In[] Define hyperparameters
def data_split(data,steps):
    in_,out_=[],[]
    samples=len(data)-steps
    for i in range(samples):
        in_.append(data[i:i+steps])
        out_.append(data[i+steps])
    in_=np.array(in_).reshape(len(in_),steps)
    out_=np.array(out_).reshape(len(out_),1)
    return in_,out_
pop=loadmat('result/woa_gru_result.mat')['pop'].reshape(-1,)
print("The optimal hyperparameters are：", [pop[i] if i == 0 else int(pop[i]) for i in range(len(pop))])

alpha = pop[0]# Learning rate
num_epochs = int(pop[1])#Number of iterations
hidden_nodes0 = int(pop[2])#First hidden layer neuron
hidden_nodes = int(pop[3])#Second hidden layer neuron
batch_size = int(pop[4])# batchsize
steps=int(pop[5])#Step size
# In[] 加载数据
data=pd.read_csv('yantailaizhou8192.csv',engine='python')[['salt']].fillna(0).values

in_,out_=data_split(data,steps)
n=range(in_.shape[0])
m=-900
train_data = in_[n[0:m],]
test_data = in_[n[m:],]
train_label = out_[n[0:m],]
test_label = out_[n[m:],]
# Normalized
ss_X=MinMaxScaler(feature_range=(0,1)).fit(train_data)
ss_Y=MinMaxScaler(feature_range=(0,1)).fit(train_label)
train_data = ss_X.transform(train_data)
train_label = ss_Y.transform(train_label)

test_data = ss_X.transform(test_data)
test_label = ss_Y.transform(test_label)

in_num=train_data.shape[1]
out_num=train_label.shape[1]

input_features = in_num
output_class = out_num

# placeholder
X = tf.compat.v1.placeholder("float", [None, input_features])
Y = tf.compat.v1.placeholder("float", [None, output_class])

def RNN(x):
    x = tf.reshape(x , [-1, 1,input_features])
    weights = {'out': tf.Variable(tf.compat.v1.random_normal([hidden_nodes, output_class]))}
    biases = {'out': tf.Variable(tf.compat.v1.random_normal([output_class]))}
    gru_cell0 = tf.compat.v1.nn.rnn_cell.GRUCell(hidden_nodes0)
    gru_cell = tf.compat.v1.nn.rnn_cell.GRUCell(hidden_nodes)
    gru_cell = tf.compat.v1.nn.rnn_cell.MultiRNNCell([gru_cell0,gru_cell])

    init_state = gru_cell.zero_state(tf.shape(x)[0], dtype=tf.float32)
    outputs, _ = tf.compat.v1.nn.dynamic_rnn(gru_cell, x, dtype=tf.float32, initial_state=init_state)
    
    output_sequence = tf.matmul(tf.reshape(outputs, [-1, hidden_nodes]), weights['out']) + biases['out']
    return tf.reshape(output_sequence, [-1, output_class])

logits = RNN(X)
loss = tf.compat.v1.losses.mean_squared_error(predictions = logits, labels = Y)
global_step = tf.Variable(0)
learning_rate = tf.compat.v1.train.exponential_decay(
                alpha,
                global_step,
                num_epochs, 0.99,
                staircase=True)
optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate, epsilon = 1e-10).minimize(loss,global_step=global_step)
init = tf.compat.v1.global_variables_initializer()

# In[]train
train = []
valid=[]
with tf.compat.v1.Session() as sess:
    sess.run(init)
    N = train_data.shape[0]
    for epoch in range(num_epochs):
        total_batch = int(math.ceil(N / batch_size))
        indices = np.arange(N)
        np.random.shuffle(indices)
        avg_loss = 0
        # Iterative training, and calculate the loss of the training set
        for i in range(total_batch):
            rand_index = indices[batch_size*i:batch_size*(i+1)]
            x = train_data[rand_index]
            y = train_label[rand_index]
            _, cost = sess.run([optimizer, loss],
                                feed_dict={X: x, Y: y})
            avg_loss += cost / total_batch
            
        # Calculate the loss of the test set
        valid_data=test_data.reshape(-1,input_features)
        valid_y = test_label.reshape(-1, output_class)
        valid_loss = sess.run(loss, feed_dict={X: valid_data, Y: valid_y})
        
        train.append(avg_loss)
        valid.append(valid_loss)
        print('epoch:',epoch,' ,train loss ',avg_loss,' ,valid loss: ',valid_loss)

    #Calculate the predicted value of the test set
    test_data=test_data.reshape(-1,input_features)
    test_pred = sess.run(logits, feed_dict={X: test_data})
    test_pred = test_pred.reshape(-1, output_class)
# Inverse normalization of test results
test_label  = ss_Y.inverse_transform(test_label)
test_pred   = ss_Y.inverse_transform(test_pred)
savemat('result/woa_gru_result2.mat',{'true':test_label,'pred':test_pred})

# In[] Draw loss curve
g = plt.figure()
plt.ylabel('MSE')
plt.xlabel('Epoch')
plt.plot( train, label='training')
plt.plot( valid, label='testing')
plt.title('loss curve')
plt.legend()
plt.show()

test_pred1=test_pred.reshape(-1,1)
test_label1=test_label.reshape(-1,1)

# mape
test_mape=np.mean(np.abs((test_pred1-test_label1)/test_label1))
# rmse
test_rmse=np.sqrt(np.mean(np.square(test_pred1-test_label1)))
# mae
test_mae=np.mean(np.abs(test_pred1-test_label1))
# R2
test_r2=r2_score(test_label1,test_pred1)

print('mape:',test_mape,' rmse:',test_rmse,' mae:',test_mae,' R2:',test_r2)
# plot test_set result
plt.figure()
plt.plot(test_label1,c='r', label='true')
plt.plot(test_pred1,c='b',label='predict')

plt.legend()
plt.show()