import keras
import numpy as np


class EpochLogRecorder(keras.callbacks.Callback):
    def on_train_begin(self, logs=None):
        self.seen_epochs = []
        self.log_key_history = []

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        keys = sorted(logs.keys())
        self.seen_epochs.append(epoch + 1)
        self.log_key_history.append(keys)
        print(f"epoch={epoch + 1} log_keys={','.join(keys)}")


keras.utils.set_random_seed(7)

x_train = np.arange(0, 8, dtype="float32").reshape(-1, 1)
y_train = (2 * x_train) + 1

model = keras.Sequential(
    [
        keras.layers.Input(shape=(1,)),
        keras.layers.Dense(1),
    ]
)
model.compile(
    optimizer=keras.optimizers.SGD(learning_rate=0.05),
    loss="mse",
    metrics=["mae"],
)

callback = EpochLogRecorder()
history = model.fit(
    x_train,
    y_train,
    batch_size=4,
    epochs=3,
    verbose=0,
    callbacks=[callback],
)

print(f"epochs recorded: {callback.seen_epochs}")
print(f"log snapshots: {len(callback.log_key_history)}")
print(f"history keys: {sorted(history.history.keys())}")