"""
MNIST classification with a fully-connected (dense) neural network.

Each 28x28 image is flattened into a linear list of 784 inputs, which is
fed through three small hidden layers and a 10-way softmax output.

Architecture (~27k trainable parameters, 4 weight layers):
    Input 784 -> Dense(32, relu) -> Dense(32, relu) -> 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


# Flatten each 28x28 image into a vector of 784 inputs
x_train = x_train.reshape(-1, 28 * 28)
x_test = x_test.reshape(-1, 28 * 28)
print("Training input shape:", x_train.shape)   # (60000, 784)
print("Test input shape:    ", x_test.shape)    # (10000, 784)

#################################

model = keras.Sequential([
    keras.Input(shape=(784,)),
    layers.Dense(32, activation="relu"),
    layers.Dense(32, activation="relu"),
    layers.Dense(24, activation="relu"),
    layers.Dense(10, activation="softmax"),
], name="dense_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"\nDense NN parameters: {model.count_params():,}")
print(f"Dense NN test accuracy: {test_acc:.4f}")