"""Huber Loss for Regression or Binary Classification. Robust loss, combination of quadratic loss and linear loss."""
import typing
import numpy as np
from h2oaicore.metrics import CustomScorer
from sklearn.preprocessing import LabelEncoder

# temp. references:
# Wikipedia: https://en.wikipedia.org/wiki/Huber_loss
# https://stackoverflow.com/questions/45006341/xgboost-how-to-use-mae-as-objective-function
# Tukey loss: https://web.as.uky.edu/statistics/users/pbreheny/764-F11/notes/12-1.pdf

class MyHuberLossScorer(CustomScorer):
    '''
    Huber Loss Scorer is a loss function used in robust regression, that is less
    sensitive to outliers in data than the squared error loss. This custom
    scorer supports both regression and binary binary classification problems
    using different formulas and different defaults for delta (see below).
    For more details see: https://en.wikipedia.org/wiki/Huber_loss

    Parameters
    ----------
    delta : numeric
        Hyperparameter with defaults =1.345 for regression and =0.1 for binary
        classification
    '''
    _delta_regression = 1.345
    _delta_binary = 0.1
    _description = "My Huber Loss for Regression or Binary Classification [delta=%f or %f]." % (
        _delta_regression, _delta_binary)
    _binary = True
    _regression = True
    _maximize = False
    _perfect_score = 0
    _display_name = "Huber"

    def score(self,
              actual: np.array,
              predicted: np.array,
              sample_weight: typing.Optional[np.array] = None,
              labels: typing.Optional[np.array] = None) -> float:

        if sample_weight is None:
            sample_weight = np.ones(actual.shape[0])

        isRegression = labels is None
        delta = MyHuberLossScorer._delta_regression if isRegression else MyHuberLossScorer._delta_binary
        if isRegression:
            abs_error = np.abs(np.subtract(actual, predicted))
            loss = np.where(abs_error < delta, .5 * (abs_error) ** 2, delta * (abs_error - 0.5 * delta))
        else:
            lb = LabelEncoder()
            labels = lb.fit_transform(labels)
            actual = lb.transform(actual)
            all0s = np.zeros(actual.shape[0])
            predicted = np.subtract(np.multiply(predicted, 2), 1)
            actual = np.where(actual == 0, -1, 1)
            actual_mult_predict = np.multiply(actual, predicted)
            loss = np.where(actual_mult_predict >= -1,
                            np.square(np.maximum(all0s, np.subtract(1, actual_mult_predict))),
                            -4 * actual_mult_predict)

        loss = np.sum(np.multiply(sample_weight, loss)) / np.sum(sample_weight)
        return float(loss) if actual.shape[0] > 0 else 0


import numpy as np

def sigmoid_sol(x):
    """
    Compute the sigmoid function for the input here.
    """
    ### YOUR CODE HERE
    x = 1. / (1 + np.exp(-x))
    ### END YOUR CODE
    return x

def sigmoid_grad_sol(f):
    """
    Compute the gradient for the sigmoid function here. Note that
    for this implementation, the input f should be the sigmoid
    function value of your original input x.
    """
    ### YOUR CODE HERE
    f = f * (1-f)
    ### END YOUR CODE
    return f


"""1.3.
Un consumidor cuyas preferencias viene representadas por la función de utilidad 𝑈(𝑥, 𝑦) = 3𝑙𝑛 𝑥 + 𝑙𝑛 𝑦, desea maximizar la utilidad que obtiene del consumo de ambos bienes. Siendo los precios a los que se enfrenta px = py = $3  y disponiendo de unos ingresos de $120, ¿cuáles serían las cantidades óptimas a consumir?
"""

import sympy as sp
sp.init_printing()
import funcionesAuxiliares as aux

m, x, y, px, py, t = sp.symbols('m,x,y,px,py,t')
Presupuesto = m - px*x - py*y
Utilidad = 3*sp.log(x) + sp.log(y)

im = 120
ipx = ipy = 3

ej13 = aux.maximizarRestriccion(Utilidad,
                         Presupuesto,
                         im,
                         ipx,
                         ipy )

ej13


import numpy as np

def si_sdr(reference, estimation):
    """
    Scale-Invariant Signal-to-Distortion Ratio (SI-SDR)

    Args:
        reference: numpy.ndarray, [..., T]
        estimation: numpy.ndarray, [..., T]

    Returns:
        SI-SDR

    [1] SDR– Half- Baked or Well Done?
    http://www.merl.com/publications/docs/TR2019-013.pdf

    >>> np.random.seed(0)
    >>> reference = np.random.randn(100)
    >>> si_sdr(reference, reference)
    inf
    >>> si_sdr(reference, reference * 2)
    inf
    >>> si_sdr(reference, np.flip(reference))
    -25.127672346460717
    >>> si_sdr(reference, reference + np.flip(reference))
    0.481070445785553
    >>> si_sdr(reference, reference + 0.5)
    6.3704606032577304
    >>> si_sdr(reference, reference * 2 + 1)
    6.3704606032577304
    >>> si_sdr([1., 0], [0., 0])  # never predict only zeros
    nan
    >>> si_sdr([reference, reference], [reference * 2 + 1, reference * 1 + 0.5])
    array([6.3704606, 6.3704606])

    """
    estimation, reference = np.broadcast_arrays(estimation, reference)

    assert reference.dtype == np.float64, reference.dtype
    assert estimation.dtype == np.float64, estimation.dtype

    reference_energy = np.sum(reference ** 2, axis=-1, keepdims=True)

    # This is $\alpha$ after Equation (3) in [1].
    optimal_scaling = np.sum(reference * estimation, axis=-1, keepdims=True) \
        / reference_energy

    # This is $e_{\text{target}}$ in Equation (4) in [1].
    projection = optimal_scaling * reference

    # This is $e_{\text{res}}$ in Equation (4) in [1].
    noise = estimation - projection

    ratio = np.sum(projection ** 2, axis=-1) / np.sum(noise ** 2, axis=-1)
    return 10 * np.log10(ratio)


import numpy as np
import re

