# Author: Dimitrios Damopoulos
# MIT license (see LICENCE.txt in the top-level folder)

# TODO: documentation

import numpy as np

def sigmoid(z):
    """
    Computes the sigmoid function

    Args:
        z (np.ndarray): the values where the sigmoid should be calculated on

    Returns:
        An np.ndarray with the same shape as the input `z' with the values of 
        sigmoid
    """
    sigm = 1. / (1. + np.exp(-z))
    return sigm

def cross_entropy(p_hat, p):
    """
    Computer the cross entropy between two distributions

    Args:
        p_hat (np.ndarray of shape m,): The "predicted" probability distribution
        p (np.ndarray of shape m,): The "true" probability distribution

    Returns:
        The cross-entropy between the two distributions
    """
    # get values which are very close to zero a bit higher, in order to avoid
    # NaNs in the computation of the logarithm later.
    eps = np.finfo(float).eps
    p_hat[p_hat < eps] = eps

    ce = - np.log(p_hat).T.dot(p)
    return ce

def cross_entropy_with_one_hot_vector(p_hat, p):
    """
    Same as math_utils.cross-entropy, but for a "true" probability distribution
    has arose for a binary random variable (i.e., a process that outputs True 
    or False for every sample. It should output the same as 
    math_utils.cross_entropy, but faster.

    Args:
        p_hat (np.ndarray of shape m,): The "predicted" probability distribution
        p (np.ndarray of shape m,): The "true" probability distribution. Its 
            values are converted to binary before the computation of the 
            cross-entropy, by just comparing them with 0: anything above zero is 
            treated as have value "1" and the rest as having value "0"
    
    Returns:
        The cross-entropy between the two distributions
    """
    eps = np.finfo(float).eps
    p1 = p > 0
    p0 = np.logical_not(p1)
    part1 = np.log(p_hat[p1] + eps).sum()
    part0 = np.log(1. - p_hat[p0] + eps).sum()
    
    return -(part0 + part1)

def rmse(y_hat, y):
    """ 
    Returns the Root Mean Square error between two series of measurements 
    """
    losses = y_hat - y
    se = losses.T.dot(losses)
    return np.sqrt(se / y.shape[0])

def gradient_of_rmse(y_hat, y, Xn):
    """ 
    Returns the gradient of the Root Mean Square error with respect to the 
    parameters of the linear model that generated the prediction `y_hat'. 
    Hence, y_hat should have been generated by a linear process of the form
    Xn.T.dot(theta)

    Args:
        
       y_hat (np.array of shape N,): The predictions of the linear model
       y (np.array of shape N,): The "ground-truth" values.

    Returns:
        The RMSE between y_hat and y
    """
        
    N = y.shape[0]
    assert N > 0, ('At least one sample is required in order to compute the '
                  'RMSE loss')
   
    losses = y - y_hat
    gradient = - 2 * Xn.T.dot(losses) / N

    return gradient
