__author__ = 'Lukas Leufen'
__date__ = '2019-10-21'


import logging
import keras
import keras.backend as K
import math
from typing import Union
import numpy as np


def to_list(arg):
    if not isinstance(arg, list):
        arg = [arg]
    return arg


def l_p_loss(power: int):
    """
    Calculate the L<p> loss for given power p. L1 (p=1) is equal to mean absolute error (MAE), L2 (p=2) is to mean
    squared error (MSE), ...
    :param power: set the power of the error calculus
    :return: loss for given power
    """
    def loss(y_true, y_pred):
        return K.mean(K.pow(K.abs(y_pred - y_true), power), axis=-1)
    return loss


class LearningRateDecay(keras.callbacks.History):
    """
    Decay learning rate during model training. Start with a base learning rate and lower this rate after every
    n(=epochs_drop) epochs by drop value (0, 1], drop value = 1 means no decay in learning rate.
    """

    def __init__(self, base_lr: float = 0.01, drop: float = 0.96, epochs_drop: int = 8):
        super().__init__()
        self.lr = {'lr': []}
        self.base_lr = self.check_param(base_lr, 'base_lr')
        self.drop = self.check_param(drop, 'drop')
        self.epochs_drop = self.check_param(epochs_drop, 'epochs_drop', upper=None)

    @staticmethod
    def check_param(value: float, name: str, lower: Union[float, None] = 0, upper: Union[float, None] = 1):
        """
        Check if given value is in interval. The left (lower) endpoint is open, right (upper) endpoint is closed. To
        only one side of the interval, set the other endpoint to None. If both ends are set to None, just return the
        value without any check.
        :param value: value to check
        :param name: name of the variable to display in error message
        :param lower: left (lower) endpoint of interval, opened
        :param upper: right (upper) endpoint of interval, closed
        :return: unchanged value or raise ValueError
        """
        if lower is None:
            lower = -np.inf
        if upper is None:
            upper = np.inf
        if lower < value <= upper:
            return value
        else:
            raise ValueError(f"{name} is out of allowed range ({lower}, {upper}{')' if upper == np.inf else ']'}: "
                             f"{name}={value}")

    def on_epoch_begin(self, epoch: int, logs=None):
        """
        Lower learning rate every epochs_drop epochs by factor drop.
        :param epoch: current epoch
        :param logs: ?
        :return: update keras learning rate
        """
        current_lr = self.base_lr * math.pow(self.drop, math.floor(epoch / self.epochs_drop))
        K.set_value(self.model.optimizer.lr, current_lr)
        self.lr['lr'].append(current_lr)
        logging.info(f"Set learning rate to {current_lr}")
        return K.get_value(self.model.optimizer.lr)