class Variable():
    def __init__(self, var, der = 1):
        """
        Initialize a Variable object with following attributes.
        Derivative is 1 by default if not specified.

        :param var: attribute representing evaluated value.
        :param der: attribute representing evaluated derivative/gradient.
        """
        if isinstance(var, int) or isinstance(var, float):
            self.var = var
            self.der = der
        else:
            raise TypeError("Input is not a real number.")

    def __str__(self):
        return f"value = {self.var}, derivative = {self.der}"

    def __repr__(self):
        return f"value = {self.var}, derivative = {self.der}"

    def __add__(self, other):
        """
        dunder method for adding a Variable object.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be added to self.var.
        :return: Variable object with value and derivative of the sum of self and other.
        """
        try:
           new_add = self.var + other.var
           new_der = self.der + other.der
           return Variable(new_add, new_der)
        except:
            if isinstance(other, int) or isinstance(other, float):
                # other is not a variable and the addition could complete if it is a real number
                return Variable(self.var + other, self.der)
            else:
                raise TypeError("Input is not a real number.")

    def __mul__(self, other):
        """
        dunder method for multiplying a Variable object or constant.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be multiplied to self.var.
        :return: Variable object with value and derivative of the product of self and other.
        """
        try:
            new_mul = other.var * self.var
            new_der = self.der * other.var + other.der * self.var
            return Variable(new_mul, new_der)
        except:
            if isinstance(other, int) or isinstance(other, float):
                # other is not a variable and the multiplication could complete if it is a real number
                new_mul = other * self.var
                new_der = self.der * other
                return Variable(new_mul, new_der)
            else:
                raise TypeError("Input is not a real number.")

    def __radd__(self, other):
        """
        dunder method for adding a Variable object and left object without __add__ mehtod.

        :param self: Variable object.
        :param other: an object that does not have an __add__ method or not implemented.
        :return: __add__ function call.
        """
        return self.__add__(other)

    def __rmul__(self, other):
        """
        dunder method for multiplying a Variable object and left object without __mul__ mehtod.

        :param self: Variable object.
        :param other: an object that does not have an __mul__ method or not implemented.
        :return: __mul__ function call.
        """
        return self.__mul__(other)

    def __sub__(self, other):
        """
        dunder method for subtracting a Variable object or constant.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be subtracted from self.var.
        :return: __add__ function call that passes the other object with negative sign.
        """
        return self.__add__(-other)

    def __rsub__(self, other):
        """
        dunder method for subtracting a Variable object and left object without __sub__ mehtod.

        :param self: Variable object.
        :param other: an object that does not have an __sub__ method or not implemented.
        :return: __add__ function call that passes the self object with negative sign.
        """
        return (-self).__add__(other)

    def __truediv__(self, other):
        """
        dunder method for dividing a Variable object or constant.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be divided from self.var.
        :return: Variable object with value and derivative of the fraction of self and other.
        """
        try:
            new_div = self.var / other.var
            new_der = (self.der * other.var - other.der * self.var) / other.var**2
            return Variable(new_div, new_der)
        except AttributeError:
            if isinstance(other, int) or isinstance(other, float):
                new_div = self.var / other
                new_der = self.der / other
                return Variable(new_div, new_der)
            else:
                raise TypeError(f"Input {other} is not valid.")

    def __neg__(self):
        """
        dunder method for taking the negative of a Variable object or constant.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be taken the negative.
        :return: Variable object with negative value and derivative of itself.
        """
        return Variable(-self.var, -self.der)

    def __rtruediv__(self, other):
        """
        dunder method for subtracting a Variable object and left object without __truediv__ mehtod.

        :param self: Variable object.
        :param other: an object that does not have an __truediv__ method or not implemented.
        :return: Variable object with value and derivative of the fraction of self and other.
        """
        try:
            new_div = other.var / self.var
            new_der = (other.der * self.var - other.var * self.der) / self.var**2
            return Variable(new_div, new_der)
        except:
            if isinstance(other, int) or isinstance(other, float):
                new_div = other / self.var
                new_der = other * (self.var**(-2)) * self.der
                return Variable(new_div, new_der)
            else:
                raise TypeError(f"Input {other} is not valid.")

    def __lt__(self, other):
        '''
        dunder method for less than comparator.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be compared with.
        :return: True if self is less than other, otherwise False.
        '''
        try:
            return self.var < other.var
        except AttributeError:
            return self.var < other

    def __gt__(self, other):
        '''
        dunder method for greater than comparator.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be compared with.
        :return: True if self is greater than other, otherwise False.
        '''
        try:
            return self.var > other.var
        except AttributeError:
            return self.var > other

    def __le__(self, other):
        '''
        dunder method for less than or equal to comparator.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be compared with.
        :return: True if self is less than or equal to other, otherwise False.
        '''
        try:
            return self.var <= other.var
        except AttributeError:
            return self.var <= other

    def __ge__(self, other):
        '''
        dunder method for greater than or equal to comparator.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be compared with.
        :return: True if self is greater than or equal to other, otherwise False.
        '''
        try:
            return self.var >= other.var
        except AttributeError:
            return self.var >= other

    def __eq__(self, other):
        '''
        dunder method for equality comparator.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be compared with.
        :return: True if self is equal to other, otherwise False.
        '''
        try:
            return self.var == other.var
        except:
            raise TypeError('Input is not comparable.')

    def __ne__(self, other):
        '''
        dunder method for inequality comparator.

        :param self: Variable object.
        :param other: Variable object, or int/float, to be compared with.
        :return: negation of __eq__ function call.
        '''
        return not self.__eq__(other)

    def __abs__(self):
        '''
        dunder method for absolute value.

        :param self: Variable object.
        :return: Variable object with value and derivative of the absolute value of self.
        '''
        return Variable(abs(self.var), abs(self.der))

    def __pow__(self, other):
        """
        dunder method for taking Variable object to the value of other object's power.

        :param self: Variable object.
        :param other: Variable object, or int/float, to the power of.
        :return: Variable object with value and derivative of the power of self with power of other's value.
        """
        try:
            new_val = self.var ** other.var
            new_der = other.var * self.var ** (other.var-1)
            return Variable(new_val, new_der)
        except:
            if isinstance(other, int):
                new_val = self.var ** other
                new_der = other * self.var ** (other - 1) * self.der
                return Variable(new_val, new_der)
            else:
                raise TypeError(f"Exponent {other} is not valid.")

    def __rpow__(self, other):
        """
        dunder method for taking the other object to the Variable object's power.

        :param self: Variable object.
        :param other: Variable object, or int/float, as the base.
        :return: Variable object with value and derivative of the power of other with power of self's value.
        """
        try:
            new_val = other ** self.var
        except:
            raise ValueError("{} must be a number.".format(other))
        new_der = other**self.var * np.log(other)
        return Variable(new_val, new_der)

    @staticmethod
    def log(var):
        """
        takes the log of the Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the log of var.
        """
        try:
            if var.var <= 0:
                raise ValueError('Input needs to be greater than 0.')
        except:
            raise TypeError(f"Input not valid.")
        log_var = np.log(var.var)
        log_der = (1. / var.var) * var.der
        return Variable(log_var, log_der)

    @staticmethod
    def sqrt(var):
        """
        takes the square root of the Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the square root of var.
        """
        if var < 0:
            raise ValueError("Square root only takes positive number in the current implementation.")
        else:
            try:
                sqrt_var = var.var**(1/2)
                sqrt_der = (1/2)*var.var**(-1/2)
            except:
                raise TypeError(f"Input is not an Variable object.")
        return Variable(sqrt_var, sqrt_der)

    @staticmethod
    def exp(var):
        """
        natural exponential function with the value of Variable object as the power.

        :param var: Variable object.
        :return: Variable object with value and derivative of natural exponential function with power var.
        """
        try:
            new_val = np.exp(var.var)
            new_der = np.exp(var.var) * var.der
            return Variable(new_val, new_der)
        except:
            if not isinstance(var, int) and not isinstance(var, float):
                raise TypeError(f"Input {var} is not valid.")

            return Variable(np.exp(var), np.exp(var))

    @staticmethod
    def sin(var):
        """
        calculates the sine of Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the sine of var.
        """
        try:
            new_val = np.sin(var.var)
            new_der = var.der * np.cos(var.var)
            return Variable(new_val, new_der)
        except:
            if not isinstance(var, int) and not isinstance(var, float):
                raise TypeError(f"Input {var} is not valid.")

            return Variable(np.sin(var), np.cos(var))

    @staticmethod
    def cos(var):
        """
        calculates the cosine of Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the cosine of var.
        """
        try:
            new_val = np.cos(var.var)
            new_der = var.der * -np.sin(var.var)
            return Variable(new_val, new_der)
        except:
            if not isinstance(var, int) and not isinstance(var, float):
                raise TypeError(f"Input {var} is not valid.")

            return Variable(np.cos(var), -np.sin(var))

    @staticmethod
    def tan(var):
        """
        calculates the tangent of Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the tangent of var.
        """
        try:
            new_val = np.tan(var.var)
            new_der = var.der * 1 / np.power(np.cos(var.var), 2)
            return Variable(new_val, new_der)
        except:
            if not isinstance(var, int) and not isinstance(var, float):
                raise TypeError(f"Input {var} is not valid.")

            return Variable(np.tan(var), 1/np.cos(var)**2)

    @staticmethod
    def arcsin(var):
        """
        calculates the arcsine of Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the arcsine of var.
        """
        try:
            if var.var > 1 or var.var < -1:
                raise ValueError('Please input -1 <= x <=1')
            else:
                new_val = np.arcsin(var.var)
                new_der = 1 / np.sqrt(1 - (var.var ** 2))
                return Variable(new_val, new_der)
        except:
            if not isinstance(var, int) and not isinstance(var, float):
                raise TypeError(f"Input {var} is not valid.")
            return Variable(np.arcsin(var), 1 / np.sqrt(1 - (var ** 2)))

    @staticmethod
    def arccos(var):
        """
        calculates the arccosine of Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the arccosine of var.
        """
        try:
            if isinstance(var, int) or isinstance(var, float):
                return np.arccos(var)

            if var.var > 1 or var.var < -1:
                raise ValueError('Please input -1 <= x <=1')
            else:
                new_val = np.arccos(var.var)
                new_der = -1 / np.sqrt(1 - (var.var ** 2))
            return Variable(new_val, new_der)
        except:
                raise TypeError(f"Input {var} is not valid.")

    @staticmethod
    def arctan(var):
        """
        calculates the arctangent of Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the arctangent of var.
        """
        try:
            new_val = np.arctan(var.var)
            new_der = var.der * 1 / (1 + np.power(var.var, 2))

            return Variable(new_val, new_der)

        except AttributeError:
            return np.arctan(var)

    @staticmethod
    def sinh(var):
        """
        calculates the hyperbolic sine of Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the hyperbolic sine of var.
        """
        try:
            new_val = np.sinh(var.var)
            new_der = var.der * np.cosh(var.var)
            return Variable(new_val, new_der)

        except AttributeError:
            return np.sinh(var)

    @staticmethod
    def cosh(var):
        """
        calculates the hyperbolic cosine of Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the hyperbolic cosine of var.
        """
        try:
            new_val = np.cosh(var.var)
            new_der = var.der * np.sinh(var.var)

            return Variable(new_val, new_der)

        except AttributeError:
            return np.cosh(var)

    @staticmethod
    def tanh(var):
        """
        calculates the hyperbolic tangent of Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the hyperbolic tangent of var.
        """
        try:
            new_val = np.tanh(var.var)
            new_der = var.der * 1 / np.power(np.cosh(var.var), 2)
            return Variable(new_val, new_der)
        except AttributeError:
            return np.tanh(var)

    def sigmoid(var):
        """
        calculates the sigmoid/logistic function of Variable object.

        :param var: Variable object.
        :return: Variable object with value and derivative of the sigmoid/logistic function of var.
        """
        try:
            logistic_var = 1 / (1 + np.exp(-var.var))
            logistic_der = logistic_var * (1-logistic_var) * var.der
            return Variable(logistic_var, logistic_der)
        except:
            raise TypeError(f"Input {var} not valid.")

class SimpleAutoDiff:
    def __init__(self, dict_val, list_funct):
        """
        How To Use SimpleAutoDiff
        ---------------------------------------------------------------------------------------------------------------------

        Inputs:

        dict_val: a dictionary object
        dict_val represents the input which is a dictionary of variables and their values formatted like the examples below
        example1={'x':7,'y':2, 'z':3}
        example2 ={'x':4}
        You can input as many variables as needed so long as it is more than 1

        list_funct: a list object
        list_funct represents the input which is a list of functions input as strings
        example_a = ['x**2','cos(np.pi*y)', '8*z']
        example_b = ['sin(x)']

        An example of a sample return using dict_val example1 and list_funct example_a would be

        ---AutoDifferentiation---
        Value: {'x': 7, 'y': 2, 'z': 3}
        Function 1:
        Expression = x**2
        Value = 49
        Gradient = 14
        Function 2:
        Expression = cos(np.pi*y)
        Value = 1.0
        Gradient = 7.694682774887159e-16
        Function 3:
        Expression = 8*z
        Value = 24
        Gradient = 8
        ---------------------------------------------------------------------------------------------------------------------

        """
        for func in list_funct:
            if not isinstance(func, str):
                raise TypeError('Invalid function input.')

        static_elem_funct = ['log', 'sqrt', 'exp', 'sin', 'cos', 'tan', 'arcsin', 'arccos', 'arctan', 'sinh', 'cosh', 'tanh', 'sigmoid']
        func_vals = []
        dict_keys =[]
        dict_vals = []
        self.jacobian = np.zeros((len(list_funct), len(dict_val)))
        count = 0

        for val, key in enumerate(dict_val):
            dict_keys.append(key)
            dict_vals.append(val)

        for pair in range(0,len(dict_val)):
            for _ in range(0,len(dict_val)):
                if _ == count:
                    exec(dict_keys[_] + "= Variable(dict_val[dict_keys[_]], der=1)")
                else:
                    exec(dict_keys[_] + "= Variable(dict_val[dict_keys[_]], der=0)")

            for fun in range(0,len(list_funct)):
                for elem_funct in static_elem_funct:
                    if elem_funct in list_funct[fun]: # e.g. log is in log(x)
                        func = 'Variable.' + list_funct[fun]
                        break
                    else:
                        func = list_funct[fun]
                func_vals.append(eval(func).var)
                self.jacobian[fun,count]= eval(func).der
            count+=1
        self.functions = func_vals
        self.dict_val = dict_val
        self.list_funct = list_funct

    def __repr__ (self):
        output = '---AutoDifferentiation---\n'
        added_output = ''
        added_output += f"Value: {self.dict_val}\n\n"
        for i in range(0, len(self.functions)):
            added_output += f"Function {i+1}: \nExpression = {self.list_funct[i]}\nValue = {str(self.functions[i])}\nGradient = {str(self.jacobian[i])}\n\n"

        return output+added_output

    def __str__(self):
        output = '---AutoDifferentiation---\n'
        added_output = ''
        added_output += f"Value: {self.dict_val}\n\n"
        for i in range(0, self.jacobian.shape[0]):
            added_output += f"Function {i+1}: \nExpression = {self.list_funct[i]}\nValue = {str(self.functions[i])}\nGradient = {str(self.jacobian[i])}\n\n"

        return output+added_output

