import numpy as np
from tensorflow import keras
import matplotlib.pyplot as plt

dataset = keras.datasets.mnist

(x_train, train_labels), (x_test, test_labels) = dataset.load_data()
number_of_classes = np.max([np.max(train_labels), np.max(test_labels)]) + 1

# Scale images to the [0, 1] range
x_train = x_train.astype("float32") / 255
x_test = x_test.astype("float32") / 255

# Make sure images have shape (28, 28, 1)
x_train = np.expand_dims(x_train, -1)
x_test = np.expand_dims(x_test, -1)

print("x_train shape:", x_train.shape)
print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")
input_shape = x_train[0].shape

#%%

model = keras.Sequential(    [
        keras.Input(shape=input_shape),
        keras.layers.Conv2D(32, kernel_size=(3, 3), activation="relu"),
        keras.layers.MaxPooling2D(pool_size=(2, 2)),
        keras.layers.Conv2D(64, kernel_size=(3, 3), activation="relu"),
        keras.layers.MaxPooling2D(pool_size=(2, 2)),
        keras.layers.Flatten(),
        keras.layers.Dropout(0.5),
        keras.layers.Dense(number_of_classes, activation="softmax"),
    ])

model.summary()

batch_size = 128
epochs = 15

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


model.fit(x_train, train_labels, epochs=epochs, batch_size=batch_size)
test_loss, test_acc = model.evaluate(x_test,  test_labels, verbose=2)
print('\nTest accuracy: %.2f%%' % (test_acc * 100))

#%%

predictions = model.predict(x_test)

classes = np.argmax(predictions, 1)
test_accuracy = np.mean(classes == test_labels)
print('\nTest accuracy: %.2f' % (test_accuracy * 100))