"""Copyright © 2020-present, Swisscom (Schweiz) AG.
All rights reserved.

SingleObjectiveCDV, used to represent a Common
Descent Vector for single objective

The SingleObjectiveCDV class is mostly a wrapper
for a single objective loss
"""

from commondescentvector.common_descent_vector import CommonDescentVector


class SingleObjectiveCDV(CommonDescentVector):
    """ SingleObjectiveCDV class. Inherits from the CommonDescentVector
    class.

    SingleObjectiveCDV is used to return the common descent vector for
    a single objective

    Attributes:
        max_empirical_loss: A PyTorch Tensor loss representing the max empirical
            loss for the objective
        loss: A PyTorch Tensor loss representing the actual loss of the objective
    """

    def __init__(self, max_empirical_loss=None, normalized=False):
        """Inits SingleObjectiveCDV

        Args:
            max_empirical_loss: A Pytorch Tensor of size 1 containing the max empirical
                loss for the objective. Must be greater than 0 if set. Default: None
            normalized: Boolean, set to True if the output of get_descent_vector
            must be normalized. Default value: False
        Raises:
            ValueError: An error occured when max_empirical_loss is less than 0.
        """

        if max_empirical_loss is not None:
            if max_empirical_loss <= 0:
                raise ValueError('Argument: max_empirical_loss must be greater than 0')

        self.__max_empirical_loss = max_empirical_loss
        self.__normalized = normalized

    def get_descent_vector(self, losses, gradients=None):
        """Return the common descent vector for a single objective.
        Normalized if normalized has been set to True during initialization


        Args:
            losses: A numpy array containing 1 PyTorch Tensor representing the loss
            gradients: unused
        Return:
            loss: A pytorch tensor representing the loss
            alpha: Always returns 1 by this function
        Raise:
            TypeError: An error occured while accessing the argument - the
            argument is NoneType
        """

        if losses is None:
            raise TypeError('Argument: losses cannot be NoneType')
        if len(losses) != 1:
            raise ValueError('Argument: losses must be of size 1')

        if self.__normalized:
            if self.__max_empirical_loss is None:
                raise TypeError('max_empirical_loss must be set')
            else:
                return (losses[0]/self.__max_empirical_loss)
        else:
            return losses[0], 1

    def set_max_empirical_loss(self, max_empirical_loss):
        """Set the max empirical loss

        Args:
            max_empirical_loss: A PyTorch Tensor loss representing the max empirical
            loss for the objective
        """

        if max_empirical_loss <= 0:
            raise ValueError('Argument: max_empirical_loss must be greater than 0')

        self.__max_empirical_loss = max_empirical_loss

    def set_normalized(self, normalized):
        """Set if the CommonDescentVector should be normalized

        Args:
            normalized: True if the descent vector should be normalized
        """

        self.__normalized = normalized