class Node():
    def __init__(self, var):
        """
        Initialize a Node object with follow attributes:
        child_node: a list that holds tuples of all depending Nodes and derivatives
        derivative: attribute representing evaluated derivative/gradient. Derivative is 1 by default.

        :param var: attribute representing evaluated value.
        """
        if isinstance(var, int) or isinstance(var, float):
            self.var = var
            self.child_node = []
            self.derivative = None
        else:
            raise TypeError("Input is not a real number.")

    def get_derivatives(self, inputs):
        """
        Method to get derivatives for each variable used in the function.
        This function uses:
        var_val: a variable which stored the function values
        der_list: a list of derivatives with respect to each variable

        :param inputs: the list of input functions.
        """
        # self.der = 1
        var_val = self.var
        der_list = np.array([var_i.partials() for var_i in inputs])
        return var_val, der_list

    def partials(self):
        """
        Method to compute derivative for variables.
        Uses self.derivative to determine whether to use list comprehension.
        For finding partial derivatives with respect to each function.
        """
        if len(self.child_node) == 0:
            return 1
        if self.derivative is not None:
            return self.derivative
        else:
            self.derivative = sum([child.partials() * partial for child, partial in self.child_node])
            return self.derivative

    def __add__(self, other):
        """
        dunder method for adding a Node object.
        internally appends derivatives of self and other and the return object to the .child_node.

        :param self: Node object.
        :param other: Node object, or int/float, to be added to self.var.
        :return: Node object with value of the sum of self and other.
        """
        try:
            new_add = Node(self.var + other.var)
            self.child_node.append((new_add, 1))
            other.child_node.append((new_add, 1))

            return new_add
        except:
            if isinstance(other, int) or isinstance(other, float):
                # other is not a Node and the addition could complete if it is a real number
                new_add = Node(self.var + other)
                self.child_node.append((new_add, 1))
                return new_add
            else:
                raise TypeError("Input is not a real number.")

    def __mul__(self, other):
        """
        dunder method for multiplying a Node object or constant.
        internally appends derivatives of self and other and the return object to the .child_node.

        :param self: Node object.
        :param other: Node object, or int/float, to be multiplied to self.var.
        :return: Node object with value of the product of self and other.
        """
        try:
            new_mul = Node(other.var * self.var)
            self.child_node.append((new_mul, other.var))
            other.child_node.append((new_mul, self.var))
            return new_mul
        except:
            if isinstance(other, int) or isinstance(other, float):
                # other is not a Node and the multiplication could complete if it is a real number
                new_mul = Node(other * self.var)
                self.child_node.append((new_mul, other))
                return new_mul
            else:
                raise TypeError("Input is not a real number.")

    def __radd__(self, other):
        """
        dunder method for adding a Node object and left object without __add__ mehtod.

        :param self: Node object.
        :param other: an object that does not have an __add__ method or not implemented.
        :return: __add__ function call.
        """
        return self.__add__(other)

    def __rmul__(self, other):
        """
        dunder method for multiplying a Node object and left object without __mul__ mehtod.

        :param self: Node object.
        :param other: an object that does not have an __mul__ method or not implemented.
        :return: __mul__ function call.
        """
        return self.__mul__(other)

    def __sub__(self, other):
        """
        dunder method for subtracting a Node object or constant.

        :param self: Node object.
        :param other: Node object, or int/float, to be subtracted from self.var.
        :return: __add__ function call that passes the other object with negative sign.
        """
        return self.__add__(-other)

    def __rsub__(self, other):
        """
        dunder method for subtracting a Node object and left object without __sub__ mehtod.

        :param self: Node object.
        :param other: an object that does not have an __sub__ method or not implemented.
        :return: __add__ function call that passes the self object with negative sign.
        """
        return (-self).__add__(other)

    def __truediv__(self, other):
        """
        dunder method for dividing a Node object or constant.
        internally appends derivatives of self and other and the return object to the .child_node.

        :param self: Node object.
        :param other: Node object, or int/float, to be divided from self.var.
        :return: Node object with value and the fraction of self and other.
        """
        try:
            new_div = Node(self.var / other.var)
            self.child_node.append((new_div,((1 * other.var - 0 * self.var) / other.var**2)))
            other.child_node.append((new_div, (-self.var/(other.var**2))))
            return new_div
        except AttributeError:
            if isinstance(other, int) or isinstance(other, float):
                new_div = Node(self.var / other)
                self.child_node.append((new_div,((1 * other - 0 * self.var) / other**2)))
                return new_div
            else:
                raise TypeError(f"Input {other} is not valid.")

    def __neg__(self):
        """
        dunder method for taking the negative of an Node object or constant.
        internally appends derivatives of self and other and the return object to the .child_node.

        :param self: Node object.
        :return: Node object with negative value of itself.
        """
        new_neg = Node(-self.var)
        self.child_node.append((new_neg, -1))
        return new_neg

    def __rtruediv__(self, other):
        """
        dunder method for dividing a Node object and left other object without __truediv__ method.
        internally appends derivatives of self and other and the return object to the .child_node.

        :param self: Node object.
        :param other: an object that does not have an __truediv__ method or not implemented.
        :return: Node object with value of the fraction of self and other.
        """
        try:
            new_div = Node(other.var / self.var)
            self.child_node.append((new_div, ((0 * self.var - other.var * 1) / self.var**2)))
            other.child_node.append((new_div, 1/self.var))
            return new_div
        except:
            if isinstance(other, int) or isinstance(other, float):
                new_div = Node(other / self.var)
                self.child_node.append((new_div, ((0 * self.var - other * 1) / self.var**2)))
                return new_div
            else:
                raise TypeError(f"Input {other} is not valid.")

    def __lt__(self, other):
        '''
        dunder method for less than comparator.

        :param self: Node object.
        :param other: Node object, or int/float, to be compared with.
        :return: True if self is less than other, otherwise False.
        '''
        try:
            return self.var < other.var
        except AttributeError:
            return self.var < other

    def __gt__(self, other):
        '''
        dunder method for greater than comparator.

        :param self: Node object.
        :param other: Node object, or int/float, to be compared with.
        :return: True if self is greater than other, otherwise False.
        '''
        try:
            return self.var > other.var
        except AttributeError:
            return self.var > other

    def __le__(self, other):

        '''
        dunder method for less than or equal to comparator.

        :param self: Node object.
        :param other: Node object, or int/float, to be compared with.
        :return: True if self is less than or equal to other, otherwise False.
        '''

        try:
            return self.var <= other.var
        except AttributeError:
            return self.var <= other

    def __ge__(self, other):

        '''
        dunder method for greater than or equal to comparator.

        :param self: Node object.
        :param other: Node object, or int/float, to be compared with.
        :return: True if self is greater than or equal to other, otherwise False.
        '''
        try:
            return self.var >= other.var
        except AttributeError:
            return self.var >= other

    def __eq__(self, other):
        '''
        dunder method for equality comparator.

        :param self: Node object.
        :param other: Node object, or int/float, to be compared with.
        :return: True if self is equal to other, otherwise False.
        '''
        try:
            return self.var == other.var
        except:
            raise TypeError('Input is not comparable.')

    def __ne__(self, other):
        '''
        dunder method for inequality comparator.

        :param self: Node object.
        :param other: Node object, or int/float, to be compared with.
        :return: negation of __eq__ function call.
        '''
        return not self.__eq__(other)

    def __abs__(self):

        '''
        dunder method for absolute value.
        internally appends derivatives of self and other and the return object to the .child_node.

        :param self: Node object.
        :return: Node object with value of the absolute value of self.
        '''
        new_abs = Node(abs(self.var))
        self.child_node.append((1, new_abs))
        return new_abs

    def __pow__(self, other):

        """
        dunder method for taking Variable object to the value of other object's power.
        internally appends derivatives of self and other and the return object to the .child_node.

        :param self: Node object.
        :param other: Node object, or int/float, to the power of.
        :return: Node object with value of the power of self with power of other's value.
        """

        try:
            new_val = Node(self.var ** other.var)
            self.child_node.append((new_val, (other.var) * self.var ** (other.var-1)))
            other.child_node.append((new_val, self.var ** other.var * (np.log(self.var))))
            return new_val
        except:
            if isinstance(other, int):
                new_val = Node(self.var ** other)
                self.child_node.append((new_val, (other) * self.var ** (other-1)))
                return new_val
            else:
                raise TypeError(f"Exponent {other} is not valid.")

    def __rpow__(self, other):

        """
        dunder method for taking the other object to the Node object's power.
        internally appends derivative of self and the return object to self.child_node.

        :param self: Node object.
        :param other: Node object, or int/float, as the base.
        :return: Node object with value of the power of other with power of self's value.
        """
        try:
            new_val = Node(other ** self.var)
        except:
            raise ValueError("{} must be a number.".format(other))
        self.child_node.append((new_val, other**self.var * np.log(other)))
        return new_val

    @staticmethod
    def log(var):
        """
        takes the log of the Node object.
        internally appends derivative of var and the return object to var.child_node.

        :param var: Node object.
        :return: Node object with value of the natural log of var.var
        """
        try:
            if var.var <= 0:
                raise ValueError('Input needs to be greater than 0.')
        except:
            raise TypeError(f"Input not valid.")
        log_var = Node(np.log(var.var))
        var.child_node.append((log_var, (1. / var.var) * 1))
        return log_var

    @staticmethod
    def sqrt(var):
        """
        takes the square root of the Node object.
        internally appends derivative of var and the return object to var.child_node.

        :var: Node object.
        :return: Node object with value of the square root of var.
        """
        if var < 0:
            raise ValueError("Square root can only takes positive values.")
        else:
            try:
                sqrt_var = Node(var.var**(1/2))
                var.child_node.append((sqrt_var, (1/2)*var.var**(-1/2)))
            except:
                raise TypeError(f"Input is not an Node object.")
        return sqrt_var

    @staticmethod
    def exp(var):
        """
        natural exponential function with the value of Node object as the power.
        internally appends derivative of var and the return object to var.child_node.

        :var: Node object.
        :return: Node object with value of natural exponential function with power var.
        """
        try:
            new_val = Node(np.exp(var.var))
            var.child_node.append((new_val, np.exp(var.var) * 1))
            return new_val
        except:
            if not isinstance(var, int) and not isinstance(var, float):
                raise TypeError(f"Input {var} is not valid.")

            return np.exp(var)

    @staticmethod
    def sin(var):
        """
        calculates the sine of Node object.
        internally appends derivative of var and the return object to var.child_node.

        :var: Node object.
        :return: Node object with value of the sine of var.
        """
        try:
            new_val = Node(np.sin(var.var))
            var.child_node.append((new_val, 1 * np.cos(var.var)))
            return new_val
        except:
            if not isinstance(var, int) and not isinstance(var, float):
                raise TypeError(f"Input {var} is not valid.")

            return np.sin(var)

    @staticmethod
    def cos(var):
        """
        calculates the cosine of Node object.
        internally appends derivative of var and the return object to var.child_node

        :var: Node object.
        :return: Node object with value of the cosine of var.
        """
        try:
            new_val = Node(np.cos(var.var))
            var.child_node.append((new_val, 1 * -np.sin(var.var)))
            return new_val
        except:
            if not isinstance(var, int) and not isinstance(var, float):
                raise TypeError(f"Input {var} is not valid.")

            return np.cos(var)

    @staticmethod
    def tan(var):
        """
        calculates the tangent of Node object.
        internally appends derivative of var and the return object to var.child_node.

        :var: Node object.
        :return: Node object with value of the tangent of var.
        """
        try:
            new_val = Node(np.tan(var.var))
            var.child_node.append((new_val, 1 * 1 / np.power(np.cos(var.var), 2)))
            return new_val
        except:
            if not isinstance(var, int) and not isinstance(var, float):
                raise TypeError(f"Input {var} is not valid.")

            return np.tan(var)

    @staticmethod
    def arcsin(var):
        """
        calculates the arcsine of Node object.
        internally appends derivative of var and the return object to var.child_node.

        :var: Node object.
        :return: Node object with value of the arcsine of var.
        """
        try:
            if var.var > 1 or var.var < -1:
                raise ValueError('Please input -1 <= x <=1')
            else:
                new_val = Node(np.arcsin(var.var))
                var.child_node.append((new_val, 1 / np.sqrt(1 - (var.var ** 2))))
                return new_val
        except:
            if not isinstance(var, int) and not isinstance(var, float):
                raise TypeError(f"Input {var} is not valid.")
            return np.arcsin(var)

    @staticmethod
    def arccos(var):
        """
        calculates the arccosine of Node object.
        internally appends derivative of var and the return object to var.child_node

        :var: Node object.
        :return: Node object with value of the arccosine of var.
        """

        try:
            if isinstance(var, int) or isinstance(var, float):
                return np.arccos(var)

            if var.var > 1 or var.var < -1:
                raise ValueError('Please input -1 <= x <=1')
            else:
                new_val = Node(np.arccos(var.var))
                var.child_node.append((new_val, -1 / np.sqrt(1 - (var.var ** 2))))
            return new_val
        except:
                raise TypeError(f"Input {var} is not valid.")

    @staticmethod
    def arctan(var):
        """
        calculates the arctangent of Node object.
        internally appends derivative of var and the return object to var.child_node

        :var: Node object.
        :return: Node object with value of the arctangent of var.
        """
        try:
            new_val = Node(np.arctan(var.var))
            var.child_node.append((new_val, 1 * 1 / (1 + np.power(var.var, 2))))

            return new_val

        except AttributeError:
            return np.arctan(var)

    @staticmethod
    def sinh(var):
        """
        calculates the hyperbolic sine of Node object.
        internally appends derivative of var and the return object to var.child_node

        :var: Node object.
        :return: Node object with value and derivative of the hyperbolic sine of var.
        """
        try:
            new_val = Node(np.sinh(var.var))
            var.child_node.append((new_val, 1 * np.cosh(var.var)))
            return new_val

        except AttributeError:
            return np.sinh(var)

    @staticmethod
    def cosh(var):
        """
        calculates the hyperbolic cosine of Node object.
        internally appends derivative of var and the return object to var.child_node

        :param var: Node object.
        :return: Node object with value and derivative of the hyperbolic cosine of var.
        """
        try:
            new_val = Node(np.cosh(var.var))
            var.child_node.append((new_val, 1 * np.sinh(var.var)))

            return new_val

        except AttributeError:
            return np.cosh(var)

    @staticmethod
    def tanh(var):
        """
        calculates the hyperbolic tangent of Node object.
        internally appends derivative of var and the return object to var.child_node

        :param var: Node object.
        :return: Node object with value and derivative of the hyperbolic cosine of var.
        """
        try:
            new_val = Node(np.tanh(var.var))
            var.child_node.append((new_val, 1 * 1 / np.power(np.cosh(var.var), 2)))
            return new_val
        except AttributeError:
            return np.tanh(var)

    def sigmoid(var):
        """
        calculates the sigmoid/logistic function of Node object.
        internally appends derivative of var and the return object to var.child_node

        :param var: Node object.
        :return: Node object with value of the sigmoid/logistic function of var.
        """
        try:
            logistic_var = Node(1 / (1 + np.exp(-var.var)))
            var.child_node.append((logistic_var, 1 / (1 + np.exp(-var.var)) * (1-(1 / (1 + np.exp(-var.var)) * 1))))
            return logistic_var
        except:
            raise TypeError(f"Input {var} not valid.")

    def __str__(self):
        return f"value = {self.var}, derivative = {self.partials()}"

    def __repr__(self):
        return f"value = {self.var}, derivative = {self.partials()}"

