"""
LeNet-5 example
"""
import gzip
from time import time
import numpy as np
from requests import get
import matplotlib.pyplot as plt
import pandas as pd
from sklearn.model_selection import train_test_split
import keras
from keras.models import Sequential
import keras.layers as layers
from keras.preprocessing.image import ImageDataGenerator
from keras.utils.np_utils import to_categorical
from keras.callbacks import TensorBoard
from keras.models import load_model

EPOCHS = 10
BATCH_SIZE = 128

train = {}
test = {}
validation = {}

def download_file(url, file_name):
    """
    Download files and stores them locally.
    """
    with open(file_name, "wb") as file:
        response = get(url)
        file.write(response.content)

def download_files():
    """
    Download all data and label files.
    """
    # train-images-idx3-ubyte.gz: training set images (9912422 bytes)
    download_file('http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz',
                  'train-images-idx3-ubyte.gz')
    # train-labels-idx1-ubyte.gz: training set labels (28881 bytes)
    download_file('http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz',
                  'train-labels-idx1-ubyte.gz')
    # t10k-images-idx3-ubyte.gz: test set images (1648877 bytes)
    download_file('http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz',
                  't10k-images-idx3-ubyte.gz')
    # t10k-labels-idx1-ubyte.gz: test set labels (4542 bytes)
    download_file('http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz',
                  't10k-labels-idx1-ubyte.gz')

def read_mnist(images_path: str, labels_path: str):
    """
    Read data and labels of the MNIST dataset.
    """
    with gzip.open(labels_path, 'rb') as labels_file:
        labels = np.frombuffer(labels_file.read(), dtype=np.uint8, offset=8)

    with gzip.open(images_path,'rb') as images_file:
        length = len(labels)
        # Load flat 28x28 px images (784 px), and convert them to 28x28 px
        features = np.frombuffer(images_file.read(), dtype=np.uint8, offset=16) \
                        .reshape(length, 784)                                   \
                        .reshape(length, 28, 28, 1)

    return features, labels

def display_image(dataset, position):
    """
    Display image at position of the given dataset.
    """
    image = dataset['features'][position].squeeze()
    plt.title('Example %d. Label: %d' % (position, dataset['labels'][position]))
    plt.imshow(image, cmap=plt.get_cmap('gray_r'))
    plt.show()

def create_lenet5():
    """
    Creates the LeNet-5 model.
    """
    model = Sequential()

    # C1: (None,32,32,1) -> (None,28,28,6).
    model.add(layers.Conv2D(6, kernel_size=(5, 5), strides=(1, 1), activation='tanh',
                            input_shape=(32,32,1), padding='valid'))

    # S2: (None,28,28,6) -> (None,14,14,6).
    model.add(layers.AveragePooling2D(pool_size=(2, 2), strides=(2, 2), padding='valid'))

    # C3: (None,14,14,6) -> (None,10,10,16).
    model.add(layers.Conv2D(16, kernel_size=(5, 5), strides=(1, 1), activation='tanh',
                            padding='valid'))

    # S4: (None,10,10,16) -> (None,5,5,16).
    model.add(layers.AveragePooling2D(pool_size=(2, 2), strides=(2, 2), padding='valid'))

    # Flatten: (None,5,5,16) -> (None, 400).
    model.add(layers.Flatten())

    # C5: (None, 400) -> (None,120).
    model.add(layers.Dense(120, activation='tanh'))

    # F6: (None,120) -> (None,84).
    model.add(layers.Dense(84, activation='tanh'))

    # Output: (None,84) -> (None,10).
    model.add(layers.Dense(10, activation='softmax'))

    # Compile the model
    # model.compile(loss='sparse_categorical_crossentropy', optimizer='sgd', metrics=['accuracy'])
    model.compile(loss=keras.losses.categorical_crossentropy,
                  optimizer=keras.optimizers.Adam(), metrics=['accuracy'])

    print(model.summary())

    return model

