# -*- coding: utf-8 -*-
"""
Created on Wed Aug  5 00:43:27 2020

@author: jisuk
"""

# %%  import importance module

from __future__ import absolute_import, division, print_function

import tensorflow as tf
from tensorflow.keras import Model, layers
import numpy as np

# %% MNIST dataset parameters
num_classes = 10  # total classes (0-9 digits)

# training parameters
learning_rate = 0.001
training_steps = 500
batch_size = 128
display_step = 10

# network parameters
conv1_filters = 32  # number of filters for 1st conv layer
conv2_filters = 64  # number of filters for 2nd conv layer
fc1_units = 1024  # number of neurons for 1st fully-connected layer

# %% prepare NIMST data
#from tensorflow.keras.datasets import mnist
from tensorflow.keras.datasets import mnist
(x_train, y_train), (x_test, y_test) = mnist.load_data()
# convert to float32
x_train, x_test = np.array(x_train, np.float32), np.array(x_test, np.float32)
# normalize images value from [0 255] to [0 1]
x_train, x_test = x_train / 255.0, x_test/255.0

# %% use tf.data API to shuffle and batch data.
train_data = tf.data.Dataset.from_tensor_slices((x_train, y_train))
train_data = train_data.repeat().shuffle(5000).batch(batch_size).prefetch(1)

# %% Create TF Model


class ConvNet(Model):
    # set layers
    def __init__(self):
        super(ConvNet, self).__init__()
        # convolutional layer with 32 filters and a kernel size o f 5
        self.conv1 = layers.Conv2D(32, kernel_size=5, activation=tf.nn.relu)
        # max pooling (down-sampling) with kernel size of 2 and strides of 2.
        self.maxpool1 = layers.MaxPool2D(2, strides=2)

        # convolution layer with 64 filters and a kernel size of 3
        self.conv2 = layers.Conv2D(64, kernel_size=3, activation=tf.nn.relu)
        # max pooling (down-sampling) with kernel size of 2 and strides of 2.
        self.maxpool2 = layers.MaxPool2D(2, strides=2)

        # Flatten thedata to a 1-D vector for the fully connected layer.
        self.flatten = layers.Flatten()

        # fully connected layer.
        self.fc1 = layers.Dense(1024)
        # apply dropout (tf is_training is False, dropout is not applied)
        self.dropout = layers.Dropout(rate=0.5)

        # output layer class prediction
        self.out = layers.Dense(num_classes)

    # set forward pass
    def call(self, x, is_training=False):
        x = tf.reshape(x, [-1, 28, 28, 1])
        x = self.conv1(x)
        x = self.maxpool1(x)
        x = self.conv2(x)
        x = self.maxpool2(x)
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.dropout(x, training=is_training)
        x = self.out(x)
        if not is_training:
            # tf cross entropy expect logits without softmax so only
            # apply softmax when not training
            x = tf.nn.softmax(x)
        return x


# build neural network model.
conv_net = ConvNet()

# %% cross entrolpy loss
# note that this will apply softmax to the logits


def cross_entropy_loss(x, y):
    # convert labels to int 64 for tf cross-entropy function.
    y = tf.cast(y, tf.int64)
    # apply softmax to logits and compute cross-entropy
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=y, logits=x)
    # average loss across the batch.
    return tf.reduce_mean(loss)

# accuracy metric


def accuracy(y_pred, y_true):
    # predicted class is the index fo highest score in prediction vector
    correct_prediction = tf. equal(
        tf.argmax(y_pred, 1), tf.cast(y_true, tf.int64))
    return tf.reduce_mean(tf.cast(correct_prediction, tf.float32), axis=-1)


# stochasitc gradient descent optimizer
optimizer = tf.optimizers.Adam(learning_rate)

# %% optimization result


def run_optimization(x, y):
    # wrap computation inside a gradientTape for automatic differentiation
    with tf.GradientTape() as g:
        # forward pass
        pred = conv_net(x, is_training=True)
        # compute loss
        loss = cross_entropy_loss(pred, y)

    # variables to update, i.e. trainable variables
    trainable_variables = conv_net.trainable_variables

    # compute gradients
    gradients = g.gradient(loss, trainable_variables)

    # update W and b following gradients
    optimizer.apply_gradients(zip(gradients, trainable_variables))


# %% run training for the given number of steps
for step, (batch_x, batch_y) in enumerate(train_data.take(training_steps), 1):
    # run the optimization to update W and b values.
    run_optimization(batch_x, batch_y)

    if step % display_step == 0:
        pred = conv_net(batch_x)
        loss = cross_entropy_loss(pred, batch_y)
        acc = accuracy(pred, batch_y)
        print("step: %i, loss: %f, accuracy: %f" % (step, loss, acc))

# %% test model on validation set
pred = conv_net(x_test)
print("test Accuracy: %f" % accuracy(pred, y_test))

# %% visualize predictions
#import matplotlib.pyplot as plt
import matplotlib.pyplot as plt
# %% predict 5 images from validation set
n_images = 5
test_images = x_test[:n_images]
predictions = conv_net(test_images)

# display image and model prediction
for i in range(n_images):
    plt.imshow(np.reshape(test_images[i], [28, 28]), cmap='gray')
    plt.show()
    print("Model prediction: %i" % np.argmax(predictions.numpy()[i]))
