from __future__ import absolute_import, division, print_function, unicode_literals

import os
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Conv2D, MaxPool2D
from tensorflow.keras.layers import Activation, Flatten, Dropout, Dense
from tensorflow.keras.datasets import mnist
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.callbacks import TensorBoard

def lenet(input_shape, num_classes):
    model = Sequential()
    # extract image features by convolution and max pooling layers
    model.add(Conv2D(20, kernel_size = 5, padding = "same", input_shape = input_shape, activation = "relu"))
    model.add(MaxPool2D(pool_size= (2,2)))
    model.add(Conv2D(50, kernel_size = 5, padding = "same", activation = "relu"))
    model.add(MaxPool2D(pool_size= (2,2)))
    #classfy the calss by fully-connected layers
    model.add(Flatten())
    model.add(Dense(500, activation= "relu"))
    model.add(Dense(num_classes))
    model.add(Activation("softmax"))
    return model

class MNISTDataset():
    def __init__(self):
        self.image_shape = (28, 28, 1) #image is 28x28x1 (grayscale)
        self.num_classes = 10
    
    def get_batch(self):
        (x_train, y_train), (x_test, y_test) = mnist.load_data()
        x_train, x_test = [self.preprocess(d) for d in [x_train, x_test]]
        y_train, y_test = [self.preprocess(d, label_data = True) for d in [y_train, y_test]]

        return x_train, y_train, x_test, y_test
    
    def preprocess(self, data, label_data = False):
        if label_data:
            #convert class vectors to binary class matrices
            data = to_categorical(data, self.num_classes)
        else:
            data = data.astype("float32")
            data /= 255 #convert the value to 0~1 scale
            shape = (data.shape[0],) + self.image_shape #add dataset length to top
            data = data.reshape(shape)

        return data

class Trainer():
    def __init__(self, model, loss, optimizer):
        self._target = model
        self._target.compile(loss = loss, optimizer = optimizer, metrics = ["accuracy"])
        self.verbose = 1
        self.log_dir = os.path.join(os.path.dirname(__file__), "logdir")
    
    def train(self, x_train, y_train, batch_size, epochs, validation_split):
        if os.path.exists(self.log_dir):
            import shutil
            shutil.rmtree(self.log_dir) #remove previous extension
        os.mkdir(self.log_dir)
        self._target.fit(
            x_train, y_train, 
            batch_size = batch_size, epochs = epochs, 
            validation_split = validation_split, 
            callbacks = [TensorBoard(log_dir = self.log_dir)],
            verbose = self.verbose
        )

dataset = MNISTDataset()

#make model
model = lenet(dataset.image_shape, dataset.num_classes)

#train the model
x_train, y_train, x_test, y_test = dataset.get_batch()
trainer = Trainer(model, loss = "categorical_crossentropy", optimizer = Adam())
trainer.train(x_train, y_train, batch_size = 128, epochs = 12, validation_split = 0.2)

#show_result
score = model.evaluate(x_test, y_test, verbose = 0)
print("Test loss:", score[0])
print("Test accuracy", score[1])