class Reverse:
    def __init__(self, dict_val, list_funct):
        """
        Inputs:

        dict_val: a dictionary object
        dict_val represents the input which is a dictionary of variables and their values formatted like the examples below
        example1={'x':7,'y':2, 'z':3}
        example2 ={'x':4}
        You can input as many variables as needed so long as it is more than 1

        list_funct: a list object
        list_funct represents the input which is a list of functions input as strings
        example_a = ['x**2','cos(np.pi*y)', '8*z']
        example_b = ['sin(x)']

        Demo Reverse Mode
        ______________________________________________________________
        INPUT

        dict_val = {'x': 3, 'y': 2, 'z':1}
        list_funct = ['x * y + exp(x * y)+ z**2', 'x + 3 * y + 4*x*z']
        reverse_out = Reverse(dict_val, list_funct)
        print(reverse_out)

        OUTPUT

        ---Reverse Differentiation---
        Function 1:
        Expression = x * y + exp(x * y)+ z**2
        Value = 410.4287934927351
        Gradient = [ 808.85758699 1213.28638048    2.        ]

        Function 2:
        Expression = x + 3 * y + 4*x*z
        Value = 21.0
        Gradient = [ 5.  3. 12.]

        """
        # checking for string type
        for func in list_funct:
            if not isinstance(func, str):
                raise TypeError('Invalid function input, must be string.')
        # checking for dictionary type
        if not isinstance(dict_val, dict):
            raise TypeError('Variable Input must be a dictionary')

        self.var = []
        self.der = []
        self.list_funct = list_funct

        static_elem_funct = ['log', 'sqrt', 'exp', 'sin', 'cos', 'tan', 'arcsin', 'arccos', 'arctan', 'sinh', 'cosh', 'tanh', 'sigmoid']

        for func in list_funct:
            for i in static_elem_funct:
                if i in func:
                    func = re.sub(i + r'\(', 'Node.' + i + '(', func)
                    func = re.sub('arcNode.', 'arc', func)

            for var_name, var_value in dict_val.items():
                exec(f'{var_name} = Node(float(var_value))')

            func_eval = eval(func)

            value_keys = str(list(dict_val.keys())).replace('\'','')
            val_1, der_1 = eval(f'func_eval.get_derivatives({value_keys})')

            self.var.append(val_1)
            self.der.append(der_1)

    def __repr__(self):
        output = '---Reverse Differentiation---\n'
        added_output = ''
        for i in range(0, len(self.list_funct)):
            added_output += f"Function {i+1}: \nExpression = {self.list_funct[i]}\nValue = {str(self.var[i])}\nGradient = {str(self.der[i])}\n\n"

        return output+added_output

    def __str__(self):
        output = '---Reverse Differentiation---\n'
        added_output = ''
        for i in range(0, len(self.list_funct)):
            added_output += f"Function {i+1}: \nExpression = {self.list_funct[i]}\nValue = {str(self.var[i])}\nGradient = {str(self.der[i])}\n\n"

        return output+added_output


import numpy as np
import math
from scipy.stats import entropy

#------------------------------------------------------------------
def cosine(x, y):
	"""
    Compute the cosine similarity between two vectors x and y.
	"""
	x=normalize(x)
	y=normalize(y)
	return np.dot(x,y.T)
#------------------------------------------------------------------
def normalize(x):
    """
    L2 normalize vector x.
    """
    norm_x = np.linalg.norm(x)
    return x if norm_x == 0 else (x / norm_x)
#------------------------------------------------------------------
def softmax_fun(x):
	e_x=np.exp(x - np.max(x))
	return e_x / np.sum(e_x)
#------------------------------------------------------------------
def sigmoid_fun(x):
	return 1/(1+np.exp(-x))
#------------------------------------------------------------------
def KL_divergence(a,b):
	return np.sum(a*np.log(a/b))
#------------------------------------------------------------------
def JSD(P, Q):
	M = 0.5 * (P + Q)
	return 0.5 * (entropy(P, M) + entropy(Q, M))
#------------------------------------------------------------------
def pdf_UnivariateGauss(x,mean,var):
	pi = 3.1415926
	denom=np.sqrt(2*pi*var)
	num= np.exp(-(x-mean)**2/(2*var))
	return num/denom
#------------------------------------------------------------------
def ReLU(x):
    return abs(x) * (x > 0)
#------------------------------------------------------------------


# Simulate an inhomogeneous Poisson point process on a rectangle
# This is done by simulating a homogeneous Poisson process, which is
# then thinned according to a (spatially *dependent*) p-thinning.
# The intensity function is
# lambda(x,y)=80*exp(-(x^2+y^2)/s^2)+100*exp(-(x^2+y^2)/s^2), where s>0.
# The simulations are then checked by examining the statistics of the number of
# points and the locations of points (using histograms)
# Author: H. Paul Keeler, 2019.
# Website: hpaulkeeler.com
# Repository: github.com/hpaulkeeler/posts
# For more details, see the post:
# https://hpaulkeeler.com/checking-poisson-point-process-simulations/

import numpy as np;  # NumPy package for arrays, random number generation, etc
import matplotlib.pyplot as plt  # for plotting
from matplotlib import cm  # for heatmap plotting
from mpl_toolkits import mplot3d  # for 3-D plots
from scipy.optimize import minimize  # for optimizing
from scipy import integrate  # for integrating
from scipy.stats import poisson  # for the Poisson probability mass function

plt.close('all');  # close all plots

# Simulation window parameters
xMin = -1;
xMax = 1;
yMin = -1;
yMax = 1;
xDelta = xMax - xMin;
yDelta = yMax - yMin;  # rectangle dimensions
areaTotal = xDelta * yDelta;

numbSim = 10 ** 4;  # number of simulations
numbBins = 30;  # number of bins for histogram

# Point process parameters
s = 0.5;  # scale parameter

def fun_lambda(x, y):
    # intensity function
    lambdaValue = 80 * np.exp(-((x + 0.5) ** 2 + (y + 0.5) ** 2) / s ** 2) + 100 * np.exp(
        -((x - 0.5) ** 2 + (y - 0.5) ** 2) / s ** 2);
    return lambdaValue;

###START -- find maximum lambda -- START ###
# For an intensity function lambda, given by function fun_lambda,
# finds the maximum of lambda in a rectangular region given by
# [xMin,xMax,yMin,yMax].
def fun_Neg(x):
    return -fun_lambda(x[0], x[1]);  # negative of lambda

xy0 = [(xMin + xMax) / 2, (yMin + yMax) / 2];  # initial value(ie centre)
# Find largest lambda value
resultsOpt = minimize(fun_Neg, xy0, bounds=((xMin, xMax), (yMin, yMax)));
lambdaNegMin = resultsOpt.fun;  # retrieve minimum value found by minimize
lambdaMax = -lambdaNegMin;

###END -- find maximum lambda -- END ###

# define thinning probability function
def fun_p(x, y):
    return fun_lambda(x, y) / lambdaMax;

# for collecting statistics -- set numbSim=1 for one simulation
numbPointsRetained = np.zeros(numbSim);  # vector to record number of points
xxAll = [];
yyAll = [];

### START -- Simulation section -- START ###
for ii in range(numbSim):
    # Simulate a Poisson point process
    numbPoints = np.random.poisson(lambdaMax * areaTotal);  # Poisson number of points
    xx = xDelta * np.random.uniform(0, 1, numbPoints) + xMin;  # x coordinates of Poisson points
    yy = yDelta * np.random.uniform(0, 1, numbPoints) + yMin;  # y coordinates of Poisson points

    # calculate spatially-dependent thinning probabilities
    p = fun_p(xx, yy);

    # Generate Bernoulli variables (ie coin flips) for thinning
    booleRetained = np.random.uniform(0, 1, numbPoints) < p;  # points to be thinned

    # x/y locations of retained points
    xxRetained = xx[booleRetained];
    yyRetained = yy[booleRetained];
    numbPointsRetained[ii] = xxRetained.size;
    xxAll.extend(xxRetained);
    yyAll.extend(yyRetained);
### END -- Simulation section -- END ###

# Plotting a simulation
fig1 = plt.figure();
plt.scatter(xxRetained, yyRetained, edgecolor='b', facecolor='none');
plt.xlabel('x');
plt.ylabel('y');
plt.title('A single realization of a Poisson point process');
plt.show();

# run empirical test on number of points generated
###START -- Checking number of points -- START###
# total mean measure (average number of points)
LambdaNumerical = integrate.dblquad(fun_lambda, xMin, xMax, lambda x: yMin, lambda y: yMax)[0];
# Test: as numbSim increases, numbPointsMean converges to LambdaNumerical
numbPointsMean = np.mean(numbPointsRetained);
# Test: as numbSim increases, numbPointsVar converges to LambdaNumerical
numbPointsVar = np.var(numbPointsRetained);
binEdges = np.arange(numbPointsRetained.min(), (numbPointsRetained.max() + 1)) - 0.5;
pdfEmp, binEdges = np.histogram(numbPointsRetained, bins=binEdges, density=True);

nValues = np.arange(numbPointsRetained.min(), numbPointsRetained.max());
# analytic solution of probability density
pdfExact = (poisson.pmf(nValues, LambdaNumerical));

# Plotting
fig2 = plt.figure();
plt.scatter(nValues, pdfExact, color='b', marker='s', facecolor='none', label='Exact');
plt.scatter(nValues, pdfEmp, color='r', marker='+', label='Empirical');
plt.autoscale(enable=True, axis='y', tight=True)
plt.xlabel('n');
plt.ylabel('P(N=n)');
plt.title('Distribution of the number of points');
plt.legend();
plt.show();
###END -- Checking number of points -- END###

###START -- Checking locations -- START###
# 2-D Histogram section
p_Estimate, xxEdges, yyEdges = np.histogram2d(xxAll, yyAll, bins=numbBins, density=True);
lambda_Estimate = p_Estimate * numbPointsMean;

xxValues = (xxEdges[1:] + xxEdges[0:xxEdges.size - 1]) / 2;
yyValues = (yyEdges[1:] + yyEdges[0:yyEdges.size - 1]) / 2;
X, Y = np.meshgrid(xxValues, yyValues)  # create x/y matrices for plotting

