# -*- coding: utf-8 -*-
"""
LeNet - Deep Learning Neural Network
"""

import tensorflow as tf
import numpy as np
from tensorflow.keras.layers import Conv2D, AveragePooling2D, Dense, Flatten
from tensorflow.keras.models import Sequential
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical 
import matplotlib.pyplot as plt

# Load the MNIST dataset using the Keras library:
(x_train , y_train), (x_test, y_test) = mnist.load_data()

# Split dataset into training and testing, padding and normalizing images:
x_train = np.pad(x_train,((0,0),(2,2),(2,2)))   # padding
x_test  = np.pad(x_test,((0,0),(2,2),(2,2)))
x_train = x_train/255.0                         # Normalizing the values
x_test  = x_test/255.0
y_train = to_categorical(y_train,10)            # converting labels into one-hot-encoded vectors
y_test  = to_categorical(y_test,10)

# Explore Dataset:
plt.imshow(x_train[0], cmap='gray')
plt.show()

# Expand dimensions of dataset, because LeNet expects images of size 32x32x1 instead of 32x32:
x_train = np.expand_dims(x_train, 3)
x_test  = np.expand_dims(x_test, 3)

# LeNet Model:
lenet = Sequential(name="LeNet-5")
lenet.add(Conv2D(6,(5,5),strides=(1,1), activation='tanh', input_shape=(32,32,1), name='C1')) # C1
lenet.add(AveragePooling2D(name='S2')) # S2
lenet.add(Conv2D(16,(5,5),strides=(1,1), activation='tanh', name='C3')) # C3
lenet.add(AveragePooling2D(name='S4')) # S4
lenet.add(Flatten())
lenet.add(Dense(120,activation='tanh', name='FC5'))      # FC5
lenet.add(Dense(84,activation='tanh', name='FC6'))       # FC6
lenet.add(Dense(10,activation='softmax', name='Output')) # FC7

# Compile model:
lenet.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# Model preview:
lenet.summary()

# Train model:
lenet.fit(x_train, y_train, epochs=10, batch_size=32)

# Evaluate model's performance by passing the test dataset:
_, acc = lenet.evaluate(x_test, y_test)

"""We aquire 98.58% accuracy."""