import tensorflow as tf
print("TensorFlow version: ",tf.version.VERSION)

"""
Explanation: Introduction to Logistic Regression Using TF 2.0
Learning Objectives

Build a neural network that classifies images.
Train this neural network.
Evaluate the accuracy of the model.

Introduction
This short introduction uses Keras, a high-level API to build and train models in TensoFlow.  In this lab, you Load and prepare the MNIST dataset, convert the samples from integers to floating-point numbers, build and train a neural network that classifies images and the evaluate then accuracy of the model.
Each learning objective will correspond to a #TODO in the student lab notebook -- try to complete that notebook first before reviewing this solution notebook.
Load necessary libraries
We will start by importing the necessary libraries for this lab.
End of explanation
"""


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

"""
Explanation: Load and prepare the MNIST dataset. Convert the samples from integers to floating-point numbers:
End of explanation
"""


model = # TODO 1 -- Your code here.

"""
Explanation: Lab Task 1: Build the tf.keras.Sequential model by stacking layers. Choose an optimizer and loss function for training:
End of explanation
"""


predictions = model(x_train[:1]).numpy()
predictions

"""
Explanation: For each example the model returns a vector of "logits" or "log-odds" scores, one for each class.
End of explanation
"""


tf.nn.softmax(predictions).numpy()

"""
Explanation: The tf.nn.softmax function converts these logits to "probabilities" for each class:
End of explanation
"""


loss_fn = # TODO 2 -- Your code here.

"""
Explanation: Note: It is possible to bake this tf.nn.softmax in as the activation function for the last layer of the network. While this can make the model output more directly interpretable, this approach is discouraged as it's impossible to
provide an exact and numerically stable loss calculation for all models when using a softmax output. 
The losses.SparseCategoricalCrossentropy loss takes a vector of logits and a True index and returns a scalar loss for each example.
Lab Task #2: Usage of losses.SparseCategoricalCrossentropy with logits vectors and a True index.
End of explanation
"""


loss_fn(y_train[:1], predictions).numpy()

model.compile(optimizer='adam',
              loss=loss_fn,
              metrics=['accuracy'])

"""
Explanation: This loss is equal to the negative log probability of the true class:
It is zero if the model is sure of the correct class.
This untrained model gives probabilities close to random (1/10 for each class), so the initial loss should be close to -tf.log(1/10) ~= 2.3.
End of explanation
"""


model.fit(x_train, y_train, epochs=5)

"""
Explanation: The Model.fit method adjusts the model parameters to minimize the loss:
End of explanation
"""


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

"""
Explanation: The Model.evaluate method checks the models performance, usually on a "Validation-set" or "Test-set".
End of explanation
"""


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

probability_model(x_test[:5])

"""
Explanation: The image classifier is now trained to ~98% accuracy on this dataset. To learn more, read the TensorFlow tutorials.
If you want your model to return a probability, you can wrap the trained model, and attach the softmax to it:
End of explanation
"""


