# tensorflow2.0
import tensorflow as tf
from tensorflow.keras import Input, Model
from tensorflow.keras.layers import Conv2D, MaxPooling2D, UpSampling2D


class CAE:
    def __init__(self, w, h, input_channel):
        self.w, self.h, self.input_channel = w, h, input_channel

    def compile(self):
        input_img = Input(shape=(self.w, self.h, self.input_channel))
        x = Conv2D(16, (3, 3),
                   strides=(2, 2),
                   activation='relu',
                   padding='same')(input_img)
        x = Conv2D(8, (3, 3),
                   strides=(2, 2),
                   activation='relu',
                   padding='same')(x)
        encoded = Conv2D(8, (3, 3),
                         strides=(2, 2),
                         activation='relu',
                         padding='same')(x)

        x = Conv2D(8, (3, 3), activation='relu', padding='same')(encoded)
        t_b, t_h, t_w, t_c = x.shape
        x = tf.image.resize(x, (t_w * 2, t_h * 2))
        x = Conv2D(8, (3, 3), activation='relu', padding='same')(x)
        t_b, t_h, t_w, t_c = x.shape
        x = tf.image.resize(x, (t_w * 2, t_h * 2))
        x = Conv2D(16, (3, 3), activation='relu', padding='same')(x)
        t_b, t_h, t_w, t_c = x.shape
        x = tf.image.resize(x, (t_w * 2, t_h * 2))
        decoded = Conv2D(1, (3, 3), activation='sigmoid', padding='same')(x)

        self.autoencoder = Model(input_img, decoded)
        self.autoencoder.compile(optimizer='adam', loss='binary_crossentropy')

    def fit(self, x_train, x_test, epochs=5, batch_size=128):

        self.autoencoder.fit(x_train,
                             x_train,
                             epochs=epochs,
                             batch_size=batch_size,
                             shuffle=True,
                             validation_data=(x_test, x_test))


def TEST(SIZE=64):
    '''
	SIZE: Specify image size
	'''
    import sys
    sys.path.append('../..')
    from Utils.load_data import load_mnist
    import matplotlib.pyplot as plt
    x_train, x_test = load_mnist(SIZE=SIZE)
    tb, th, tw, tc = x_test.shape
    w, h, input_channel = th, tw, 1
    cae = CAE(w, h, input_channel)
    cae.compile()
    cae.fit(x_train, x_test, epochs=1)

    decoded_imgs = cae.autoencoder.predict(x_test[:20])

    n = 5
    plt.figure(figsize=(10, 4))
    for i in range(1, n + 1):
        # display original
        ax = plt.subplot(2, n, i)
        plt.imshow(x_test[i].reshape(w, h))
        plt.gray()
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)

        # display reconstruction
        ax = plt.subplot(2, n, i + n)
        plt.imshow(decoded_imgs[i].reshape(w, h))
        plt.gray()
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)
    plt.show()


if __name__ == '__main__':
    TEST()