# analytic solution of probability density
lambda_Exact = fun_lambda(X, Y);

# Plot empirical estimate
fig3 = plt.figure();
plt.rc('text', usetex=True);
plt.rc('font', family='serif');
ax = plt.subplot(211, projection='3d');
surf = ax.plot_surface(X, Y, lambda_Estimate, cmap=plt.cm.plasma);
plt.xlabel('x');
plt.ylabel('y');
plt.title('Estimate of $\lambda(x)$');
plt.locator_params(axis='x', nbins=3);
plt.locator_params(axis='y', nbins=3);
plt.locator_params(axis='z', nbins=3);
# Plot exact expression
ax = plt.subplot(212, projection='3d');
surf = ax.plot_surface(X, Y, lambda_Exact, cmap=plt.cm.plasma);
plt.xlabel('x');
plt.ylabel('y');
plt.title('True $\lambda(x)$');
plt.locator_params(axis='x', nbins=3);
plt.locator_params(axis='y', nbins=3);
plt.locator_params(axis='z', nbins=3);
###END -- Checking locations -- END###


# Illustrate einstein summation
# https://rockt.github.io/2018/04/30/einsum
# https://docs.scipy.org/doc/numpy/reference/generated/numpy.einsum.html

import superimport

import numpy as np

np.random.seed(42)
a = np.arange(3)
b = np.arange(3)
A = np.arange(6).reshape(2,3)
B = np.arange(15).reshape(3,5)
S = np.arange(9).reshape(3,3)
T = np.random.randn(2,2,2,2)

## Single argument

# Matrix transpose
assert np.allclose(A.T, np.einsum('ij->ji', A))

# Sum all elements
assert np.allclose(np.sum(A), np.einsum('ij->', A))

# Sum across rows
assert np.allclose(np.sum(A, axis=0), np.einsum('ij->j', A))

# Sum across columns
assert np.allclose(np.sum(A, axis=1), np.einsum('ij->i', A))

# Sum specific axis of tensor
assert np.allclose(np.sum(T, axis=1), np.einsum('ijkl->ikl', T))
assert np.allclose(np.sum(np.sum(T, axis=0), axis=0), np.einsum('ijkl->kl', T))

# repeated indices with one arg extracts diagonals
assert np.allclose(np.diag(S), np.einsum('ii->i', S))

# Trace
assert np.allclose(np.trace(S), np.einsum('ii->', S))

## Two arguments

# Matrix vector multiplication
assert np.allclose(np.dot(A, b), np.einsum('ik,k->i', A, b))

# Matrix matrix multiplication
assert np.allclose(np.dot(A, B), np.einsum('ik,kj->ij', A, B))
assert np.allclose(np.matmul(A, B), np.einsum('ik,kj->ij', A, B))

# Inner product
assert np.allclose(np.dot(a, b), np.einsum('i,i->', a, b))
assert np.allclose(np.inner(a, b), np.einsum('i,i->', a, b))

# Outer product
assert np.allclose(np.outer(a, b), np.einsum('i,j->ij', a, b))

# Elementwise product
assert np.allclose(a * a, np.einsum('i,i->i', a, a))
assert np.allclose(A * A, np.einsum('ij,ij->ij', A, A))
assert np.allclose(np.multiply(A, A), np.einsum('ij,ij->ij', A, A))

# Batch matrix multiplication
I= 3; J = 2; K = 5; L = 3;
AA = np.random.randn(I,J,K)
BB = np.random.randn(I,K,L)
# C[ijl] = sum_k A[ijk] B[ikl]
CC = np.zeros((I,J,L))
for i in range(I):
    for j in range(J):
        for l in range(L):
            s = 0
            for k in range(K):
                s += AA[i,j,k] * BB[i,k,l]
            CC[i,j,l] = s
assert np.allclose(CC, np.einsum('ijk,ikl->ijl', AA, BB))

## >2 arguments

# Batch sentence embedding and averaging
N = 2; C = 3; D = 4; K = 5; T = 6;
S = np.random.randn(N, T, K)
W = np.random.randn(K, D)
V = np.random.randn(D, C)
L = np.zeros((N,C))
for n in range(N):
    for c in range(C):
        s = 0
        for d in range(D):
            for k in range(K):
                for t in range(T):
                    s += S[n,t,k] * W[k,d] * V[d,c]
        L[n,c] = s
assert np.allclose(L, np.einsum('ntk,kd,dc->nc', S, W, V))

path = np.einsum_path('ntk,kd,dc->nc', S, W, V, optimize='optimal')[0]
assert np.allclose(L, np.einsum('ntk,kd,dc->nc', S, W, V, optimize=path))

import jax.numpy as jnp
path = jnp.einsum_path('ntk,kd,dc->nc', S, W, V, optimize='optimal')[0]
assert np.allclose(L, jnp.einsum('ntk,kd,dc->nc', S, W, V, optimize=path))

# Use full student network from KOller and Friedman
str = 'c,dc,gdi,si,lg,jls,hgj->'
K = 5
cptC = np.random.randn(K)
cptD = np.random.randn(K,K)
cptG = np.random.randn(K,K,K)
cptS = np.random.randn(K,K)
cptL = np.random.randn(K,K)
cptJ = np.random.randn(K,K,K)
cptH = np.random.randn(K,K,K)
cpts = [cptC, cptD, cptG, cptS, cptL, cptJ, cptH]
path_info = np.einsum_path(str, *cpts, optimize='optimal')
print(path_info[0]) # 'einsum_path', (0, 1), (0, 5), (0, 4), (0, 3), (0, 2), (0, 1)]
print(path_info[1])
'''
  Complete contraction:  c,dc,gdi,si,lg,jls,hgj->
         Naive scaling:  8
     Optimized scaling:  4
      Naive FLOP count:  2.734e+06
  Optimized FLOP count:  2.176e+03
   Theoretical speedup:  1256.606
  Largest intermediate:  1.250e+02 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
   2                     dc,c->d                    gdi,si,lg,jls,hgj,d->
   3                   d,gdi->gi                       si,lg,jls,hgj,gi->
   3                   gi,si->gs                          lg,jls,hgj,gs->
   3                  gs,lg->gls                            jls,hgj,gls->
   4                 gls,jls->gj                                 hgj,gj->
   3                    gj,hgj->                                       ->
'''

path_info = np.einsum_path(str, *cpts, optimize='greedy')
print(path_info[1])
'''
  Complete contraction:  c,dc,gdi,si,lg,jls,hgj->
         Naive scaling:  8
     Optimized scaling:  5
      Naive FLOP count:  2.734e+06
  Optimized FLOP count:  7.101e+03
   Theoretical speedup:  385.069
  Largest intermediate:  1.250e+02 elements
--------------------------------------------------------------------------
scaling                  current                                remaining
--------------------------------------------------------------------------
   5                hgj,jls->gls                     c,dc,gdi,si,lg,gls->
   3                  gls,lg->gs                         c,dc,gdi,si,gs->
   2                     dc,c->d                            gdi,si,gs,d->
   3                   d,gdi->gi                               si,gs,gi->
   3                   gs,si->gi                                  gi,gi->
   2                     gi,gi->                                       ->
'''


#!/usr/bin/env python3
# -*- coding: utf-8 -*-

""" Documentation

"""
__author__ = "R. Adda, M.Charpentier, E. Vazquez"
__copyright__ = "CentraleSupelec, 2020"
__license__ = "MIT"
__maintainer__ = "E. Vazquez"
__email__ = "emmanuel.vazquez@centralesupelec.fr"
__status__ = "alpha"

import math

import matplotlib.pyplot as plt
import numpy as np
from scipy.integrate import ode

def SEIRD(t, x, theta, const_params):
    # Computes...
    #
    # Inputs:
    #   t
    #   x
    #   theta = [R, T_inf, T_inc , pfatal]
    #   ...
    # Output:
    #   dx
    # ...
    # ...

    R0 = theta[0]    # Basic Reproduction Rate
    Tinf = theta[1]  # Infection Time
    Tinc = theta[2]  # Incubation Time
    pfatal = math.pow(10, theta[3])  # Death proportion for I compartment

    N = const_params[0]
    Gamma = const_params[1]
    mu = const_params[2]

    S = x[0]
    E = x[1]
    I = x[2]
    R = x[3]
    D = x[4]

    gamma = 1 / Tinf
    a = 1 / Tinc
    beta = gamma * R0

    dS = Gamma - mu * S - beta / N * I * S
    dE = beta / N * I * S - (mu + a) * E
    dI = a * E - gamma * (1 + pfatal) * I
    dR = gamma * I - mu * R
    dD = gamma * pfatal * I

    dx = [dS, dE, dI, dR, dD]

    return dx

def SEIRD_with_cutoff(t, x, theta, const_params):
    # Computes...
    #
    # Inputs:
    #   t
    #   x
    #   theta = [R, beta_cut, T_inf, T_inc , pfatal]
    #   ...
    # Output:
    #   dx
    # ...
    # ...

    R0 = theta[0]    # Basic Reproduction Rate
    beta_cut = theta[1]
    Tinf = theta[2]  # Infection Time
    Tinc = theta[3]  # Incubation Time
    pfatal = pow(10, theta[4])  # Death proportion for I compartment

    N = const_params[0]
    Gamma = const_params[1]
    mu = const_params[2]
    cutoff_time = const_params[3]
    lift_time = const_params[4]

    S = x[0]
    E = x[1]
    I = x[2]
    R = x[3]
    D = x[4]

    gamma = 1 / Tinf
    a = 1 / Tinc
    if t >= cutoff_time and t < lift_time:
        beta = beta_cut * gamma * R0
    else:
        beta = gamma * R0

    dS = Gamma - mu * S - beta / N * I * S
    dE = beta / N * I * S - (mu + a) * E
    dI = a * E - gamma * (1 + pfatal) * I
    dR = gamma * I - mu * R
    dD = gamma * pfatal * I

    dx = [dS, dE, dI, dR, dD]

    return dx

def SIRD_with_cutoff(t, x, theta, const_params):
    # Computes...
    #
    # Inputs:
    #   t
    #   x
    #   theta = [R, beta_cut, T_inf, T_inc , pfatal]
    #   ...
    # Output:
    #   dx
    # ...
    # ...

    R0 = theta[0]               # Basic Reproduction Rate
    beta_cut = theta[1]
    Tinf = theta[2]             # Infection Time
    pfatal = pow(10, theta[3])  # Death proportion for I compartment

    N = const_params[0]
    Gamma = const_params[1]
    mu = const_params[2]
    cutoff_time = const_params[3]
    lift_time = const_params[4]

    S = x[0]
    I = x[1]
    R = x[2]
    D = x[3]

    gamma = 1 / Tinf
    if t >= cutoff_time and t < lift_time:
        beta = beta_cut * gamma * R0
    else:
        beta = gamma * R0

    dS = Gamma - mu * S - beta / N * I * S
    dI = beta / N * I * S - mu * I - gamma * (1 + pfatal) * I
    dR = gamma * I - mu * R
    dD = gamma * pfatal * I

    dx = [dS, dI, dR, dD]

    return dx

