import math
import os
import traceback
from functools import partial

from keras.callbacks import ModelCheckpoint, CSVLogger, LearningRateScheduler, ReduceLROnPlateau, EarlyStopping, \
    TensorBoard


def get_callbacks(model_file, initial_learning_rate=1e-4, learning_rate_drop=0.5, learning_rate_epochs=None,
                  learning_rate_patience=10, logging_file="training.log", verbosity=1, early_stopping_patience=None):
    callbacks = list()
    callbacks.append(
        ModelCheckpoint(
            model_file, monitor='val_loss', save_best_only=True))
    callbacks.append(
        CSVLogger(
            os.path.join(os.path.dirname(model_file), logging_file), append=True))
    callbacks.append(
        TensorBoard(
            log_dir=os.path.join(os.path.dirname(os.path.dirname(model_file)), 'log'),
            histogram_freq=0, write_graph=False, write_images=False,
            embeddings_freq=0, embeddings_layer_names=None,
            embeddings_metadata=None))
    if learning_rate_epochs:
        callbacks.append(
            LearningRateScheduler(
                partial(step_decay, initial_lrate=initial_learning_rate,
                        drop=learning_rate_drop, epochs_drop=learning_rate_epochs)))
    else:
        callbacks.append(
            ReduceLROnPlateau(
                factor=learning_rate_drop, patience=learning_rate_patience, verbose=verbosity))
    if early_stopping_patience:
        callbacks.append(
            EarlyStopping(
                verbose=verbosity, patience=early_stopping_patience))

    return callbacks