import tensorflow as tf
from tensorflow.python.saved_model import builder as saved_model_builder
from tensorflow.python.saved_model import tag_constants, signature_constants, signature_def_utils_impl

import numpy as np
import pandas as pd
import os

from tensorflow.python.keras.backend import set_session
from sklearn.preprocessing import MinMaxScaler
from tensorflow.python.keras.models import Sequential
from tensorflow.python.keras.layers import Input, Dense, GRU, Embedding, LSTM, Lambda
from tensorflow.python.keras.optimizers import RMSprop
from tensorflow.python.keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard, ReduceLROnPlateau
from tensorflow.python.keras.models import model_from_json

# for plotting
import matplotlib
import matplotlib.pyplot as plt
import time

import os
import shutil
from google.cloud import storage
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data

###############################################################
# SETUP
###############################################################

# define input arguments
tf.app.flags.DEFINE_string('version', '1',
'model version')
tf.app.flags.DEFINE_string('bucket', 'reti-iot-gke-prod-cluster-ml-bucket',
'name of the GCS bucket')
tf.app.flags.DEFINE_string('dest', 'metering-ac/site-1-hour/1',
'subfolder of the GCS bucket')
tf.app.flags.DEFINE_integer('steps', 2000,
'number of runs through entire training set')
arg_version = tf.app.flags.FLAGS.version
arg_bucket = tf.app.flags.FLAGS.bucket
arg_dest = tf.app.flags.FLAGS.dest
arg_steps = tf.app.flags.FLAGS.steps

print("arg_version: " + arg_version + " arg_bucket: " + arg_bucket + " arg_dest: " + arg_dest)
print("tensorflow version:" + tf.__version__)
# model_version = str(int(time.time()))
# model_path = './models/'+model_version
DO_TRAINING = True

WARMUP_STEPS = 20
###############################################################
# TRAINNING FUNCTIONS
###############################################################

def scaleMinMax(x, min, max):
return (x - min) / (max - min)

def inverseScaleMinMax(x, min, max):
return x * (max - min) + min

def batch_generator(batch_size, sequence_length, num_x_signals, num_y_signals, num_train, x_train, y_train):
"""
Generator function for creating random batches of training-data.
"""

# Infinite loop.
while True:
# Allocate a new array for the batch of input-signals.
x_shape = (batch_size, sequence_length, num_x_signals)
x_batch = np.zeros(shape=x_shape, dtype=np.float16)

# Allocate a new array for the batch of output-signals.
y_shape = (batch_size, sequence_length, num_y_signals)
y_batch = np.zeros(shape=y_shape, dtype=np.float16)

# Fill the batch with random sequences of data.
for i in range(batch_size):
# Get a random start-index.
# This points somewhere into the training-data.
idx = np.random.randint(num_train - sequence_length)
# Copy the sequences of data starting at this index.
# x_batch[i] = x_train_scaled[idx:idx+sequence_length]
x_batch[i] = x_train[idx:idx+sequence_length]
y_batch[i] = y_train[idx:idx+sequence_length]
yield (x_batch, y_batch)
def loss_mse_warmup (y_true, y_pred):

warmup_steps = WARMUP_STEPS
"""
Calculate the Mean Squared Error between y_true and y_pred,
but ignore the beginning "warmup" part of the sequences.
y_true is the desired output.
y_pred is the model's output.
"""

# The shape of both input tensors are:
# [batch_size, sequence_length, num_y_signals].

# Ignore the "warmup" parts of the sequences
# by taking slices of the tensors.
y_true_slice = y_true[:, warmup_steps:, :]
y_pred_slice = y_pred[:, warmup_steps:, :]

# These sliced tensors both have this shape:
# [batch_size, sequence_length - warmup_steps, num_y_signals]

# Calculate the MSE loss for each value in these tensors.
# This outputs a 3-rank tensor of the same shape.
loss = tf.losses.mean_squared_error(labels=y_true_slice,
predictions=y_pred_slice)

# Keras may reduce this across the first axis (the batch)
# but the semantics are unclear, so to be sure we use
# the loss across the entire tensor, we reduce it to a
# single scalar with the mean function.
loss_mean = tf.reduce_mean(loss)

return loss_mean


###############################################################
# PLOT FUNCTIONS
###############################################################

