# Library import
import tensorflow as tf
import pandas as pd
from tensorflow import keras
import matplotlib as plt

# Load MNIST dataset. Convert integer to float.
mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

# Build a tf.keras.model with layers.
model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation='softmax')
])

# Choose an optimizer and loss function to train the model
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# Train and evaluate the model  
early_stopping = keras.callbacks.EarlyStopping(
    patience=10,
    min_delta=0.001,
    restore_best_weights=True,
)

history = model.fit(
    x_train, y_train,
    validation_data=(x_test, y_test),
    batch_size=512,
    epochs=500,
    callbacks=[early_stopping],
    verbose=0, # hide the output
)

history_df = pd.DataFrame(history.history)
ax = history_df.loc[0:, ['loss']].plot()
ax.figure.savefig('./loss.png')
ax = history_df.loc[0:, ['accuracy']].plot()
ax.figure.savefig('./accuracy.png')

print(("Best Validation Loss: {:0.4f}" +\
      "\nBest Validation Accuracy: {:0.4f}")\
      .format(history_df['val_loss'].min(),
              history_df['val_accuracy'].max()))

model.evaluate(x_test,  y_test, verbose=2)