def train_lenet5(model):
    """
    Trains the LeNet-5 model.
    """
    x_train, y_train = train['features'], to_categorical(train['labels'])
    x_validation, y_validation = validation['features'], to_categorical(validation['labels'])

    train_generator = ImageDataGenerator().flow(x_train, y_train, batch_size=BATCH_SIZE)
    validation_generator = ImageDataGenerator().flow(x_validation, y_validation,
                                                     batch_size=BATCH_SIZE)

    steps_per_epoch = x_train.shape[0] // BATCH_SIZE
    validation_steps = x_validation.shape[0] // BATCH_SIZE

    print('Number of training images:', train['features'].shape[0])
    print('Number of validation images:', validation['features'].shape[0])

    tensorboard = TensorBoard(log_dir="logs/{}".format(time()))


    model.fit_generator(train_generator, steps_per_epoch=steps_per_epoch, epochs=EPOCHS,
                    validation_data=validation_generator, validation_steps=validation_steps,
                    shuffle=True, callbacks=[tensorboard])

def print_netout(i, netout, maximum_value):
    """
    Prints the netout of LeNet-5.
    """
    print("[%.4d] : [%.0f %.0f %.0f %.0f %.0f %.0f %.0f %.0f %.0f %.0f] = >>> %d <<<"
          % (i,
             netout[0], netout[1], netout[2], netout[3], netout[4],
             netout[5], netout[6], netout[7], netout[8], netout[9],
             maximum_value))

def main():
    """
    Defined starting point of source code.
    """

    # Step 1:
    # Download the MNIST dataset with consist of labeled handwritten images (28x28 px).
    download_files()

    # Step 2:
    # Read MNIST dataset (training and testing)
    train['features'], train['labels'] = read_mnist('train-images-idx3-ubyte.gz',
                                                    'train-labels-idx1-ubyte.gz')
    test['features'], test['labels'] = read_mnist('t10k-images-idx3-ubyte.gz',
                                                  't10k-labels-idx1-ubyte.gz')

    # Step 3:
    # Explore the dataset
    print('Number of training images:', train['features'].shape[0])
    print('Number of test images:', test['features'].shape[0])

    # Step 4:
    # Display three images of training dataset
    for i in range(3):
        display_image(train, i)

    # Step 5:
    # Plot information about the training data
    train_labels_count = np.unique(train['labels'], return_counts=True)
    dataframe_train_labels = pd.DataFrame({'Label':train_labels_count[0],
                                           'Count':train_labels_count[1]})
    print(dataframe_train_labels)

    # Step 5:
    # Split training data into training and validation
    train['features'], validation['features'], train['labels'], validation['labels'] \
        = train_test_split(train['features'], train['labels'], test_size=0.2, random_state=0)

    print('Number of training images:', train['features'].shape[0])
    print('Number of validation images:', validation['features'].shape[0])

    # Step 6:
    # Prepare our input features.
    # The LeNet-5 architecture accepts 32x32 pixel images as input, but MNIST data is 28x28 pixels.
    # We simply pad the imges with zeros to overcome that.
    train['features']      = np.pad(train['features'], ((0,0),(2,2),(2,2),(0,0)), 'constant')
    test['features']       = np.pad(test['features'], ((0,0),(2,2),(2,2),(0,0)), 'constant')
    validation['features'] = np.pad(validation['features'], ((0,0),(2,2),(2,2),(0,0)), 'constant')

    print("Updated Image Shape: {}".format(train['features'][0].shape))

    # Step 7:
    # Create and compile the LeNet-5 model
    model = create_lenet5()

    # Step 8:
    # Train, save and load the LeNet-5 model
    train_lenet5(model)
    model.save('lenet5.h5')
    model = load_model('lenet5.h5')

    # Step 9:
    # PEvaluate with test data and print results
    score = model.evaluate(test['features'], to_categorical(test['labels']))
    print('Test loss:', score[0])
    print('Test accuracy:', score[1])

    # Step 10:
    # Execute model and predict results
    netouts = model.predict(test['features'])
    maximum_values = netouts.argmax(axis=1)

    # Step 11:
    # Investigate ten predictions (randomly selected)
    for i in np.random.choice(np.arange(0, len(test['labels'])), size=(10,)):
        print_netout(i, netouts[i], maximum_values[i])
        display_image(test, i)

if __name__ == "__main__":
    main()