import numpy as np
from Function import Function


class Activation(Function):
    def __init__(self, Z, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.Z = Z
        self.Activation_Output = self.forward_pass(Z)
        self.grad = self.local_gradient(Z)

    def backward_pass(self, dY):
        # dY is the global gradient from the nec=xt layer propagated back to this layer
        return dY * self.grad['dA']

    def Sigmoid(Z):
        Sigmoid_Act = Sigmoid(Z)
        return Sigmoid_Act

    def Relu(Z):
        Relu_Act = Relu(Z)
        return Relu_Act

    def BinaryStep(Z):
        BinaryStep_Act = BinaryStep(Z)
        return BinaryStep_Act

    def Linear_act(Z):
        Linear_Act = Linear_act(Z)
        return Linear_Act

    def Softmax(Z):
        Softmax_Act = Softmax(Z)
        return Softmax_Act

    def LeakyReLU(Z):
        LeakyReLU_Act = LeakyReLU(Z)
        return LeakyReLU_Act


class Sigmoid(Activation):
    def forward_pass(self, Z):
        o = 1 / (1 + np.exp(-Z))  # o is an intermediate variable
        self.cache['Z'] = Z
        self.cache['A'] = o
        return o

    def local_gradient(self, Z):
        sigmoid = 1 / (1 + np.exp(-Z))

        self.grad = {'dA': np.multiply(sigmoid, (1 - sigmoid))}
        return self.grad


class Relu(Activation):
    def forward_pass(self, Z):
        o = np.maximum(0, Z)  # o is an intermediate variable
        self.cache['Z'] = Z
        self.cache['A'] = o

        return o

    def local_gradient(self, Z):
        Z = (Z > 0)
        self.grad = {'dA': Z}
        return self.grad


class BinaryStep(Activation):
    def forward_pass(self, Z):
        if Z >= 0:  # o is an intermediate variable
            o = 1
        elif Z < 0:
            o = -1
        self.cache['Z'] = Z
        self.cache['A'] = o
        return o

    def local_gradient(self, Z):
        self.grad = {'dA': 0}
        return self.grad


class Linear_act(Activation):
    def forward_pass(self, Z):
        self.cache['Z'] = Z
        self.cache['A'] = Z
        return Z

    def local_gradient(self, Z):
        self.grad = {'dA': np.ones_like(Z)}
        return self.grad


class Softmax(Activation):
    def forward(self, Z):
        exp_x = np.exp(Z)
        o = exp_x / np.sum(exp_x, axis=1, keepdims=True)  # o is an intermediate variable
        self.cache['Z'] = Z
        self.cache['A'] = o
        return o

    def local_grad(self, Z):
        jacobian_m = np.diag(Z)
        for i in range(len(jacobian_m)):
            for j in range(len(jacobian_m)):
                if i == j:
                    jacobian_m[i][j] = Z[i] * (1 - Z[i])
                else:
                    jacobian_m[i][j] = -Z[i] * Z[j]
        self.grad = {'dA': jacobian_m}
        return self.grad


class LeakyReLU(Activation):
    def forward(self, Z, LRP):
        o = Z * (Z > 0) + LRP * Z * (Z <= 0)  # o is an intermediate variable
        self.cache['Z'] = Z
        self.cache['A'] = o
        return o
        # Leaky Relu Parameter is LRP determines slope of negative part of the leaky relu function

    def local_grad(self, Z, alpha):
        self.grad = {'dA': (1 * (Z > 0) + alpha * (Z <= 0))}
        return self.grad


def get_activation(activation):
    if (activation == 'sigmoid'):
        return Sigmoid
    elif (activation == 'relu'):
        return Relu
    elif (activation == 'linear'):
        return Linear_act
    elif (activation == 'softmax'):
        return Softmax
    elif (activation == 'leakyrelu'):
        return LeakyReLU
    elif (activation == 'binarystep'):
        return BinaryStep
 