def plot_comparison (model, x, y, target_names, start_idx, length=100):
"""
Plot the predicted and true output-signals.
:param start_idx: Start-index for the time-series.
:param length: Sequence-length to process and plot.
:param train: Boolean whether to use training- or test-set.
"""
y_true = y
# End-index for the sequences.
end_idx = start_idx + length
# Select the sequences from the given start-index and
# of the given length.
x = x[start_idx:end_idx]
y_true = y_true[start_idx:end_idx]
# Input-signals for the model.
x = np.expand_dims(x, axis=0)

# Use the model to predict the output-signals.
y_pred = model.predict(x)
# The output of the model is between 0 and 1.
# Do an inverse map to get it back to the scale
# of the original data-set.
# y_pred_rescaled = y_scaler.inverse_transform(y_pred[0])
y_pred_rescaled = y_pred[0]
# For each output-signal.
for signal in range(len(target_names)):
# Get the output-signal predicted by the model.
signal_pred = y_pred_rescaled[:, signal]
# Get the true output-signal from the data-set.
signal_true = y_true[:, signal]

# Make the plotting-canvas bigger.
plt.figure(figsize=(15,5))
# Plot and compare the two signals.
plt.plot(signal_true, label='true')
plt.plot(signal_pred, label='pred')
# Plot grey box for warmup-period.
p = plt.axvspan(0, WARMUP_STEPS, facecolor='black', alpha=0.15)
# Plot labels etc.
plt.ylabel(target_names[signal])
plt.legend()
plt.show()

###############################################################
# MAIN FUNCTION
###############################################################
def main():
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
# config.gpu_options.per_process_gpu_memory_fraction = 0.7
# config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
set_session(sess) # set this TensorFlow session as the default session for Keras

###############################################################
# FETCH DATA
###############################################################


###############################################################
# DATA PREPROCESSING
###############################################################

df = pd.read_csv("data/myoungsung_manufacturer/myoungsung_office.csv")
df = df.dropna()
df['date'] = pd.to_datetime(df["date"])
df.set_index("date", inplace=True)

# for i in range(1, len(df)):
# df.loc[i,'power_consumption_last']
power_consumption_last = 0
for i, row in df.iterrows():
# print(i)
if i - pd.Timedelta(hours=1) in df.index:
power_consumption_last += df.at[i,'power_consumption_increment']
df.at[i,'power_consumption_last'] = power_consumption_last

df.drop('power_consumption_increment', axis=1, inplace=True)

df['day'] = df.index.dayofyear
df['hour'] = df.index.hour
# df = df.dropna()

target_names = ['power_consumption_last']

shift_days = 1
shift_steps = shift_days * 24 # Number of hours.

df_targets = df[target_names].shift(-shift_steps)

x_data = df.values[:-shift_steps]
y_data = df_targets.values[:-shift_steps]


num_data = len(x_data)
train_split = 0.9
num_train = int(train_split * num_data)

num_test = num_data - num_train

x_train = x_data[:num_train]
x_test = x_data[num_train:]

y_train = y_data[:num_train]
y_test = y_data[num_train:]

num_x_signals = x_data.shape[1]
num_y_signals = y_data.shape[1]

# df.to_csv("x_test.csv", encoding='utf-8')
# df_targets.to_csv("y_test.csv", encoding='utf-8')

# x_scaler = MinMaxScaler()
# x_train_scaled = x_scaler.fit_transform(x_train)
# x_test_scaled = x_scaler.transform(x_test)

min_X = x_train.min()
max_X = x_train.max()

# y_scaler = MinMaxScaler()
# y_train_scaled = y_scaler.fit_transform(y_train)
# y_test_scaled = y_scaler.transform(y_test)

min_Y = y_train.min()
max_Y = y_train.max()


# batch_size = int(256 / 4)
# sequence_length = int(24 * 7 * 8 / 8)
batch_size = 150
sequence_length = 24

generator = batch_generator(batch_size=batch_size, sequence_length=sequence_length, num_x_signals=num_x_signals, num_y_signals=num_y_signals, num_train=num_train, x_train=x_train, y_train=y_train)

# x_batch, y_batch = next(generator)

# batch = 0 # First sequence in the batch.
# signal = 0 # First signal from the 20 input-signals.
# seq = x_batch[batch, :, signal]
# plt.plot(seq)
# plt.show()

# validation_data = (np.expand_dims(x_test_scaled, axis=0), np.expand_dims(y_test_scaled, axis=0))
validation_data = (np.expand_dims(x_test, axis=0), np.expand_dims(y_test, axis=0))
# warmup_steps = 20

optimizer = RMSprop(lr=1e-3)


if DO_TRAINING :

model = Sequential()

model.add(Lambda(lambda x: scaleMinMax(x, min_X, max_X ), input_shape=(None, num_x_signals,)))