class model:
    '''
    Attributes
        model_type :
        f : RHS of the ODE
        ...
        ...
        theta : parameters
        t : time axis for simulation
        x : [S (Susceptible), E (Exposed), I (Infectious), R (Recovered), D (Death)] along time

    Methods
        simulate
        ...
        plot
    '''

    def __init__(self, config, tobs):
        # Initialization
        # 1. Model choice
        self.set_model_type(config)

        # 2. set regional parameters
        self.set_regional_params(config)

        # 3. simulation parameters
        self.set_simulation_params(config)

        # 4. build t and preallocate x
        self.prepare_simulations(config, tobs)

    def set_model_type(self, config):
        # Model type
        self.model_type = config['model_type']

        if self.model_type == 'SEIRD':
            self.f = SEIRD
            self.state_dim = 5  # state dim
            self.state_ref = {'S': 0, 'E':1, 'I':2, 'R':3, 'D':4}
            self.theta_dim = 5  # theta dim

        elif self.model_type == 'SEIRD_with_cutoff':
            self.f = SEIRD_with_cutoff
            self.state_dim = 5  # state dim
            self.state_ref = {'S': 0, 'E':1, 'I':2, 'R':3, 'D':4}
            self.theta_dim = 6  # theta dim

        elif self.model_type == 'SIRD_with_cutoff':
            self.f = SIRD_with_cutoff
            self.state_dim = 4  # state dim
            self.state_ref = {'S': 0, 'I':1, 'R':2, 'D':3}
            self.theta_dim = 5  # theta dim

        else:
            raise('model type is not implemented')

    def set_regional_params(self, config):
        # Regional parameters

        if self.model_type == 'SEIRD':
            self.regional_params = [config['N'], config['Gamma'], config['mu']]

        elif self.model_type == 'SEIRD_with_cutoff':
            self.regional_params = [config['N'], config['Gamma'], config['mu'],
                                    config['cutoff_time'],
                                    config['lift_time']]

        elif self.model_type == 'SIRD_with_cutoff':
            self.regional_params = [config['N'], config['Gamma'], config['mu'],
                                    config['cutoff_time'],
                                    config['lift_time']]

    def set_simulation_params(self, config):
        # Simulation parameters

        self.time_step = config['sim_step']
        self.sim_duration = config['sim_duration']

    def prepare_simulations(self, config, tobs):
        # build t and preallocate x

        self.t = np.arange(0, self.sim_duration, self.time_step)
        self.x = np.zeros([self.t.shape[0], self.state_dim])  # preallocation

        # start date, assume a a datetime object for t0_ref and convert to numpy datetime object
        self.t0_refdate = np.datetime64(config['t0_refdate'])

        # assume a pandas dataframe of timestamp objects and convert to numpy datetime objects
        tobs_rel = [tobs[i] - self.t0_refdate for i in range(tobs.shape[0])]
        # pdb.set_trace()

        # build a tobs_idx array in order to retrieve observations with self.x[self.tobs_idx]
        self.tobs_rel = [tobs_rel[i].days for i in range(len(tobs_rel))]
        k = 0
        self.tobs_idx = []
        for i in range(len(self.tobs_rel)):
            t = self.tobs_rel[i]
            while k < self.t.shape[0]:
                if self.t[k] >= t:
                    self.tobs_idx.append(k)
                    break
                k += 1
        # pdb.set_trace()

    def y_from_tobs(self):
        # Helper function to retrieve fatalities at tobs

        fatalities_idx = self.state_ref['D']
        return self.x[self.tobs_idx, fatalities_idx]

    def simulate(self, theta):
        # ODE integration happens here

        # Set theta
        self.theta = theta

        # State initialization
        t0 = theta[-1]
        if self.model_type in ['SEIRD', 'SEIRD_with_cutoff']:
            x0 = [self.regional_params[0] - 1, 0, 1, 0, 0]
        elif self.model_type == 'SIRD_with_cutoff':
            x0 = [self.regional_params[0] - 1, 1, 0, 0]

        # ODE initialization
        # r = ode(f).set_integrator('vode', method='adams', with_jacobian=False)
        r = ode(self.f).set_integrator('lsoda', with_jacobian=False)
        r.set_initial_value(x0, t0).set_f_params(self.theta, self.regional_params)

        idx_start = 0
        while self.t[idx_start] < t0:
            idx_start += 1

        # Run simulation
        for i in range(idx_start + 1, self.t.shape[0]):
            # NB1: r.successful() is not tested
            # NB2: r.t is self-incrementing
            t = self.t[i]
            r.integrate(t)
            self.x[i, :] = r.y

    def plot(self, semilog=True):
        # Simple plot

        if semilog:
            plt.semilogy(self.t, self.x[:, self.state_ref['S']], label="Susceptible")
            plt.semilogy(self.t, self.x[:, self.state_ref['E']], label="Exposed")
            plt.semilogy(self.t, self.x[:, self.state_ref['I']], label="Infectious")
            plt.semilogy(self.t, self.x[:, self.state_ref['R']], label="Recovered")
            plt.semilogy(self.t, self.x[:, self.state_ref['D']], label="Death")
            plt.ylim((0.1, np.max(self.x[:, self.state_ref['S']])))
        else:
            plt.plot(self.t, self.x[:, self.state_ref['E']], label="Exposed")
            plt.plot(self.t, self.x[:, self.state_ref['I']], label="Infectious")
            plt.plot(self.t, self.x[:, self.state_ref['D']], label="Death")

        plt.title("Propagation model SEIR(D)")
        plt.grid(True)

        plt.legend()
        plt.show()

    def plot_with_obs(self, yobs, ax=None, transparency=1):
        # Simple plot

        if ax is None:
            fig, ax = plt.subplots()
            plt.xlabel('Time (days)')
            plt.ylabel(r'$\log_{10}(y(t))$')

        idx = np.where(yobs > 0)
        tobs = np.array(self.tobs_rel)
        tobs = tobs[idx[0]]
        yobs = yobs[idx[0]]

        # pdb.set_trace()
        plt.semilogy(self.t, self.x[:, self.state_ref['I']], label='Prediction', color='C4', alpha=transparency / 2)
        plt.semilogy(self.t, self.x[:, self.state_ref['D']], label='Prediction', color='C7', alpha=transparency)
        plt.semilogy(tobs, yobs, label='Observed', linestyle='dashed', marker='o', color='C1')
        return ax


import numpy as np
from scipy.spatial import ConvexHull
from shapely.geometry import Point
from scipy.ndimage.interpolation import rotate

def minimum_bounding_rectangle(df):
    """
    Taken from: https://gis.stackexchange.com/questions/22895/finding-minimum-area-rectangle-for-given-points/22904
    This finds the smallest rectangle which bounds a set of points, and calculates the
    coordinates of the corners of the box.

    :param df: A dataframe of coordinates stored in X and Y columns
    :value: an nx2 matrix of coordinates representing the corners of the bounding box.
    """
    pi2 = np.pi/2.

    # Convert dataframe to array
    points = np.array(df[['x', 'y']])

    # Get the convex hull for the points
    # Convex hull = the smallest convex polygon which contains all the points.
    hull_points = points[ConvexHull(points).vertices]

    # calculate edge angles
    edges = np.zeros((len(hull_points)-1, 2))
    edges = hull_points[1:] - hull_points[:-1]

    angles = np.zeros((len(edges)))
    angles = np.arctan2(edges[:, 1], edges[:, 0])

    angles = np.abs(np.mod(angles, pi2))
    angles = np.unique(angles)

    # find rotation matrices
    # XXX both work
    rotations = np.vstack([
        np.cos(angles),
        np.cos(angles-pi2),
        np.cos(angles+pi2),
        np.cos(angles)]).T
#     rotations = np.vstack([
#         np.cos(angles),
#         -np.sin(angles),
#         np.sin(angles),
#         np.cos(angles)]).T
    rotations = rotations.reshape((-1, 2, 2))

    # apply rotations to the hull
    rot_points = np.dot(rotations, hull_points.T)

    # find the bounding points
    min_x = np.nanmin(rot_points[:, 0], axis=1)
    max_x = np.nanmax(rot_points[:, 0], axis=1)
    min_y = np.nanmin(rot_points[:, 1], axis=1)
    max_y = np.nanmax(rot_points[:, 1], axis=1)

    # find the box with the best area
    areas = (max_x - min_x) * (max_y - min_y)
    best_idx = np.argmin(areas)

    # return the best box
    x1 = max_x[best_idx]
    x2 = min_x[best_idx]
    y1 = max_y[best_idx]
    y2 = min_y[best_idx]
    r = rotations[best_idx]

    rval = np.zeros((4, 2))
    rval[0] = np.dot([x1, y2], r)
    rval[1] = np.dot([x2, y2], r)
    rval[2] = np.dot([x2, y1], r)
    rval[3] = np.dot([x1, y1], r)

    return rval

def two_opt(locations,improvement_threshold):
    """
    Adapted from: https://stackoverflow.com/questions/25585401/travelling-salesman-in-scipy.
    This uses the 2-opt algorithm adapted from https://en.wikipedia.org/wiki/2-opt to calculate
    the optimum order in which to visit a set of coordinates in order to minimise the distance travelled.
    This also is based on the route ending where it started.
    : param locations: Nx2 array containing X and Y coordinates.
    : param improvement_threshold:
    : return route: An array containing the order of index values of df in which sites should be visited.
    """
    # Calculate the euclidian distance in n-space of the route r traversing locations c, ending at the path start.
    path_distance = lambda r,c: np.sum([np.linalg.norm(c[r[p]]-c[r[p-1]]) for p in range(len(r))])
    # Reverse the order of all elements from element i to element k in array r.
    two_opt_swap = lambda r,i,k: np.concatenate((r[0:i],r[k:-len(r)+i-1:-1],r[k+1:len(r)]))
    # Make an array of row numbers corresponding to locations.
    route = np.arange(locations.shape[0])
    # Initialize the improvement factor.
    improvement_factor = 1
    # Calculate the distance of the initial path
    best_distance = path_distance(route,locations)
    while improvement_factor > improvement_threshold: # If the route is still improving, keep going!
        # Record the distance at the beginning of the loop.
        distance_to_beat = best_distance
        # From each location except the first and last
        for swap_first in range(1,len(route)-2):
            # to each of the locations following,
            for swap_last in range(swap_first+1,len(route)):
                # try reversing the order of these locations
                new_route = two_opt_swap(route,swap_first,swap_last)
                # and check the total distance with this modification.
                new_distance = path_distance(new_route,locations)
                # If the path distance is an improvement,
                if new_distance < best_distance:
                    # Make this the accepted best route
                    route = new_route
                    # and update the distance corresponding to this route.
                    best_distance = new_distance
         # Calculate how much the route has improved.
        improvement_factor = 1 - best_distance/distance_to_beat
    # When the route is no longer improving substantially, stop searching and return the route -->
    return route


import torch
import math
import numpy as np

def lerp_np(x,y,w):
    fin_out = (y-x)*w + x
    return fin_out

def generate_fractal_noise_2d(shape, res, octaves=1, persistence=0.5):
    noise = np.zeros(shape)
    frequency = 1
    amplitude = 1
    for _ in range(octaves):
        noise += amplitude * generate_perlin_noise_2d(shape, (frequency*res[0], frequency*res[1]))
        frequency *= 2
        amplitude *= persistence
    return noise

def generate_perlin_noise_2d(shape, res):
    def f(t):
        return 6 * t ** 5 - 15 * t ** 4 + 10 * t ** 3

    delta = (res[0] / shape[0], res[1] / shape[1])
    d = (shape[0] // res[0], shape[1] // res[1])
    grid = np.mgrid[0:res[0]:delta[0], 0:res[1]:delta[1]].transpose(1, 2, 0) % 1
    # Gradients
    angles = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1)
    gradients = np.dstack((np.cos(angles), np.sin(angles)))
    g00 = gradients[0:-1, 0:-1].repeat(d[0], 0).repeat(d[1], 1)
    g10 = gradients[1:, 0:-1].repeat(d[0], 0).repeat(d[1], 1)
    g01 = gradients[0:-1, 1:].repeat(d[0], 0).repeat(d[1], 1)
    g11 = gradients[1:, 1:].repeat(d[0], 0).repeat(d[1], 1)
    # Ramps
    n00 = np.sum(grid * g00, 2)
    n10 = np.sum(np.dstack((grid[:, :, 0] - 1, grid[:, :, 1])) * g10, 2)
    n01 = np.sum(np.dstack((grid[:, :, 0], grid[:, :, 1] - 1)) * g01, 2)
    n11 = np.sum(np.dstack((grid[:, :, 0] - 1, grid[:, :, 1] - 1)) * g11, 2)
    # Interpolation
    t = f(grid)
    n0 = n00 * (1 - t[:, :, 0]) + t[:, :, 0] * n10
    n1 = n01 * (1 - t[:, :, 0]) + t[:, :, 0] * n11
    return np.sqrt(2) * ((1 - t[:, :, 1]) * n0 + t[:, :, 1] * n1)

