"""
Lecture 14 -- CNN in Practice: Flattening, FC Layers, Softmax & Implementation
A complete, runnable Keras CNN for MNIST digit classification (10 classes,
28x28x1 input), using the instructor's own layer choices:
Conv2D(32,5x5) -> MaxPool(2x2) -> Conv2D(64,5x5) -> MaxPool(2x2)
-> Flatten -> Dense(1000, relu) -> Dense(num_classes, softmax)

CPU training note: ~1-2 minutes/epoch on a typical laptop CPU for MNIST.
GPU (local or Colab) note: a few seconds/epoch -- prefer GPU for anything
beyond a quick correctness check. See resources.html for a link to the
official TensorFlow CNN tutorial (Colab-enabled) for a full training run.
"""
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Flatten, Dense
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical  # not used with sparse CE, kept for reference

# ---- data: MNIST handwritten digits ----
(x_train, y_train), (x_test, y_test) = mnist.load_data()
input_shape = (28, 28, 1)
num_classes = 10

x_train = x_train.reshape(-1, 28, 28, 1).astype("float32") / 255.0
x_test  = x_test.reshape(-1, 28, 28, 1).astype("float32") / 255.0
# labels stay as plain integers 0-9 -- matches sparse_categorical_crossentropy below

# ---- model: instructor's layer stack ----
model = Sequential()
model.add(Conv2D(32, kernel_size=(5, 5), strides=(1, 1), activation='relu', input_shape=input_shape))
model.add(MaxPooling2D(pool_size=(2, 2), strides=(2, 2)))
model.add(Conv2D(64, (5, 5), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(1000, activation='relu'))
model.add(Dense(num_classes, activation='softmax'))

model.summary()

# ---- compile: Adam optimizer (Lecture 9) + categorical cross-entropy (Section 3) ----
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

# ---- train ----
# NOTE: run a small number of epochs / a data subset for a quick CPU check;
# see the box in Lecture 14 Section 6 for why a full run belongs on GPU.
model.fit(x_train, y_train,
          batch_size=128,
          epochs=5,
          validation_split=0.1)

test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0)
print(f"Test accuracy: {test_acc:.4f}")