## taken/adapted from "Deep Learning with Python (Third Edition)” by Chollet
## CNN for MNIST

import os

# Sets the environment variable from within the Python runtime
os.environ["KERAS_BACKEND"] = "torch"

import keras
from keras import layers

##asked ChatGPT to generate it....

import matplotlib.pyplot as plt   # <-- IMPORTANT: fixes NameError

# Load MNIST
(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()

# Preprocess: add channel dimension
x_train = x_train[..., None].astype("float32") / 255.0
x_test  = x_test[..., None].astype("float32") / 255.0

input_shape = (28, 28, 1)   

# Build CNN model
model = keras.Sequential([
    layers.Input(shape=input_shape),          # <-- cleaner input declaration

    layers.Conv2D(32, (3,3), activation="relu", padding="same"),
    layers.MaxPooling2D((2,2)),

    layers.Conv2D(64, (3,3), activation="relu", padding="same"),
    layers.MaxPooling2D((2,2)),

    layers.Flatten(),
    layers.Dense(128, activation="relu"),
    layers.Dropout(0.3),
    layers.Dense(10, activation="softmax")
])

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

# Train CNN
history = model.fit(
    x_train, y_train,
    validation_split=0.1,
    epochs=10,
    batch_size=128,
    verbose=2
)

# Report accuracies
train_acc = history.history["accuracy"][-1]
val_acc = history.history["val_accuracy"][-1]
test_loss, test_acc = model.evaluate(x_test, y_test, verbose=0)

print(f"Final training accuracy:     {train_acc:.4f}")
print(f"Final validation accuracy:   {val_acc:.4f}")
print(f"Final test accuracy:         {test_acc:.4f}")

# ---- Plot learning curves ----
epochs = range(1, len(history.history["accuracy"]) + 1)

plt.figure(figsize=(14, 5))

# Accuracy
plt.subplot(1, 2, 1)
plt.plot(epochs, history.history["accuracy"], label="Training Accuracy")
plt.plot(epochs, history.history["val_accuracy"], label="Validation Accuracy")
plt.title("CNN Accuracy")
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.legend()

# Loss
plt.subplot(1, 2, 2)
plt.plot(epochs, history.history["loss"], label="Training Loss")
plt.plot(epochs, history.history["val_loss"], label="Validation Loss")
plt.title("CNN Loss")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()

plt.tight_layout()
plt.show()

##show the # of paras in each layer:
model.summary()