%pip install -U "comet_ml>=3.44.0"

import comet_ml

comet_ml.login(project_name="keras-image-classification")

from tensorflow import keras
from tensorflow.keras.datasets import mnist
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping

# create an experiment with your api key
experiment = comet_ml.start()

# the data, shuffled and split between train and test sets
(x_train, y_train), (x_test, y_test) = mnist.load_data()

x_train = x_train.reshape(60000, 784)
x_test = x_test.reshape(10000, 784)

x_train = x_train.astype("float32")
x_test = x_test.astype("float32")

x_train /= 255
x_test /= 255

print(x_train.shape[0], "train samples")
print(x_test.shape[0], "test samples")

# these will all get logged
params = {
    "batch_size": 128,
    "epochs": 2,
    "layer1_type": "Dense",
    "layer1_num_nodes": 64,
    "layer1_activation": "relu",
    "optimizer": "adam",
}

num_classes = 10

# convert class vectors to binary class matrices
y_train = keras.utils.to_categorical(y_train, num_classes)
y_test = keras.utils.to_categorical(y_test, num_classes)

model = Sequential()
model.add(Dense(64, activation="relu", input_shape=(784,)))
model.add(Dense(num_classes, activation="softmax"))

# print model.summary() to preserve automatically in `Output` tab
print(model.summary())

model.compile(
    loss="categorical_crossentropy", optimizer=params["optimizer"], metrics=["accuracy"]
)

experiment.log_parameters(params)
experiment.log_dataset_hash(x_train)  # creates and logs a hash of your data

# will log metrics with the prefix 'train_'
model.fit(
    x_train,
    y_train,
    batch_size=params["batch_size"],
    epochs=params["epochs"],
    verbose=1,
    validation_data=(x_test, y_test),
)

loss, accuracy = model.evaluate(x_test, y_test)
metrics = {"loss": loss, "accuracy": accuracy}
with experiment.test():
    experiment.log_metrics(metrics)

predictions = model.predict(x_test)

# Logs the image corresponding to the model prediction
def index_to_example(index):
    image_array = x_test[index].reshape(28, 28)
    image_name = "confusion-matrix-%05d.png" % index
    results = experiment.log_image(image_array, name=image_name)
    # Return sample, assetId (index is added automatically)
    return {"sample": image_name, "assetId": results["imageId"]}


LABELS = [f"class_{i}" for i in range(10)]

experiment.log_confusion_matrix(
    y_test[:100],
    predictions[:100],
    labels=LABELS,
    index_to_example_function=index_to_example,
    title="Confusion Matrix: Evaluation",
    file_name="confusion-matrix-eval.json",
)

experiment.end()