import keras
import numpy as np


class EpochLogRecorder(keras.callbacks.Callback):
    def __init__(self):
        super().__init__()
        self.epochs = []
        self.log_keys = []

    def on_train_begin(self, logs=None):
        self.epochs.clear()
        self.log_keys.clear()

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