import glob
import os
import time

import imageio
import matplotlib.pyplot as plt
import numpy as np
import PIL
import tensorflow as tf
from IPython import display
from tensorflow.keras import layers

"""
Explanation: Convolutional Autoencoder on MNIST dataset
Learning Objective
1. Build an autoencoder architecture (consisting of an encoder and decoder) in Keras
2. Define the loss using the reconstructive error
3. Define a training step for the autoencoder using tf.GradientTape()
4. Train the autoencoder on the MNIST dataset
Introduction
This notebook demonstrates how to build and train a convolutional autoencoder.
Autoencoders consist of two models: an encoder and a decoder.
<img src="../assets/autoencoder2.png" width="600">
In this notebook we'll build an autoencoder to recreate MNIST digits. This notebook demonstrates this process on the MNIST dataset. The following animation shows a series of images produced by the generator as it was trained for 100 epochs. The images increasingly resemble hand written digits as the autoencoder learns to reconstruct the original images.
<img src="../assets/autoencoder.gif">
Import TensorFlow and other libraries
End of explanation
"""


np.random.seed(1)
tf.random.set_seed(1)
BATCH_SIZE = 128
BUFFER_SIZE = 60000
EPOCHS = 60
LR = 1e-2
EMBED_DIM = 64  # intermediate_dim

"""
Explanation: Next, we'll define some of the environment variables we'll use in this notebook. Note that we are setting the EMBED_DIM to be 64. This is the dimension of the latent space for our autoencoder.
End of explanation
"""


(train_images, _), (test_images, _) = tf.keras.datasets.mnist.load_data()

train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype(
    "float32"
)
train_images = (train_images - 127.5) / 127.5  # Normalize the images to [-1, 1]

"""
Explanation: Load and prepare the dataset
For this notebook, we will use the MNIST dataset to train the autoencoder. The encoder will map the handwritten digits into the latent space, to force a lower dimensional representation and the decoder will then map the encoding back.
End of explanation
"""


# Batch and shuffle the data
train_dataset = tf.data.Dataset.from_tensor_slices(train_images)
train_dataset = train_dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
train_dataset = train_dataset.prefetch(BATCH_SIZE * 4)

test_images = test_images.reshape(test_images.shape[0], 28, 28, 1).astype(
    "float32"
)
test_images = (test_images - 127.5) / 127.5  # Normalize the images to [-1, 1]

"""
Explanation: Next, we define our input pipeline using tf.data. The pipeline below reads in train_images as tensor slices and then shuffles and batches the examples for training.
End of explanation
"""


# TODO 1.
def make_encoder(embed_dim):
    model = tf.keras.Sequential(name="encoder")
    model.add(
        layers.Conv2D(
            64, (5, 5), strides=(2, 2), padding="same", input_shape=[28, 28, 1]
        )
    )
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Conv2D(128, (5, 5), strides=(2, 2), padding="same"))
    model.add(layers.LeakyReLU())
    model.add(layers.Dropout(0.3))

    model.add(layers.Flatten())
    model.add(layers.Dense(embed_dim))
    assert model.output_shape == (None, embed_dim)

    return model

"""
Explanation: Create the encoder and decoder models
Both our encoder and decoder models will be defined using the Keras Sequential API.
The Encoder
The encoder uses tf.keras.layers.Conv2D layers to map the image into a lower-dimensional latent space. We will start with an image of size 28x28x1 and then use convolution layers to map into a final Dense layer.
Exercise. Complete the code below to create the CNN-based encoder model. Your model should have input_shape to be 28x28x1 and end with a final Dense layer the size of embed_dim.
End of explanation
"""


# TODO 1.
def make_decoder(embed_dim):
    model = tf.keras.Sequential(name="decoder")
    model.add(layers.Dense(embed_dim, use_bias=False, input_shape=(embed_dim,)))
    model.add(layers.Dense(6272, use_bias=False, input_shape=(embed_dim,)))
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(layers.Reshape((7, 7, 128)))

    model.add(
        layers.Conv2DTranspose(
            128, (5, 5), strides=(1, 1), padding="same", use_bias=False
        )
    )
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(
        layers.Conv2DTranspose(
            64, (5, 5), strides=(2, 2), padding="same", use_bias=False
        )
    )
    model.add(layers.BatchNormalization())
    model.add(layers.LeakyReLU())

    model.add(
        layers.Conv2DTranspose(
            1,
            (5, 5),
            strides=(2, 2),
            padding="same",
            use_bias=False,
            activation="tanh",
        )
    )
    assert model.output_shape == (None, 28, 28, 1)

    return model

"""
Explanation: The Decoder
The decoder uses tf.keras.layers.Conv2DTranspose (upsampling) layers to produce an image from the latent space. We will start with a Dense layer with the same input shape as embed_dim, then upsample several times until you reach the desired image size of 28x28x1.
Exercise. Complete the code below to create the decoder model. Start with a Dense layer that takes as input a tensor of size embed_dim. Use tf.keras.layers.Conv2DTranspose over multiple layers to upsample so that the final layer has shape 28x28x1 (the shape of our original MNIST digits).
Hint: Experiment with using BatchNormalization or different activation functions like LeakyReLU.
End of explanation
"""


ae_model = tf.keras.models.Sequential(
    [make_encoder(EMBED_DIM), make_decoder(EMBED_DIM)]
)

"""
Explanation: Finally, we stitch the encoder and decoder models together to create our autoencoder.
End of explanation
"""