def rand_perlin_2d_np(shape, res, fade=lambda t: 6 * t ** 5 - 15 * t ** 4 + 10 * t ** 3):
    delta = (res[0] / shape[0], res[1] / shape[1])
    d = (shape[0] // res[0], shape[1] // res[1])
    grid = np.mgrid[0:res[0]:delta[0], 0:res[1]:delta[1]].transpose(1, 2, 0) % 1

    angles = 2 * math.pi * np.random.rand(res[0] + 1, res[1] + 1)
    gradients = np.stack((np.cos(angles), np.sin(angles)), axis=-1)
    tt = np.repeat(np.repeat(gradients,d[0],axis=0),d[1],axis=1)

    tile_grads = lambda slice1, slice2: np.repeat(np.repeat(gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]],d[0],axis=0),d[1],axis=1)
    dot = lambda grad, shift: (
                np.stack((grid[:shape[0], :shape[1], 0] + shift[0], grid[:shape[0], :shape[1], 1] + shift[1]),
                            axis=-1) * grad[:shape[0], :shape[1]]).sum(axis=-1)

    n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0])
    n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0])
    n01 = dot(tile_grads([0, -1], [1, None]), [0, -1])
    n11 = dot(tile_grads([1, None], [1, None]), [-1, -1])
    t = fade(grid[:shape[0], :shape[1]])
    return math.sqrt(2) * lerp_np(lerp_np(n00, n10, t[..., 0]), lerp_np(n01, n11, t[..., 0]), t[..., 1])

def rand_perlin_2d(shape, res, fade=lambda t: 6 * t ** 5 - 15 * t ** 4 + 10 * t ** 3):
    delta = (res[0] / shape[0], res[1] / shape[1])
    d = (shape[0] // res[0], shape[1] // res[1])

    grid = torch.stack(torch.meshgrid(torch.arange(0, res[0], delta[0]), torch.arange(0, res[1], delta[1])), dim=-1) % 1
    angles = 2 * math.pi * torch.rand(res[0] + 1, res[1] + 1)
    gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim=-1)

    tile_grads = lambda slice1, slice2: gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]].repeat_interleave(d[0],
                                                                                0).repeat_interleave(
        d[1], 1)
    dot = lambda grad, shift: (
                torch.stack((grid[:shape[0], :shape[1], 0] + shift[0], grid[:shape[0], :shape[1], 1] + shift[1]),
                            dim=-1) * grad[:shape[0], :shape[1]]).sum(dim=-1)

    n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0])

    n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0])
    n01 = dot(tile_grads([0, -1], [1, None]), [0, -1])
    n11 = dot(tile_grads([1, None], [1, None]), [-1, -1])
    t = fade(grid[:shape[0], :shape[1]])
    return math.sqrt(2) * torch.lerp(torch.lerp(n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1])

def rand_perlin_2d_octaves(shape, res, octaves=1, persistence=0.5):
    noise = torch.zeros(shape)
    frequency = 1
    amplitude = 1
    for _ in range(octaves):
        noise += amplitude * rand_perlin_2d(shape, (frequency * res[0], frequency * res[1]))
        frequency *= 2
        amplitude *= persistence
    return noise


# Chapter 5 Compressing Data via Dimensionality Reduction
# Using kernel principal component analysis for nonlinear mappings
'''
Using kernel PCA, we will learn how to transform data that is not linearly spearable
onto a new, lower-dimensional subspace that is suitable for linear classifiers.'''

# Kernel functions and the kernel trick

# implementing a kernel principal component analysis in Python
'''
Now we are going to implement an RBF kernel PCA in Python following the three
steps that summaried the kernal PCA approach. Using the SciPy and Numpy helper
function's, we will see that implementing a kernel PCA is actually really simple.'''
from scipy.spatial.distance import pdist, squareform
from scipy import exp
from scipy.linalg import eigh
import numpy as np
def rbf_kernel_pca(X, gamma, n_components):

    '''
    RBF kernel PCA implementation.

    Parameters
    --------------------
    X: {NumPy ndarray}, shape = [n_samples, n_features]

    gamma: float
        Tuning parameter of the RBF kernel

    n_components: int
        Number of principal components to return

    Returns
    -------------------
    X_pc: {Numpy ndarray}, shape = [n_samples, k_features]
        Projected dataset

    '''
    # Calculate pairwise squared Euclidean distance
    # in the MxN dimensional dataset
    sq_dists = pdist(X, 'sqeuclidean' )

    # Convert pairwise distances into a square matrix
    mat_sq_dists = squareform(sq_dists)

    # Compute the symmetric kernel matrix
    K = exp(-gamma * mat_sq_dists)

    # Center the kernel matrix.
    N = K.shape[0]
    one_n = np.ones((N, N)) / N
    K = K - one_n.dot(K) - K.dot(one_n) + one_n.dot(K).dot(one_n)

    # Obtaining eigenpairs from the centered kernel matrix
    # numpy.eigh returns them in sorted order
    eigvals, eigvecs = eigh(K)

    # Collect the top k eigenvectors (projected samples)
    X_pc = np.column_stack((eigvecs[:, -i] for i in range(1, n_components + 1)))

    return X_pc

# Example 1 - separating half-moon shapes
'''
Now, let's apply our rbf_kernel_pca on some nonlinear example datasets.
We will start by creating a two-dimensional dataset of 100 sample points
representing two half-moon shapes:'''
from sklearn.datasets import make_moons
X, y = make_moons(n_samples=100, random_state=123)
plt.scatter(X[y==0, 0], X[y==0, 1], color='red', marker='^', alpha=0.5)
plt.scatter(X[y==1, 0], X[y==1, 1], color='blue', marker='o', alpha=0.5)
plt.show()

'''
For the purposes of illustration, the half-moon of triangular symbols shall
represent one class and the half-moon depicted by the circular symbols
represent the samples from another class:'''
'''
Clearly, these two half-moon shapes are not linearly separable and our goal
is to unfold the half-moons via kernel PCA so that the dataset can serve as a
suitable input for a linear classifier. But first, let's see what the dataset looks
like if we project it onto the principal components via standard PCA:'''
from sklearn.decomposition import PCA
scikit_pca = PCA(n_components=2)
X_spca = scikit_pca.fit_transform(X)
fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(7,3))
ax[0].scatter(X_spca[y==0, 0], X_spca[y==0, 1],color='red', marker='^', alpha=0.5)
ax[0].scatter(X_spca[y==1, 0], X_spca[y==1, 1],color='blue', marker='o', alpha=0.5)
ax[1].scatter(X_spca[y==0, 0], np.zeros((50,1))+0.02, color='red', marker='^', alpha=0.5)
ax[1].scatter(X_spca[y==1, 0], np.zeros((50,1)),color='blue', marker='o', alpha=0.5)
ax[0].set_xlabel('PC1')
ax[0].set_ylabel('PC2')
ax[1].set_ylim([-1, 1])
ax[1].set_yticks([])
ax[1].set_xlabel('PC1')
plt.show()

'''
Note that when we plotted the first principal component only (right subplot),
we shifted the triangular samples slightly upwards and the circular samples
slightly downwards to better visualize the class overlap'''

'''
Now let's try out our kernel PCA function rbf_kernel_pca, which we implemented
in the previous subsection:'''
from matplotlib.ticker import FormatStrFormatter
X_kpca = rbf_kernel_pca(X, gamma=15, n_components=2)
fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(7,3))
ax[0].scatter(X_kpca[y==0, 0], X_kpca[y==0, 1], color='red', marker='^', alpha=0.5)
ax[0].scatter(X_kpca[y==1, 0], X_kpca[y==1, 1], color='blue', marker='o', alpha=0.5)
ax[1].scatter(X_kpca[y==0, 0], np.zeros((50, 1)) + 0.02, color='red', marker='^', alpha=0.5)
ax[1].scatter(X_kpca[y==1, 0], np.zeros((50, 1)) - 0.02, color='blue', marker='o', alpha=0.5)
ax[0].set_xlabel('PC1')
ax[0].set_ylabel('PC2')
ax[1].set_ylim([-1, 1])
ax[1].set_yticks([])
ax[1].set_xlabel('PC1')
ax[0].xaxis.set_major_formatter(FormatStrFormatter('%0.1f'))
ax[1].xaxis.set_major_formatter(FormatStrFormatter('%0.1f'))
plt.show()

'''
We can now see that the two classes (circles and triangles) are linearly well separated
so that it becomes a suitable training dataset for linear classifiers:'''
'''
Unfortuately, there is no universal value for the tuning parameter gamma that works
well for different datasets. To find a gamma value that is appropriate for a given problem
requires experimentation. '''

# Example 2 - separating concentric circles
'''
In the previous subsection, we showed how to separate half-moon shapes via
kernel-PCA. Since we put so much effort into understanding the concept of  kernel
PCA, let's take a look at another interesting example of a nonlinear problem:
concetric circles.'''
from sklearn.datasets import make_circles
X, y = make_circles(n_samples=1000, random_state=123, noise=0.1, factor=0.2)
plt.scatter(X[y==0, 0], X[y==0, 1], color='red', marker='^', alpha=0.5)
plt.scatter(X[y==1, 0], X[y==1, 1], color='blue', marker='o', alpha=0.5)
plt.show()

'''Let's start with the standard PCA approach to compare it with the results of the RBF
kernel PCA:'''
scikit_pca = PCA(n_components=2)
X_spca = scikit_pca.fit_transform(X)
fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(7, 3))
ax[0].scatter(X_spca[y==0, 0], X_spca[y==0, 1], color='red', marker='^', alpha=0.5)
ax[0].scatter(X_spca[y==1, 0], X_spca[y==1, 1], color='blue', marker='o', alpha=0.5)
ax[1].scatter(X_spca[y==0, 0], np.zeros((500, 1))+0.02, color='red', marker='^', alpha=0.5)
ax[1].scatter(X_spca[y==1, 0], np.zeros((500, 1))-0.02, color='blue', marker='o', alpha=0.5)
ax[0].set_xlabel('PC1')
ax[0].set_ylabel('PC2')
ax[1].set_ylim([-1, 1])
ax[1].set_yticks([])
ax[1].set_xlabel('PC1')
plt.show()
''' Again, we can see that standard PCA is not able to produce results suitable for
training a linear classifier.'''
''' Given an appropriate value for gamma, let's see if we are luckier using the RBF
kernel PCA implementation:'''
X_kpca = rbf_kernel_pca(X, gamma=15, n_components=2)
fig, ax = plt.subplots(nrows=1, ncols=2, figsize=(7, 3))
ax[0].scatter(X_kpca[y==0, 0], X_kpca[y==0, 1], color='red', marker='^', alpha=0.5)
ax[0].scatter(X_kpca[y==1, 0], X_kpca[y==1, 1], color='blue', marker='o', alpha=0.5)
ax[1].scatter(X_kpca[y==0, 0], np.zeros((500, 1))+0.02, color='red', marker='^', alpha=0.5)
ax[1].scatter(X_kpca[y==1, 0], np.zeros((500, 1))-0.02, color='blue', marker='o', alpha=0.5)
ax[0].set_xlabel('PC1')
ax[0].set_ylabel('PC2')
ax[1].set_ylim([-1, 1])
ax[1].set_yticks([])
ax[1].set_xlabel('PC1')
plt.show()

