"""
MNIST classification with a larger, regularized CNN aimed at the highest
practical accuracy (no limit on parameters or layers).

Improvements over mnist_cnn.py:
  * Deeper network: five 3x3 conv layers in three blocks (a VGG-style
    design) with more filters (32 -> 64 -> 128).
  * Batch normalization after each conv layer for faster, more stable training.
  * Dropout to reduce overfitting.
  * Data augmentation (small random rotations, shifts and zooms), applied
    only during training, so the network sees many slightly different
    versions of each digit.
  * More epochs with a learning rate that is cut when validation loss
    stops improving, plus early stopping that keeps the best weights.

Architecture (~437k trainable parameters):
    Input 28x28x1 -> augmentation
    -> [Conv(32) -> BN -> Conv(32) -> BN -> MaxPool -> Dropout]   # 14x14x32
    -> [Conv(64) -> BN -> Conv(64) -> BN -> MaxPool -> Dropout]   # 7x7x64
    -> [Conv(128) -> BN -> MaxPool -> Dropout]                    # 3x3x128
    -> Flatten -> Dense(256) -> BN -> Dropout -> Dense(10, softmax)
"""
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

EPOCHS = 30
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)


def conv_bn(filters):
    """3x3 convolution followed by batch normalization and ReLU."""
    return [
        layers.Conv2D(filters, (3, 3), padding="same", use_bias=False),
        layers.BatchNormalization(),
        layers.Activation("relu"),
    ]


model = keras.Sequential([
    keras.Input(shape=(28, 28, 1)),

    # Data augmentation (only active during training)
    layers.RandomRotation(0.03),               # about +/- 10 degrees
    layers.RandomTranslation(0.1, 0.1),        # up to +/- 10% shift
    layers.RandomZoom(0.1),

    *conv_bn(32), *conv_bn(32),
    layers.MaxPooling2D((2, 2)),
    layers.Dropout(0.25),

    *conv_bn(64), *conv_bn(64),
    layers.MaxPooling2D((2, 2)),
    layers.Dropout(0.25),

    *conv_bn(128),
    layers.MaxPooling2D((2, 2)),
    layers.Dropout(0.25),

    layers.Flatten(),
    layers.Dense(256, use_bias=False),
    layers.BatchNormalization(),
    layers.Activation("relu"),
    layers.Dropout(0.4),
    layers.Dense(10, activation="softmax"),
], name="cnn_mnist_best")

model.summary()

model.compile(optimizer=keras.optimizers.Adam(learning_rate=1e-3),
              loss="sparse_categorical_crossentropy",
              metrics=["accuracy"])

callbacks = [
    # Cut the learning rate when validation loss stops improving
    keras.callbacks.ReduceLROnPlateau(monitor="val_loss", factor=0.3,
                                      patience=2, min_lr=1e-5, verbose=1),
    # Stop once it has stalled for a while and keep the best weights seen
    keras.callbacks.EarlyStopping(monitor="val_loss", patience=6,
                                  restore_best_weights=True, verbose=1),
]

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

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