from keras.models import load_model as load_model_keras
from tensorflow.python.client import device_lib
from multiprocessing import cpu_count
from keras.callbacks import Callback
from six import iteritems
import numpy as np


def load_model(model_path):
    return load_model_keras(model_path)


def use_multiprocessing():
    if __get_available_gpus():
        # if GPU is available, use all available CPUs for batch preprocessing
        use_multiprocessing = True
        num_workers = cpu_count()
    else:
        # device = 'CPU'
        use_multiprocessing = False
        num_workers = 1
    return use_multiprocessing, num_workers


def __get_available_gpus():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos if x.device_type == 'GPU']


class LoggingMetrics(Callback):
    """Callback for logging metrics at the end of each epoch.
    Default format looks like "Epoch: 3 - loss: 3.4123 - val_loss: 5.4321"
    # Arguments
        self.logger.debug_fcn: function for self.logger.debuging. default is self.logger.debug, which will self.logger.debug to standard output.
        format_epoch: format string for each epoch [default="Epoch: {} - {}"]
        format_keyvalue: format string for key value pairs [default="{}: {}"]
        format_separator= separator string between each pair [default=" - "]
    # Example
        ```python
        # Write to Python logging
        import logging
        model.fit(x, y, callbacks = [LoggingCallback(logging.info)]
        # Write to a file
        with open("log.txt", 'w') as f:
            def self.logger.debug_fcn(s):
                f.write(s)
                f.write("\n")
            model.fit(x, y, callbacks = [LoggingCallback(self.logger.debug_fcn)]
        ```
    """

    def __init__(
        self,
        logger,
        format_epoch="Epoch: {} - {}",
        format_keyvalue="{}: {:0.4f}",
        format_separator=" - ",
    ):
        Callback.__init__(self)
        self.logger = logger
        self.format_epoch = format_epoch
        self.format_keyvalue = format_keyvalue
        self.format_separator = format_separator

    def on_epoch_end(self, epoch, logs={}):
        values = self.format_separator.join(
            self.format_keyvalue.format(k, v) for k, v in iteritems(logs)
        )
        msg = self.format_epoch.format(epoch + 1, values)
        self.logger.debug(msg)


class LoggingModels(Callback):
    """
    """

    def __init__(
        self,
        filepath,
        logger,
        monitor='val_loss',
        verbose=0,
        save_best_only=False,
        save_weights_only=False,
        mode='auto',
        period=1,
    ):
        self.monitor = monitor
        self.verbose = verbose
        self.filepath = filepath
        self.save_best_only = save_best_only
        self.save_weights_only = save_weights_only
        self.period = period
        self.epochs_since_last_save = 0
        self.logger = logger

        if mode not in ['auto', 'min', 'max']:
            self.logger.warning(
                'ModelCheckpoint mode %s is unknown, ' 'fallback to auto mode.' % (mode),
                RuntimeWarning,
            )
            mode = 'auto'

        if mode == 'min':
            self.monitor_op = np.less
            self.best = np.Inf
        elif mode == 'max':
            self.monitor_op = np.greater
            self.best = -np.Inf
        else:
            if 'acc' in self.monitor or self.monitor.startswith('fmeasure'):
                self.monitor_op = np.greater
                self.best = -np.Inf
            else:
                self.monitor_op = np.less
                self.best = np.Inf

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        self.epochs_since_last_save += 1
        if self.epochs_since_last_save >= self.period:
            self.epochs_since_last_save = 0
            filepath = self.filepath.format(epoch=epoch + 1, **logs)
            if self.save_best_only:
                current = logs.get(self.monitor)
                if current is None:
                    self.logger.warning(
                        'Can save best model only with %s available, ' 'skipping.' % (self.monitor),
                        RuntimeWarning,
                    )
                else:
                    if self.monitor_op(current, self.best):
                        if self.verbose > 0:
                            self.logger.debug(
                                '\nEpoch %05d: %s improved from %0.5f to %0.5f,'
                                ' saving model to %s'
                                % (epoch + 1, self.monitor, self.best, current, filepath)
                            )
                        self.best = current
                        if self.save_weights_only:
                            self.model.save_weights(filepath, overwrite=True)
                        else:
                            self.model.save(filepath, overwrite=True)
                    else:
                        if self.verbose > 0:
                            self.logger.debug(
                                '\nEpoch %05d: %s did not improve from %0.5f'
                                % (epoch + 1, self.monitor, self.best)
                            )
            else:
                if self.verbose > 0:
                    self.logger.debug('\nEpoch %05d: saving model to %s' % (epoch + 1, filepath))
                if self.save_weights_only:
                    self.model.save_weights(filepath, overwrite=True)
                else:
                    self.model.save(filepath, overwrite=True)
