#!/usr/bin/env python
# coding: utf-8

# In[7]:


import argparse
class Deeplearn:
    def methd(self,epochh):
        from keras.preprocessing.image import ImageDataGenerator
        from keras.models import Sequential
        from keras.layers import Dense, Dropout, Activation, Flatten
        from keras.layers import Conv2D, MaxPooling2D, ZeroPadding2D
        from keras.layers.normalization import BatchNormalization
        from keras.regularizers import l2
        from keras.datasets import mnist
        from keras.utils import np_utils
        import keras
        import struct
        import numpy as np

        def read_idx(filename):
            """Credit: https://gist.github.com/tylerneylon"""
            with open(filename, 'rb') as f:
                zero, data_type, dims = struct.unpack('>HBB', f.read(4))
                shape = tuple(struct.unpack('>I', f.read(4))[0] for d in range(dims))
                return np.frombuffer(f.read(), dtype=np.uint8).reshape(shape)
        # loads the MNIST dataset
        (x_train, y_train), (x_test, y_test)  = mnist.load_data()

        # Lets store the number of rows and columns
        img_rows = x_train[0].shape[0]
        img_cols = x_train[1].shape[0]

        # Getting our date in the right 'shape' needed for Keras
        # We need to add a 4th dimenion to our date thereby changing our
        # Our original image shape of (60000,28,28) to (60000,28,28,1)
        x_train = x_train.reshape(x_train.shape[0], img_rows, img_cols, 1)
        x_test = x_test.reshape(x_test.shape[0], img_rows, img_cols, 1)
        # store the shape of a single image 
        input_shape = (img_rows, img_cols, 1)

        # change our image type to float32 data type
        x_train = x_train.astype('float32')
        x_test = x_test.astype('float32')

        # Normalize our data by changing the range from (0 to 255) to (0 to 1)
        x_train /= 255
        x_test /= 255
        # Now we one hot encode outputs
        y_train = np_utils.to_categorical(y_train)
        y_test = np_utils.to_categorical(y_test)

        num_classes = y_test.shape[1]
        num_pixels = x_train.shape[1] * x_train.shape[2]
        # create model
        model = Sequential()
        # 2 sets of CRP (Convolution, RELU, Pooling)
        model.add(Conv2D(20, (5, 5),
                         padding = "same", 
                         input_shape = input_shape))
        model.add(Activation("relu"))
        model.add(MaxPooling2D(pool_size = (2, 2), strides = (2, 2)))

        model.add(Conv2D(50, (5, 5),
                         padding = "same"))
        model.add(Activation("relu"))
        model.add(MaxPooling2D(pool_size = (2, 2), strides = (2, 2)))

        # Fully connected layers (w/ RELU)
        model.add(Flatten())
        model.add(Dense(500))
        model.add(Activation("relu"))
        # Softmax (for classification)
        model.add(Dense(num_classes))
        model.add(Activation("softmax"))

        model.compile(loss = 'categorical_crossentropy',
                      optimizer = keras.optimizers.Adadelta(),
                      metrics = ['accuracy'])
        # Training Parameters
        batch_size = 128
        epochs = epochh

        history = model.fit(x_train, y_train,
                  batch_size=batch_size,
                  epochs=epochs,
                  validation_data=(x_test, y_test),
                  shuffle=True)

        model.save("mnist_LeNet.h5")
        # Evaluate the performance of our trained model
        scores = model.evaluate(x_test, y_test, verbose=1)
        f = open("score","w")       
        g = open("Actual_score.txt","w")
        f.write(str(int(scores[1]*100)))
        g.write(str(scores[1]*100))
        g.write("Epoch: "+str(epochs))
        g.close()
        f.close()
if __name__== '__main__':
    arg = argparse.ArgumentParser(description='Minist_LeNet')

    arg.add_argument('-e', '--epoch', type = int, default = 1, help = 'Epochs(Default=1)', choices= range(1,50))

    args = arg.parse_args()
    deepLearn = Deeplearn()
    deepLearn.methd(args.epoch)


