import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten
from tensorflow.keras.optimizers import SGD, Adam, RMSprop
import matplotlib.pyplot as plt

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train / 255.0
x_test = x_test / 255.0

def create_model(optimizer):
    model = Sequential([Flatten(input_shape=(28, 28)), Dense(128, activation='relu'),Dense(10, activation='softmax')])
    model.compile(optimizer=optimizer, loss='sparse_categorical_crossentropy', metrics=['accuracy'])
    return model

optimizers = {"SGD": SGD(),"RMSprop": RMSprop(),"Adam": Adam()}
history_dict = {}

for name, opt in optimizers.items():
    print(f"nTraining with {name}")
    model = create_model(opt)
    history = model.fit(x_train, y_train, epochs=5, batch_size=32, validation_split=0.2, verbose=0)
    history_dict[name] = history

plt.figure(figsize=(8,5))
for name, history in history_dict.items():
    plt.plot(history.history['val_accuracy'], label=name)

plt.title("Optimizer Comparison - Validation Accuracy")
plt.xlabel("Epochs")
plt.ylabel("Accuracy")
plt.legend()
plt.show()