import matplotlib.pyplot as plt
import tensorflow as tf

mnist = tf.keras.datasets.mnist

(images_train, labels_train), (images_test, labels_test) = mnist.load_data()

plt.figure()
plt.imshow(images_train[0])
plt.colorbar()
plt.grid(False)
plt.show()

images_train, images_test = images_train / 255.0, images_test / 255.0

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)
])

model.compile(optimizer='adam', loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
model.fit(images_train, labels_train, epochs=5, verbose=0, use_multiprocessing=True)
model.evaluate(images_test, labels_test, verbose=1)

probability_model = tf.keras.Sequential([model, tf.keras.layers.Softmax()])
for i, n in enumerate(labels_test[0:10]):
    p: tf.Tensor = probability_model(images_test[i: i + 1])
    pv = p.numpy().flatten().tolist()
    print(pv.index(max(pv)), n)

print("doen")