# tensorflow2.0
import tensorflow as tf
from tensorflow.keras import Input, Model, layers, regularizers
from tensorflow.keras.layers import Dense
import numpy as np


class Encoder(tf.keras.Model):
    def __init__(self, w, h, dim=256):
        super().__init__()
        self.fc1 = Dense(dim, activation='relu', input_shape=(w * h, ))
        self.fc2 = Dense(dim / 2, activation='relu')
        self.fc3 = Dense(dim / 4, activation='relu')

    def call(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x


class Decoder(tf.keras.Model):
    def __init__(self, w, h, dim=256):
        super().__init__()
        Out_dim = w * h

        self.fc1 = Dense(dim / 2, activation='relu', input_shape=(dim, ))
        self.fc2 = Dense(dim, activation='relu')
        self.fc3 = Dense(Out_dim, activation='sigmoid')

    def call(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        x = self.fc3(x)
        return x


class DAE:
    def __init__(self, dim, w, h):
        self.dim = dim
        self.w, self.h = w, h

    def compile(self):
        input_img = Input(shape=(self.w * self.h), )
        self.encoder = Encoder(self.w, self.h, self.dim)
        self.encoded = self.encoder(input_img)

        self.decoder = Decoder(self.w, self.h, self.dim)
        self.decoded = self.decoder(self.encoded)

        self.autoencoder = Model(input_img, self.decoded)
        self.autoencoder(input_img)
        self.autoencoder.compile(optimizer=tf.keras.optimizers.Adam(),
                                 loss='binary_crossentropy')

    def fit(self, x_train, x_test, epochs=10, batch_size=256):
        self.autoencoder.fit(x_train,
                             x_train,
                             epochs=epochs,
                             batch_size=batch_size,
                             shuffle=True,
                             validation_data=(x_test, x_test))


def TEST():
    from tensorflow.keras.datasets import mnist
    import matplotlib.pyplot as plt
    (x_train, _), (x_test, _) = mnist.load_data()
    x_train = x_train.astype('float32') / 255.
    x_test = x_test.astype('float32') / 255.
    x_train = x_train.reshape((len(x_train), np.prod(x_train.shape[1:])))
    x_test = x_test.reshape((len(x_test), np.prod(x_test.shape[1:])))
    dim = 256
    w, h = 28, 28
    ae = DAE(dim, w, h)
    ae.compile()
    ae.fit(x_train, x_test)

    decoded_imgs = ae.autoencoder.predict(x_test)
    n = 10  # how many digits we will display
    plt.figure(figsize=(20, 4))
    for i in range(n):
        # display original
        ax = plt.subplot(2, n, i + 1)
        loc = i + 20
        plt.imshow(x_test[loc].reshape(28, 28))
        plt.gray()
        ax.get_xaxis().set_visible(False)
        ax.get_yaxis().set_visible(False)

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


if __name__ == "__main__":
    TEST()
