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

import os

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

# Only then should you import Keras.
import keras

from keras import models
from keras import layers

from keras.datasets import mnist

##asked ChatGPT to add Dropout layers and traning curves...

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

# Preprocess
x_train = x_train.reshape(-1, 28*28).astype("float32") / 255.0
x_test  = x_test.reshape(-1, 28*28).astype("float32") / 255.0

# Build MLP model with Dropout
model = keras.Sequential([
    layers.Dense(256, activation="relu"),
    layers.Dropout(0.3),
    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
history = model.fit(
    x_train, y_train,
    validation_split=0.1,
    epochs=15,
    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 ----
import matplotlib.pyplot as plt
epochs = range(1, len(history.history["accuracy"]) + 1)

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

# Accuracy subplot
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("Accuracy with Dropout")
plt.xlabel("Epoch")
plt.ylabel("Accuracy")
plt.legend()

# Loss subplot
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("Loss with Dropout")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()

plt.tight_layout()
plt.show()

# Show model summary
model.summary()