# Projecting new data points
from scipy.spatial.distance import pdist, squareform
from scipy import exp
from scipy import exp
from scipy.linalg import eigh
import numpy as np

def rbf_kernel_pca(X, gamma, n_components):
    '''
    RBF kernel PCA implementation.

    Parameters
    -----------------
    X: {Numpy ndarray}, shape = {n_samples, n_features}

    gamma: float
        Tuning parameter of the RBF kernel

    n_components: int
        number of principal components to return

    Returns
    -----------------
    X_pc: {Numpy ndarray}, shape = [n_samples, k_features}
        Projected dataset

    lambdas: list
        Eigenvalues

    '''
    # Calculate pairwise squared Euclidean distances
    # in the MxN dimensional dataset.
    sq_dists = pdist(X, 'sqeuclidean')

    # Convert pairwise distances into a square matrix.
    mat_sq_dists = squareform(sq_dists)

    # Compute the symmertric kernel matrix
    K = exp(-gamma * mat_sq_dists)

    # Center the kernel matrix
    N = K.shape[0]
    one_n = np.ones((N, N)) / N
    K = K - one_n.dot(K) - K.dot(one_n) + one_n.dot(K).dot(one_n)

    # Obtaining eigenpairs from the centered kernel matrix
    # numpy.eigh returns them in sorted order
    eigvals, eigvecs = eigh(K)

    # Collect the top k eigenvectors (projected samples)
    alphas = np.column_stack((eigvecs[:, -i] for i in range(1, n_components+1)))

    # Collect the corresponding eigenvalues
    lambdas = [eigvals[-i] for i in range(1, n_components+1)]

    return alphas, lambdas

'''
Now, let's create a new half-moon dataset and project it onto a one-dimensional
subspace using the updated RBF kernel PCA implementation:'''
X, y = make_moons(n_samples=100, random_state=123)
alphas, lambdas = rbf_kernel_pca(X, gamma=15, n_components=1)

'''
To make sure that we implement the code for projecting new samples, let's assume
that the 26th point from the half-moon dataset is a new data point x', and our task is
to project it onto this new subspace:'''
x_new = X[25]
x_new
x_proj = alphas[25] # original projection
x_proj

def project_x(x_new, X, gamma, alphas, lambdas):
    pair_dist = np.array([np.sum((x_new-row)**2) for row in X])
    k = np.exp(-gamma * pair_dist)
    return k.dot(alphas / lambdas)

'''
By executing the following code, we are able to reproduce the original projection.
Using the project_x function, we will be able to project any new data samples as
well. The code is as follows:'''
x_reproj = project_x(x_new, X, gamma=15, alphas=alphas, lambdas=lambdas)
x_reproj

'''
Lastly, let's visualize the projection on the first principal component.'''
plt.scatter(alphas[y==0, 0], np.zeros((50)), color='red', marker='^', alpha=0.5)
plt.scatter(alphas[y==1, 0], np.zeros((50)), color='blue', marker='o', alpha=0.5)
plt.scatter(x_proj, 0, color='black', label='original projection of point X[25]', marker='^', s=100)
plt.scatter(x_reproj, 0, color='green', label='remapped point X[25]', marker='x', s=500)
plt.legend(scatterpoints=1)
plt.show()
'''
As we can see in the following scatterplot, we mapped the sample x' onto the first
principal component correctly:'''

# Kernel principal component analysis in scikit-learn
'''
For our convenience, scikit-learn implements a kernel PCA class in the
sklearn.decomposition submodule. The usage is similar to the standard
PCA class, and we can specify the kernel via the kernel parameter:'''
from sklearn.decomposition import KernelPCA
X, y = make_moons(n_samples=100, random_state=123)
scikit_kpca = KernelPCA(n_components=2, kernel='rbf', gamma=15)
X_skernpca=scikit_kpca.fit_transform(X)
'''
To see if we get results that are consistent with our own kernel PCA
implementation, let's plot the transformed half-moon shape data onto the
first two principal components:'''
plt.scatter(X_skernpca[y==0, 0], X_skernpca[y==0, 1], color='red', marker='^', alpha=0.5)
plt.scatter(X_skernpca[y==1, 0], X_skernpca[y==1, 1], color='blue', marker='o', alpha=0.5)
plt.xlabel('PC1')
plt.ylabel('PC2')
plt.plt.show()


import matplotlib
import matplotlib.pyplot as plt

import numpy as np
import math

from skimage import data
from skimage.filters import try_all_threshold

from skimage import io
import os

import scipy.ndimage

filename = os.getcwd() + '/Ejercicio-3/Img/ruido03.png'

img = io.imread(filename, as_gray=False)

Gx = np.array([[-1,0,1],[-2,0,2],[-1,0,1]])
Gy = np.rot90(Gx, k=1)

G45 = np.array([[0,1,2],[-1,0,1],[-2,-1,0]])
G135 = np.rot90(G45, k=1)

def convolve_1d(array, kernel):
    ks = kernel.shape[0] # shape gives the dimensions of an array, as a tuple
    final_length = array.shape[0] - ks + 1
    return np.array([(array[i:i+ks]*kernel).sum() for i in range(final_length)])

def convolve_2d(array,kernel):
    ks = kernel.shape[1] # shape gives the dimensions of an array, as a tuple
    final_height = array.shape[1] - ks + 1
    return np.array([convolve_1d(array[:,i:i+ks],kernel) for i in range(final_height)]).T

image1 = convolve_2d(img, Gx)
image2 = convolve_2d(img, Gy)
image3 = convolve_2d(img, G45)
image4 = convolve_2d(img, G135)

image5 = image1*0
for i in range(0,image1.shape[0]):
    for j in range(0,image1.shape[1]):
        image5[i,j] = (image1[i,j]**2 + image2[i,j]**2)**0.5

fig, ax = plt.subplots(1, 1)
ax.imshow(img, cmap=plt.cm.gray)
ax.set_title('Original')
ax.axis('off')

fig1, ax1 = plt.subplots(1, 1)
ax1.imshow(image1, cmap=plt.cm.gray)
ax1.set_title('Convolución horizontal')
ax1.axis('off')

fig2, ax2 = plt.subplots(1, 1)
ax2.imshow(image2, cmap=plt.cm.gray)
ax2.set_title('Convolución vertical')
ax2.axis('off')

fig3, ax3 = plt.subplots(1, 1)
ax3.imshow(image3, cmap=plt.cm.gray)
ax3.set_title('Convolución 45º')
ax3.axis('off')

fig4, ax4 = plt.subplots(1, 1)
ax4.imshow(image4, cmap=plt.cm.gray)
ax4.set_title('Convolución 135º')
ax4.axis('off')

fig5, ax5 = plt.subplots(1, 1)
ax5.imshow(image5, cmap=plt.cm.gray)
ax5.set_title('Convolución Total')
ax5.axis('off')

plt.show()


import logging

import numpy as np
import pandas as pd

logger = logging.getLogger("hepynet")

def cov(x, y, w):
    """Calculates weighted covariance"""
    return np.sum(
        w
        * (x - np.average(x, axis=0, weights=w))
        * (y - np.average(y, axis=0, weights=w))
    ) / np.sum(w)

def corr(x, y, w):
    """Calculates weighted correlation"""
    return cov(x, y, w) / np.sqrt(cov(x, x, w) * cov(y, y, w))

def corr_matrix(x, w=None):
    """Calculates correlation coefficient matrix"""
    if w is None:
        w = np.ones(x.shape[0])
    num_features = x.shape[1]
    # np.cov can't deal with negative weights, use self-defined function for now
    corr_m = np.zeros((num_features, num_features))
    for row in range(num_features):  # TODO: need to optimize the algorithm
        for col in range(row + 1):
            corr_m[row][col] = corr(x[:, row], x[:, col], w)
            corr_m[col][row] = corr_m[row][col]
    return corr_m

def extract_bkg_df(df: pd.DataFrame):
    return df.loc[(df["is_mc"] == True) & (df["is_sig"] == False)]

def extract_sig_df(df: pd.DataFrame):
    return df.loc[df["is_sig"] == True]

def get_cut_index(np_array, cut_values, cut_types):
    """Parses cuts arguments and returns cuts indexes."""
    assert len(cut_values) == len(
        cut_types
    ), "cut_values and cut_types should have same length."
    pass_index = None
    for cut_value, cut_type in zip(cut_values, cut_types):
        temp_index = get_cut_index_value(np_array, cut_value, cut_type)
        if pass_index is None:
            pass_index = temp_index
        else:
            pass_index = np.intersect1d(pass_index, temp_index)
    return pass_index

def get_cut_index_value(np_array: np.ndarray, cut_value, cut_type):
    """Returns cut indexes based on cut_value and cut_type.

    If cut_type is "=":
        returns all indexes that have values equal to cut_value
    If cut_type is ">":
        returns all indexes that have values lager than cut_value
    If cut_type is "<":
        returns all indexes that have values smaller than cut_value

    Args:
        array_dict: numpy array
        cut_feature: str
        cut_bool: bool
    """
    logger.debug("@ data_io.array_utils.get_cut_index_value")
    assert np_array.ndim == 1
    logger.debug(f"Cut input shape: {np_array.shape}, cut value {cut_value}")
    # Make cuts
    if cut_type == "=":
        pass_index = np.argwhere(np_array == cut_value)
    elif cut_type == ">":
        pass_index = np.argwhere(np_array > cut_value)
    elif cut_type == "<":
        pass_index = np.argwhere(np_array < cut_value)
    else:
        raise ValueError("Invalid cut_type specified.")
    return pass_index.flatten()


import numpy as np

# it may cause an overflow error
def softmax_overflow(a):
    exp_a = np.exp(a)
    sum_exp_a = np.sum(exp_a)
    y = exp_a / sum_exp_a

    return y

def softmax(a):
    c = np.max(a)
    exp_a = np.exp(a - c)  # for overflow
    sum_exp_a = np.sum(exp_a)
    y = exp_a / sum_exp_a

    return y

a = np.array([0.3, 2.9, 4.0])

exp_a = np.exp(a)
print(exp_a)

sum_exp_a = np.sum(exp_a)

print(sum_exp_a)

y = exp_a / sum_exp_a
print(y)

# fix overflow
b = np.array([1010, 1000, 990])
print(np.exp(b) / np.sum(np.exp(b)))
# -> array([ nan, nan, nan])

c = np.max(b)
print(b - c)
# -> array([ 0, -10, -20])

print(np.exp(b - c) / np.sum(np.exp(b - c)))
# -> array([ 9.999546e-01, 4.53978686e-05, 2.06106005e-09])

# softmax result
d = np.array([0.3, 2.9, 4.0])
z = softmax(d)
print(z)
# -> [ 0.01821127, 0.24519181, 0.736559691]
print(np.sum(z))
# -> 1.0


import numpy as np

def get_positional_encoding(H, W, dim=128):
    """
    2D positional encoding, following
        https://papers.nips.cc/paper/7181-attention-is-all-you-need.pdf
    """

    assert dim % 4 == 0, 'dim must be a multiply of 4 (h/w x sin/cos)'
    c_period = 10000. ** np.linspace(0., 1., dim // 4)
    h_vec = np.tile(np.arange(0, H).reshape((H, 1, 1)), (1, W, 1)) / c_period
    w_vec = np.tile(np.arange(0, W).reshape((1, W, 1)), (H, 1, 1)) / c_period
    position_encoding = np.concatenate(
        (np.sin(h_vec), np.cos(h_vec), np.sin(w_vec), np.cos(w_vec)), axis=-1)
    position_encoding = position_encoding.reshape((1, H, W, dim))
    return position_encoding