!pip install tensorflow
import tensorflow as tf
from tensorflow import keras
import numpy as np
import matplotlib.pyplot as plt

fashion_mnist = keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

train_images = train_images / 255
test_images = test_images / 255

print("Train data shape: ", train_images.shape)
print("Train labels shape: ", train_labels.shape)
print("Test data shape: ", test_images.shape)
print("Test labels shape: ", test_labels.shape)

class_names = ['T-shirt', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

plt.imshow(train_images[0], cmap = "Greys", interpolation = "nearest")

# BUILD THE NEURAL NETWORK

# SET UP THE LAYERS
model = keras.Sequential([keras.layers.Flatten(input_shape = (28,28)),
                          keras.layers.Dense(128, activation = "relu"),
                          keras.layers.Dense(10)])

''' Flatten layer -- takes the array, which is currently in the shape (28,28), and flattens it into one long array of shape (1, 784)
    First Dense layer -- contains 128 fully connected nodes
    Second Dense layer -- this layer returns a logit array of length 10
    Each dense layer contains a score that the image belongs to one of the 10 classes '''

model.compile(optimizer = "adam",
              loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits = True),
              metrics = ["accuracy"])

''' loss -- measures how accurate the model is during training. We want this function to be minimized
          SparseCategoricalCrossentropy -- calculates the crossentropy between the labels and the predictions,
                                           use this when you have more than two categories to predict
    optimizer -- determines how to model will be updated based on the loss function and the data it sees
    metrics -- attribute used to monitor the training and testing steps '''

# TRAIN THE MODEL
# 50 epochs may be too much, accuracy does not go up too much after 15-20 epochs

model.fit(train_images, train_labels, epochs = 20, batch_size = 200)

print(model.summary())

# FIND THE ACCURACY WHEN USING THE TESTING DATA
# COMPARE THIS TO TRAINING ACCURACY TO SEE IF OVERFITTING HAS OCCURRED

test_loss, test_accuracy = model.evaluate(test_images, test_labels, verbose = 2)
print("\nTest Accuracy = ", test_accuracy * 100)

probability_model = tf.keras.Sequential([model,
                                         tf.keras.layers.Softmax()])

''' This softmax layer takes the logits that the original model outputs and converts
    them into probabilities which are much easier to understand and interpret '''

predictions = probability_model.predict(test_images)

# argmax finds the index of the highest prediction value
# check with the actual label to see how this prediction was
print(np.argmax(predictions[0]))
print(test_labels[0])

def plot_image(i, predictions_array, true_label, img):
  true_label, img = true_label[i], img[i]
  plt.grid(False)
  plt.yticks([])
  plt.xticks([])

  plt.imshow(img, cmap = "Greys")

  predicted_label = np.argmax(predictions_array)
  if predicted_label == true_label:
    color = "green"
  else:
    color = "red"

  plt.xlabel("{} {:2.0f}% ({})".format(class_names[predicted_label], # 1. display the predicted class
                                     100 * np.max(predictions_array), # 2. display the confidence of the prediction
                                     class_names[true_label], # in parenthesis, display the actual class label
                                     color = color)) # if the prediction is correct, the color is green. Wrong predictions are red

def plot_value_array(i, predicitons_array, true_label):
  # assign the true label
  true_label = true_label[i]
  # get rid of the grid of the plot
  plt.grid(False)
  # set the ticks of the x-axis to be the number of classes we are predicting
  plt.xticks(range(10))
  # get rid of the y ticks
  plt.yticks([])
  # set up the bar plot, the x values are 0-9 and the y values are the values that appear in the prediction array
  thisplot = plt.bar(range(10), predicitons_array)
  # set the limit of the y axis to be between 0 and 1
  plt.ylim(0,1)
  # assign the predicted label to a name 
  predicted_label = np.argmax(predicitons_array)

  # set the colors of the two special cases in our bar plot
  thisplot[predicted_label].set_color("red")
  thisplot[true_label].set_color("green")

num_rows = 4
num_cols = 3
num_images = num_rows * num_cols
plt.figure(figsize = (2*2*num_cols, 2*num_rows))

for i in range(num_images):
  plt.subplot(num_rows, 2*num_cols, 2*i+1)
  plot_image(i, predictions[i], test_labels, test_images)
  plt.subplot(num_rows, 2*num_cols, 2*i+2)
  plot_value_array(i, predictions[i], test_labels)

plt.tight_layout()
plt.show()