"""
MNIST classification with a convolutional neural network (CNN).

The 28x28 images are kept in their 2-D form so the convolution layers can
exploit spatial structure (neighbouring pixels, local strokes/edges).

Built to be comparable to mnist_dense.py: the same number of weight layers
(4) and a similar parameter count (~24k here vs ~27k for the dense net),
trained for the same number of epochs with the same optimizer.

Architecture:
    Input 28x28x1
    -> Conv2D(16, 3x3, relu) -> MaxPool(2x2)      # 26x26x16 -> 13x13x16
    -> Conv2D(32, 3x3, relu) -> MaxPool(2x2)      # 11x11x32 -> 5x5x32
    -> Flatten (800) -> Dense(24, relu) -> Dense(10, softmax)
"""
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

EPOCHS = 5
BATCH_SIZE = 128

tf.random.set_seed(0)

# Load data and scale pixel values to [0, 1]
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()
x_train = x_train.astype("float32") / 255.0
x_test = x_test.astype("float32") / 255.0

# Keep the 2-D image layout and add a single channel dimension
x_train = x_train.reshape(-1, 28, 28, 1)
x_test = x_test.reshape(-1, 28, 28, 1)
print("Training input shape:", x_train.shape)   # (60000, 28, 28, 1)
print("Test input shape:    ", x_test.shape)    # (10000, 28, 28, 1)

model = keras.Sequential([
    keras.Input(shape=(28, 28, 1)),
    layers.Conv2D(16, (3, 3), activation="relu"),
    layers.MaxPooling2D((2, 2)),
    layers.Conv2D(32, (3, 3), activation="relu"),
    layers.MaxPooling2D((2, 2)),
    layers.Flatten(),
    layers.Dense(24, activation="relu"),
    layers.Dense(10, activation="softmax"),
], name="cnn_mnist")

model.summary()

model.compile(optimizer="adam",
              loss="sparse_categorical_crossentropy",
              metrics=["accuracy"])

model.fit(x_train, y_train,
          epochs=EPOCHS,
          batch_size=BATCH_SIZE,
          validation_split=0.1,
          verbose=2)

test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0)
print(f"\nCNN parameters: {model.count_params():,}")
print(f"CNN test accuracy: {test_acc:.4f}")