__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
import os
import time
import socket


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


def check_path_and_create(path):
    try:
        os.makedirs(path)
        logging.debug(f"Created path: {path}")
    except FileExistsError:
        logging.debug(f"Path already exists: {path}")


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)


class TimeTracking(object):
    """
    Track time to measure execution time. Time tracking automatically starts on initialisation and ends by calling stop
    method. Duration can always be shown by printing the time tracking object or calling get_current_duration.
    """

    def __init__(self, start=True):
        self.start = None
        self.end = None
        if start:
            self._start()

    def _start(self):
        self.start = time.time()
        self.end = None

    def _end(self):
        self.end = time.time()

    def _duration(self):
        if self.end:
            return self.end - self.start
        else:
            return time.time() - self.start

    def __repr__(self):
        return f"{round(self._duration(), 2)}s"

    def run(self):
        self._start()

    def stop(self, get_duration=False):
        if self.end is None:
            self._end()
        else:
            msg = f"Time was already stopped {time.time() - self.end}s ago."
            logging.error(msg)
            raise AssertionError(msg)
        if get_duration:
            return self.duration()

    def duration(self):
        return self._duration()


def prepare_host(create_new=True):
    hostname = socket.gethostname()
    try:
        user = os.getlogin()
    except OSError:
        user = "default"
    if hostname == 'ZAM144':
        path = f'/home/{user}/Data/toar_daily/'
    elif hostname == 'zam347':
        path = f'/home/{user}/Data/toar_daily/'
    elif hostname == 'linux-gzsx':
        path = f'/home/{user}/machinelearningtools/data/toar_daily/'
    elif (len(hostname) > 2) and (hostname[:2] == 'jr'):
        path = f'/p/project/cjjsc42/{user}/DATA/toar_daily/'
    elif (len(hostname) > 2) and (hostname[:2] == 'jw'):
        path = f'/p/home/jusers/{user}/juwels/intelliaq/DATA/toar_daily/'
    elif "runner-6HmDp9Qd-project-2411-concurrent" in hostname:
        path = f'/home/{user}/machinelearningtools/data/toar_daily/'
    else:
        logging.error(f"unknown host '{hostname}'")
        raise OSError(f"unknown host '{hostname}'")
    if not os.path.exists(path):
        try:
            if create_new:
                check_path_and_create(path)
                return path
            else:
                raise PermissionError
        except PermissionError:
            logging.error(f"path '{path}' does not exist for host '{hostname}'.")
            raise NotADirectoryError(f"path '{path}' does not exist for host '{hostname}'.")
    else:
        logging.debug(f"set path to: {path}")
        return path


def set_experiment_name(experiment_date=None, experiment_path=None):

    if experiment_date is None:
        experiment_name = "TestExperiment"
    else:
        experiment_name = f"{experiment_date}_network/"
    if experiment_path is None:
        experiment_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", experiment_name))
    else:
        experiment_path = os.path.abspath(experiment_path)
    return experiment_name, experiment_path