ae_model.summary()

make_encoder(EMBED_DIM).summary()

make_decoder(EMBED_DIM).summary()

"""
Explanation: Using .summary() we can have a high-level summary of the full autoencoder model as well as the individual encoder and decoder. Note how the shapes of the tensors mirror each other as data is passed through the encoder and then the decoder.
End of explanation
"""


# TODO 2.
def loss(model, original):
    reconstruction_error = tf.reduce_mean(
        tf.square(tf.subtract(model(original), original))
    )
    return reconstruction_error

"""
Explanation: Next, we define the loss for our autoencoder model. The loss we will use is the reconstruction error. This loss is similar to the MSE loss we've commonly use for regression. Here we are applying this error pixel-wise to compare the original MNIST image and the image reconstructed from the decoder.
End of explanation
"""


optimizer = tf.keras.optimizers.SGD(lr=LR)

"""
Explanation: Optimizer for the autoencoder
Next we define the optimizer for model, specifying the learning rate.
End of explanation
"""


checkpoint_dir = "./ae_training_checkpoints"
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt")
checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=ae_model)

"""
Explanation: Save checkpoints
This notebook also demonstrates how to save and restore models, which can be helpful in case a long running training task is interrupted.
End of explanation
"""


# TODO 3.
@tf.function
def train_step(images):
    with tf.GradientTape() as tape:
        ae_gradients = tape.gradient(
            loss(ae_model, images), ae_model.trainable_variables
        )
    gradient_variables = zip(ae_gradients, ae_model.trainable_variables)

    optimizer.apply_gradients(gradient_variables)

"""
Explanation: Define the training loop
Next, we define the training loop for training our autoencoder. The train step will use tf.GradientTape() to keep track of gradient steps through training.
Exercise.
Complete the code below to define the training loop for our autoencoder. Notice the use of tf.function below. This annotation causes the function train_step to be "compiled". The train_step function takes as input a batch of images and passes them through the ae_model. The gradient is then computed on the loss against the ae_model output and the original image. In the code below, you should
- define ae_gradients. This is the gradient of the autoencoder loss with respect to the variables of the ae_model.
- create the gradient_variables by assigning each ae_gradient computed above to it's respective training variable.
- apply the gradient step using the optimizer
End of explanation
"""


def train(dataset, epochs):
    for epoch in range(epochs):
        start = time.time()

        for image_batch in dataset:
            train_step(image_batch)

        # Produce images for the GIF as we go
        display.clear_output(wait=True)
        generate_and_save_images(ae_model, epoch + 1, test_images[:16, :, :, :])

        # Save the model every 5 epochs
        if (epoch + 1) % 5 == 0:
            checkpoint.save(file_prefix=checkpoint_prefix)

        print(f"Time for epoch {epoch + 1} is {time.time() - start} sec")

    # Generate after the final epoch
    display.clear_output(wait=True)
    generate_and_save_images(ae_model, epochs, test_images[:16, :, :, :])

"""
Explanation: We use the train_step function above to define training of our autoencoder. Note here, the train function takes as argument the tf.data dataset and the number of epochs for training.
End of explanation
"""


def generate_and_save_images(model, epoch, test_input):
    # Notice `training` is set to False.
    # This is so all layers run in inference mode (batchnorm).
    predictions = model(test_input, training=False)

    fig = plt.figure(figsize=(4, 4))

    for i in range(predictions.shape[0]):
        plt.subplot(4, 4, i + 1)
        pixels = predictions[i, :, :] * 127.5 + 127.5
        pixels = np.array(pixels, dtype="float")
        pixels = pixels.reshape((28, 28))
        plt.imshow(pixels, cmap="gray")
        plt.axis("off")

    plt.savefig(f"image_at_epoch_{epoch:04d}.png")
    plt.show()

"""
Explanation: Generate and save images.
We'll use a small helper function to generate images and save them.
End of explanation
"""


generate_and_save_images(ae_model, 4, test_images[:16, :, :, :])

"""
Explanation: Let's see how our model performs before any training. We'll take as input the first 16 digits of the MNIST test set. Right now they just look like random noise.
End of explanation
"""


# TODO 4.
train(train_dataset, EPOCHS)

"""
Explanation: Train the model
Call the train() method defined above to train the autoencoder model.
We'll print the resulting images as training progresses. At the beginning of the training, the decoded images look like random noise. As training progresses, the model outputs will look increasingly better. After about 50 epochs, they resemble MNIST digits. This may take about one or two minutes / epoch
End of explanation
"""


# Display a single image using the epoch number
def display_image(epoch_no):
    return PIL.Image.open(f"./ae_images/image_at_epoch_{epoch_no:04d}.png")

display_image(EPOCHS)

anim_file = "autoencoder.gif"

with imageio.get_writer(anim_file, mode="I") as writer:
    filenames = glob.glob("./ae_images/image*.png")
    filenames = sorted(filenames)
    last = -1
    for i, filename in enumerate(filenames):
        frame = 2 * (i**0.5)
        if round(frame) > round(last):
            last = frame
        else:
            continue
        image = imageio.imread(filename)
        writer.append_data(image)
    image = imageio.imread(filename)
    writer.append_data(image)

import IPython

if IPython.version_info > (6, 2, 0, ""):
    display.Image(filename=anim_file)

"""
Explanation: Create a GIF
Lastly, we'll create a gif that shows the progression of our produced images through training.
End of explanation
"""