model.add(LSTM(units=20, return_sequences=True, input_shape=(None, num_x_signals,)))
model.add(Dense(num_y_signals, activation='sigmoid'))

model.add(Lambda(lambda x: inverseScaleMinMax(x, min_Y, max_Y ), input_shape=(None, num_y_signals,)))

model.compile(loss=loss_mse_warmup, optimizer=optimizer)
model.summary()

path_checkpoint = '23_checkpoint.keras'
callback_checkpoint = ModelCheckpoint(filepath=path_checkpoint,
monitor='val_loss',
verbose=1,
save_weights_only=True,
save_best_only=True)

callback_early_stopping = EarlyStopping(monitor='val_loss',
patience=5, verbose=1)

callback_tensorboard = TensorBoard(log_dir='./23_logs/',
histogram_freq=0,
write_graph=False)

callback_reduce_lr = ReduceLROnPlateau(monitor='val_loss',
factor=0.1,
min_lr=1e-4,
patience=0,
verbose=1)

callbacks = [callback_early_stopping,
callback_checkpoint,
callback_tensorboard,
callback_reduce_lr]

model.fit_generator(generator=generator,
epochs=20,
steps_per_epoch=100,
validation_data=validation_data,
callbacks=callbacks)


try:
model.load_weights(path_checkpoint)
except Exception as error:
print("Error trying to load checkpoint.")
print(error)
# serialize model to JSON
model_json = model.to_json()
with open("model.json", "w") as json_file:
json_file.write(model_json)
# serialize weights to HDF5
model.save_weights("model.h5")
print("Saved model to disk")

x = model.input
y = model.output
prediction_signature = tf.saved_model.signature_def_utils.predict_signature_def({"inputs": x}, {"predictions":y})

valid_prediction_signature = tf.saved_model.signature_def_utils.is_valid_signature(prediction_signature)
if(valid_prediction_signature == False):
raise ValueError("Error: Prediction signature not valid!")

export_path = arg_dest + "/" + arg_version
# export_path = arg_version
# builder = tf.saved_model.builder.SavedModelBuilder(export_path)
if os.path.exists(export_path):
shutil.rmtree(export_path)
builder = saved_model_builder.SavedModelBuilder(export_path)
legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
builder.add_meta_graph_and_variables(
sess, [tag_constants.SERVING],
signature_def_map={
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:prediction_signature,
},
legacy_init_op=legacy_init_op)

builder.save()

###############################################################
# UPLOAD MODEL TO GOOGLE CLOUD STORAGE
###############################################################

print("uploading to " + export_path)
client = storage.Client()
bucket = client.get_bucket(arg_bucket)
if bucket:
for root, dirs, files in os.walk(export_path):
for file in files:
path = os.path.join(root, file)
print(path)
blob = bucket.blob(path)
blob.upload_from_filename(path)


else :
# load json and create model
json_file = open('model.json', 'r')
loaded_model_json = json_file.read()
json_file.close()
model = model_from_json(loaded_model_json, custom_objects={'scaleMinMax': scaleMinMax, 'inverseScaleMinMax': inverseScaleMinMax })
# load weights into new model
model.load_weights("model.h5")
print("Loaded model from disk")
# evaluate loaded model on test data
model.compile(loss=loss_mse_warmup, optimizer=optimizer)
model.summary()

if False:
from keras.initializers import RandomUniform

# Maybe use lower init-ranges.
init = RandomUniform(minval=-0.05, maxval=0.05)

model.add(Dense(num_y_signals,
activation='linear',
kernel_initializer=init))

# result = model.evaluate(x=np.expand_dims(x_test_scaled, axis=0),
# y=np.expand_dims(y_test_scaled, axis=0))
result = model.evaluate(x=np.expand_dims(x_test, axis=0),
y=np.expand_dims(y_test, axis=0))
print("loss (test-set):", result)

# If you have several metrics you can use this instead.
if False:
for res, metric in zip(result, model.metrics_names):
print("{0}: {1:.3e}".format(metric, res))

train=False

if train:
# Use training-data.
# x = x_train_scaled
x = x_train
y_true = y_train
else:
# Use test-data.
# x = x_test_scaled
x = x_test
y_true = y_test

# plot_comparison(start_idx=100000, length=1000, train=True)
# plot_comparison(start_idx=200000, length=1000, train=True)

plot_comparison(model=model, x=x, y=y_true, target_names=target_names,start_idx=1, length=1000)

if __name__ == '__main__':
main()