from tensorflow.keras.datasets import mnist
(train_images, Y_train), (test_images, Y_test) = mnist.load_data()
train_images = train_images.reshape((train_images.shape[0], -1))/255.
test_images = test_images.reshape((test_images.shape[0], -1))/255.

"""
Explanation: Adding a custom callback for keras during embedding
This notebook shows you how to use custom callbacks during training. In this example, we use early stopping to train the network until loss reaches some desired plateau. 
load data
End of explanation
"""


from umap.parametric_umap import ParametricUMAP

keras_fit_kwargs = {"callbacks": [
    tf.keras.callbacks.EarlyStopping(
        monitor='loss',
        min_delta=10**-2,
        patience=10,
        verbose=1,
    )
]}

embedder = ParametricUMAP(
    verbose=True,
    keras_fit_kwargs = keras_fit_kwargs,
    n_training_epochs=5
)

embedding = embedder.fit_transform(train_images)

"""
Explanation: create parametric umap model
End of explanation
"""


embedding = embedder.embedding_

import matplotlib.pyplot as plt

fig, ax = plt.subplots( figsize=(8, 8))
sc = ax.scatter(
    embedding[:, 0],
    embedding[:, 1],
    c=Y_train.astype(int),
    cmap="tab10",
    s=0.1,
    alpha=0.5,
    rasterized=True,
)
ax.axis('equal')
ax.set_title("UMAP in Tensorflow embedding", fontsize=20)
plt.colorbar(sc, ax=ax);

"""
Explanation: plot results
End of explanation
"""


embedder._history.keys()

fig, ax = plt.subplots(figsize=(10,5))
ax.plot(embedder._history['loss'])
ax.set_ylabel('Cross Entropy')
ax.set_xlabel('Epoch')

"""
Explanation: plotting loss
End of explanation
"""


