#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue May  7 10:09:05 2019

@author: mingkuan
"""

import keras.backend as K

import numpy as np
import pylab as plt
from keras.utils import to_categorical
from keras.models import Model
from keras.layers import Input
from keras.layers import LSTM
from keras.layers import Dense
from keras.layers.convolutional import Conv3D
#from keras.layers.convolutional_recurrent import ConvLSTM2D
from keras.layers.normalization import BatchNormalization
from keras.models import load_model
from keras.callbacks import EarlyStopping, ModelCheckpoint
import matplotlib.pyplot as plt
from keras.utils import to_categorical
from keras.regularizers import l2
from keras.backend import clip 


#moving = np.load('/home/mingkuan/Desktop/mnist_test_seq.npy')
moving = np.load('/home/mingkuan/Desktop/arr_0.npy')
moving.shape # (20, 10000, 64, 64)

moving = moving.reshape((12000,14,64,64))


dt = moving.reshape(12000,14,64,64,1)
dt.shape

del moving


#  train set
X1_moving = dt[:10000,:7,:,:,:]
X1_moving.shape

X2_moving = np.concatenate((np.zeros((10000,1,64,64,1)), dt[:10000,7:13,:,:,:]), axis=1) 
X2_moving.shape

y_moving = dt[:10000,7:14,:,:,:]
y_moving.shape

# self-defined cross entropy loss
#def customLoss(yTrue,yPred):
#    return K.sum(-yTrue*K.log(yPred) + (1-yTrue)*K.log(1-yPred))

which = np.random.randint(12000)
track = dt[which][:,:,:,:]


for i in range(14):
    fig = plt.figure(figsize=(10, 5))

    ax = fig.add_subplot(121)

    toplot = track[i, ::, ::, 0]
    plt.imshow(toplot)





"""  
64*64 image
1-layer stacked ConvLSTM2D Encoder-Decoder
"""
def define_models_1_moving_1(n_filter, filter_size):
    # define training encoder
    encoder_inputs = Input(shape=(None, 64, 64, 1))
    encoder_1 = ConvLSTM2D(filters = n_filter, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    encoder_outputs_1, encoder_state_h_1, encoder_state_c_1 = encoder_1(encoder_inputs)
    # define training decoder
    decoder_inputs = Input(shape=(None, 64, 64, 1))
    decoder_1 = ConvLSTM2D(filters=n_filter, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    decoder_outputs_1, _, _ = decoder_1([decoder_inputs, encoder_state_h_1, encoder_state_c_1])
    decoder_conv3d = Conv3D(filters=1, kernel_size=(1,1,64), padding='same', data_format='channels_last',
                            kernel_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    decoder_outputs = decoder_conv3d(decoder_outputs_1)       
    model = Model([encoder_inputs, decoder_inputs], decoder_outputs)
    
    # define inference encoder
    encoder_model = Model(encoder_inputs, [encoder_state_h_1, encoder_state_c_1])
    
    # define inference decoder
    decoder_state_input_h_1 = Input(shape=(64,64,n_filter))
    decoder_state_input_c_1 = Input(shape=(64,64,n_filter))
    decoder_output_1, decoder_state_h_1_new, decoder_state_c_1_new = decoder_1([decoder_inputs, decoder_state_input_h_1, decoder_state_input_c_1])
    decoder_output = decoder_conv3d(decoder_output_1)
    decoder_model = Model([decoder_inputs , decoder_state_input_h_1 , decoder_state_input_c_1],
                          [decoder_output, decoder_state_h_1_new, decoder_state_c_1_new]) 
    return model, encoder_model, decoder_model


train_1_moving_1, infenc_1_moving_1, infdec_1_moving_1 = define_models_1_moving_1(n_filter=128, filter_size=5)

train_1_moving_1.compile(loss='mse', optimizer='adam', metrics=['mae'])#, metrics=['mse'])

#train_2_moving_1.fit([X1_moving,X2_moving], y_moving, batch_size=8, validation_split=0.3, epochs=1)
es = EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=5)
#cp = ModelCheckpoint('model-mnist-2layer.h5', verbose=1, save_best_only=True)
history_1_moving_1 = train_1_moving_1.fit([X1_moving,X2_moving], y_moving, batch_size=8, 
                                          validation_split=0.25, epochs=100, callbacks=[es])



plt.plot(history_1_moving_1.history['loss'][:])
plt.plot(history_1_moving_1.history['val_loss'][:])
plt.title('1-layer convlstm model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()



def predict_sequence_1(infenc, infdec, source, n_steps):
	 encode
	state_h_1, state_c_1 = infenc.predict(source)  # source_dim = ()
	decoder_input = source[:,-1,:,:,:].reshape((1,1,64,64,1))
    #decoder_input = decoder_input.astype('float')
	output = list()
	for t in range(n_steps):
		# predict next char
		yhat, h_1, c_1 = infdec.predict([decoder_input, state_h_1, state_c_1])
		# store prediction
		output.append(yhat[0,0,:])
		# update state
		state_h_1, state_c_1 = h_1, c_1
		# update target sequence
		decoder_input = yhat
	return np.array(output)






"""  
64*64 image
2-layer stacked ConvLSTM2D Encoder-Decoder
"""

def define_models_2_moving_1(n_filter, filter_size):
    # define training encoder
    encoder_inputs = Input(shape=(None, 64, 64, 1))
    encoder_1 = ConvLSTM2D(filters = n_filter, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    encoder_2 = ConvLSTM2D(filters = n_filter, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    encoder_outputs_1, encoder_state_h_1, encoder_state_c_1 = encoder_1(encoder_inputs)
    encoder_outputs_2, encoder_state_h_2, encoder_state_c_2 = encoder_2(encoder_outputs_1)
    # define training decoder
    decoder_inputs = Input(shape=(None, 64, 64, 1))
    decoder_1 = ConvLSTM2D(filters=n_filter, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    decoder_2 = ConvLSTM2D(filters=n_filter, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    decoder_outputs_1, _, _ = decoder_1([decoder_inputs, encoder_state_h_1, encoder_state_c_1])
    decoder_outputs_2, _, _ = decoder_2([decoder_outputs_1, encoder_state_h_2, encoder_state_c_2])
    decoder_conv3d = Conv3D(filters=1, kernel_size=(1,1,64), padding='same', data_format='channels_last',
                            kernel_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    decoder_outputs = decoder_conv3d(decoder_outputs_2)
    #clip(dec oder_outputs, 0, 255)
    
#    denselayer = Dense(1, activation='softmax')
#    decoder_outputs = denselayer(decoder_outputs)
    
    
    model = Model([encoder_inputs, decoder_inputs], decoder_outputs)
    #print(model.summary(line_length=250))
    
    # define inference encoder
    encoder_model = Model(encoder_inputs, [encoder_state_h_1, encoder_state_c_1, encoder_state_h_2, encoder_state_c_2])
    
    # define inference decoder
    decoder_state_input_h_1 = Input(shape=(64,64,n_filter))
    decoder_state_input_c_1 = Input(shape=(64,64,n_filter))
    decoder_state_input_h_2 = Input(shape=(64,64,n_filter))
    decoder_state_input_c_2 = Input(shape=(64,64,n_filter))
    decoder_output_1, decoder_state_h_1_new, decoder_state_c_1_new = decoder_1([decoder_inputs, decoder_state_input_h_1, decoder_state_input_c_1])
    decoder_output_2, decoder_state_h_2_new, decoder_state_c_2_new = decoder_2([decoder_output_1, decoder_state_input_h_2, decoder_state_input_c_2])
    decoder_output = decoder_conv3d(decoder_output_2)
    #clip(decoder_output, 0, 255)
    
#    decoder_output = denselayer(decoder_output)
    
    decoder_model = Model([decoder_inputs , decoder_state_input_h_1 , decoder_state_input_c_1, decoder_state_input_h_2 , decoder_state_input_c_2],
                          [decoder_output, decoder_state_h_1_new, decoder_state_c_1_new, decoder_state_h_2_new, decoder_state_c_2_new])
    
    return model, encoder_model, decoder_model


train_2_moving_1, infenc_2_moving_1, infdec_2_moving_1 = define_models_2_moving_1(n_filter=64, filter_size=5)

train_2_moving_1.compile(loss='mse', optimizer='adam', metrics=['mae'])#, metrics=['mse'])

#train_2_moving_1.fit([X1_moving,X2_moving], y_moving, batch_size=8, validation_split=0.3, epochs=1)
es = EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=5)
#cp = ModelCheckpoint('model-mnist-2layer.h5', verbose=1, save_best_only=True)
history_2_moving_1 = train_2_moving_1.fit([X1_moving,X2_moving], y_moving, batch_size=8, 
                                          validation_split=0.25, epochs=100, callbacks=[es])



plt.plot(history_2_moving_1.history['loss'][:])
plt.plot(history_2_moving_1.history['val_loss'][:])
plt.title('2-layer convlstm model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()



# generate target given source sequence
def predict_sequence_2(infenc, infdec, source, n_steps):
	# encode
	state_h_1, state_c_1, state_h_2, state_c_2 = infenc.predict(source)  # source_dim = ()
	#decoder_input = source[:,-1,:,:,:].reshape((1,1,64,64,1))
	decoder_input = np.repeat(0,64*64).reshape((1,1,64,64,1))
    #decoder_input = decoder_input.astype('float')
	# 123
	output = list()
	for t in range(n_steps):
		# predict next char
		yhat, h_1, c_1, h_2, c_2 = infdec.predict([decoder_input, state_h_1, state_c_1, state_h_2, state_c_2])
		# store prediction
		output.append(yhat[0,0,:])
		# update state
		state_h_1, state_c_1, state_h_2, state_c_2 = h_1, c_1, h_2, c_2
		# update target sequence
		decoder_input = yhat
	return np.array(output)







"""
3-layer stacked ConvLSTM2D Encoder-Decoder
"""

def define_models_3_moving(n_filter, filter_size):
    # define training encoder
    encoder_inputs = Input(shape=(None, 64, 64, 1))
    encoder_1 = ConvLSTM2D(filters = n_filter, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    encoder_2 = ConvLSTM2D(filters = 32, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    encoder_3 = ConvLSTM2D(filters = 32, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    encoder_outputs_1, encoder_state_h_1, encoder_state_c_1 = encoder_1(encoder_inputs)
    encoder_outputs_2, encoder_state_h_2, encoder_state_c_2 = encoder_2(encoder_outputs_1)
    encoder_outputs_3, encoder_state_h_3, encoder_state_c_3 = encoder_3(encoder_outputs_2)
    # define training decoder
    decoder_inputs = Input(shape=(None, 64, 64, 1))
    decoder_1 = ConvLSTM2D(filters=n_filter, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    decoder_2 = ConvLSTM2D(filters=32, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    decoder_3 = ConvLSTM2D(filters=32, kernel_size=filter_size, padding='same', return_sequences=True, return_state=True,
                           kernel_regularizer=l2(0.0005), recurrent_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    decoder_outputs_1, _, _ = decoder_1([decoder_inputs, encoder_state_h_1, encoder_state_c_1])
    decoder_outputs_2, _, _ = decoder_2([decoder_outputs_1, encoder_state_h_2, encoder_state_c_2])
    decoder_outputs_3, _, _ = decoder_3([decoder_outputs_2, encoder_state_h_3, encoder_state_c_3])
    decoder_conv3d = Conv3D(filters=1, kernel_size=(1,1,32), padding='same', data_format='channels_last',
                            kernel_regularizer=l2(0.0005), bias_regularizer=l2(0.0005))
    decoder_outputs = decoder_conv3d(decoder_outputs_3)
    
    model = Model([encoder_inputs, decoder_inputs], decoder_outputs)
    #print(model.summary(line_length=250))
    
    # define inference encoder
    encoder_model = Model(encoder_inputs, 
                          [encoder_state_h_1, encoder_state_c_1, encoder_state_h_2, encoder_state_c_2, encoder_state_h_3, encoder_state_c_3])
    
    # define inference decoder
    decoder_state_input_h_1 = Input(shape=(64,64,n_filter))
    decoder_state_input_c_1 = Input(shape=(64,64,n_filter))
    decoder_state_input_h_2 = Input(shape=(64,64,32))
    decoder_state_input_c_2 = Input(shape=(64,64,32))
    decoder_state_input_h_3 = Input(shape=(64,64,32))
    decoder_state_input_c_3 = Input(shape=(64,64,32))
    decoder_output_1, decoder_state_h_1_new, decoder_state_c_1_new = decoder_1([decoder_inputs, decoder_state_input_h_1, decoder_state_input_c_1])
    decoder_output_2, decoder_state_h_2_new, decoder_state_c_2_new = decoder_2([decoder_output_1, decoder_state_input_h_2, decoder_state_input_c_2])
    decoder_output_3, decoder_state_h_3_new, decoder_state_c_3_new = decoder_3([decoder_output_2, decoder_state_input_h_3, decoder_state_input_c_3])
    decoder_output = decoder_conv3d(decoder_output_3)
    decoder_model = Model([decoder_inputs , decoder_state_input_h_1 , decoder_state_input_c_1, decoder_state_input_h_2 , decoder_state_input_c_2, 
                           decoder_state_input_h_3 , decoder_state_input_c_3],
                          [decoder_output, decoder_state_h_1_new, decoder_state_c_1_new, decoder_state_h_2_new, decoder_state_c_2_new, 
                           decoder_state_h_3_new, decoder_state_c_3_new])
    
    return model, encoder_model, decoder_model


train_3_moving, infenc_3_moving, infdec_3_moving = define_models_3_moving(n_filter=64, filter_size=5)

train_3_moving.compile(loss='mse', optimizer='adam', metrics=['mae'])

#train_3_moving.fit([X1_moving,X2_moving], y_moving, batch_size=8, validation_split=0.1, epochs=1)


#train_2.fit([X1,X2], y, batch_size=32, epochs=300, validation_split=0.1)
es = EarlyStopping(monitor='val_loss', mode='min', verbose=1, patience=5)
history_3_moving = train_3_moving.fit([X1_moving, X2_moving], y_moving, batch_size=4, validation_split=0.25, epochs=100, callbacks=[es])


plt.plot(history_3_moving.history['loss'])
plt.plot(history_3_moving.history['val_loss'])
plt.title('3-layer convlstm model loss')
plt.ylabel('loss')
plt.xlabel('epoch')
plt.legend(['train', 'test'], loc='upper left')
plt.show()



# generate target given source sequence
def predict_sequence_3(infenc, infdec, source, n_steps):
	# encode
	state_h_1, state_c_1, state_h_2, state_c_2, state_h_3, state_c_3 = infenc.predict(source)  # source_dim = ()
#	decoder_input = source[:,-1,:,:,:].reshape((1,1,64,64,1))
	decoder_input = np.repeat(0,64*64).reshape((1,1,64,64,1))
	# collect predictions
	output = list()
	for t in range(n_steps):
		# predict next char
		yhat, h_1, c_1, h_2, c_2, h_3, c_3 = infdec.predict([decoder_input, state_h_1, state_c_1, state_h_2, state_c_2, state_h_3, state_c_3])
		# store prediction
		output.append(yhat[0,0,:])
		# update state
		state_h_1, state_c_1, state_h_2, state_c_2, state_h_3, state_c_3 = h_1, c_1, h_2, c_2, h_3, c_3
		# update target sequence
		decoder_input = yhat
	return np.array(output)
















"""
Prediction
"""  
which = np.random.randint(12000) 
track = dt[which,:10,:,:,:]
track.shape

history = track[np.newaxis, ::, ::, ::, ::]
history.shape

prediction_lstm = predict_sequence_lstm(infenc_1_lstm_1, infdec_1_lstm_1, history, 10)
prediction_1 = predict_sequence_1(infenc_1_moving_1, infdec_1_moving_1, history, 10)
prediction_2 = predict_sequence_2(infenc_2_moving_1, infdec_2_moving_1, history, 10)
prediction_3 = predict_sequence_3(infenc_3_moving, infdec_3_moving, history, 10)

prediction_lstm.shape
prediction_1.shape
prediction_2.shape
prediction_3.shape

ground_truth = dt[which][10:20, :, :, ::] 
ground_truth.shape 


fig = plt.figure(figsize=(20,24))
for i in range(7):
    ax = fig.add_subplot(6,7,i+1)
    toplot = track[i+3, ::, ::, 0]
    plt.imshow(toplot)
    
    ax = fig.add_subplot(6,7,i+8)
    toplot = ground_truth[i, ::, ::, 0]
    plt.imshow(toplot)
    
    ax = fig.add_subplot(6,7,i+15)
    toplot = prediction_lstm[i, ::, ::, 0]
    plt.imshow(toplot)
    
    ax = fig.add_subplot(6,7,i+22)
    toplot = prediction_1[i, ::, ::, 0]
    plt.imshow(toplot)
    
    ax = fig.add_subplot(6,7,i+29)
    toplot = prediction_2[i, ::, ::, 0]
    plt.imshow(toplot)
    
    ax = fig.add_subplot(6,7,i+36)
    toplot = prediction_3[i, ::, ::, 0]
    plt.imshow(toplot)


"""
try prediction on 3-digit moving MNIST
Generate 3-digit moving MNIST data and stored as arr_0.npy
"""
moving_3_digit = np.load('/home/mingkuan/Desktop/arr_0.npy')


'''
Test set performance
'''
prediction_1_list = np.zeros((2000,7,64,64,1))
prediction_2_list = np.zeros((2000,7,64,64,1))
prediction_3_list = np.zeros((2000,7,64,64,1))

for i in range(10000,12000):
    track = dt[i,:7,:,:,:]

    history = track[np.newaxis, ::, ::, ::, ::]
	
    prediction_1 = predict_sequence_1(infenc_1_moving_1, infdec_1_moving_1, history, 7)
    prediction_2 = predict_sequence_2(infenc_2_moving_1, infdec_2_moving_1, history, 7)
    prediction_3 = predict_sequence_3(infenc_3_moving, infdec_3_moving, history, 7)
    
    prediction_1_list[i-10000,:,:,:,:] = prediction_1
    prediction_2_list[i-10000,:,:,:,:] = prediction_2
    prediction_3_list[i-10000,:,:,:,:] = prediction_3
    
prediction_1_list = np.array(prediction_1_list).reshape(2000,7,64,64)
prediction_2_list = np.array(prediction_2_list).reshape(2000,7,64,64)
prediction_3_list = np.array(prediction_3_list).reshape(2000,7,64,64)

ground_truth = dt[10000:,7:,:,:,0]
ground_truth.shape

def mse(x,y):
    sum_xy=0
    for i in range(64):
        for j in range(64):
            sum_xy = sum_xy + (x[i,j]-y[i,j])**2
    mse = sum_xy/(64*64)    
    return mse

def mae(x,y):
    sum_xy=0
    for i in range(64):
        for j in range(64):
            sum_xy = sum_xy + np.abs(x[i,j]-y[i,j])
    mae = sum_xy/(64*64)    
    return mae

mse_1 = 0
mae_1 = 0
mse_2 = 0
mae_2 = 0
mse_3 = 0
mae_3 = 0

for i in range(2000): 
    for j in range(7):
        mse_1 = mse_1 + mse(prediction_1_list[i,j,:,:],ground_truth[i,j,:,:])
        mae_1 = mae_1 + mae(prediction_1_list[i,j,:,:],ground_truth[i,j,:,:])
        mse_2 = mse_2 + mse(prediction_2_list[i,j,:,:],ground_truth[i,j,:,:])
        mae_2 = mae_2 + mae(prediction_2_list[i,j,:,:],ground_truth[i,j,:,:])
        mse_3 = mse_3 + mse(prediction_3_list[i,j,:,:],ground_truth[i,j,:,:])
        mae_3 = mae_3 + mae(prediction_3_list[i,j,:,:],ground_truth[i,j,:,:])

mse_1/(2000*7)
mae_1/(2000*7)
mse_2/(2000*7)
mae_2/(2000*7)
mse_3/(2000*7)
mae_3/(2000*7)


