import numpy as np

class NeuralNet(object):
    '''
    Neural Network with tanh activation and softmax output function using Adam Optimizer.

    Parameters
    ----------
    wShape : tuple
        Shape of feedforward neural network architecture.

    Usage: nn = NeuralNet(ni, nhi, no)
        ni: Number of input features.
        nhi: Number of neurons in ith hidden layer.
        no: Number of outputs.
    '''

    def __init__(self, wShape):
        super().__init__()
        self.wShape = wShape
        self.param = {
            'W':[np.random.randn(self.wShape[i+1], self.wShape[i]) for i in range(len(self.wShape)-1)],
            'B':[np.zeros((self.wShape[i+1], 1)) for i in range(len(wShape)-1)]
        }

    def softmax(self, x):
        x = np.exp(x - np.max(x, axis=0))
        x /= np.sum(x, axis=0)
        return x

    def forwardPass(self, X):
        A, Z = [X], []
        for i in range(len(self.wShape)-1):
            Z.append(np.add(np.matmul(self.param['W'][i], A[-1]), self.param['B'][i]))
            A.append((self.softmax(Z[-1]) if i==len(self.wShape)-2 else np.tanh(Z[-1])))

        Y_hat = A[-1]; del A[0];
        cache = {
            'Z': Z,
            'A': A
        }
        return Y_hat, cache

    def crossEntropyLoss(self, Y_hat, Y, esp=1e-15):
        return -np.mean(np.log(np.clip(Y_hat[Y], esp, 1-esp)))

    def predictLabel(self, Y_hat):
        return np.argmax(Y_hat, axis=0)

    def backwardPass(self, cache, X, Y):
        dZ, dW, dB = [cache['A'][-1] - Y], [], []
        for i in range(len(self.wShape)-3, -1, -1):
            dZ.insert(0, np.matmul(np.transpose(self.param['W'][i+1]), dZ[0]) * (1 - np.power(cache['A'][i], 2)))
            dW.insert(0, np.matmul(dZ[1], np.transpose(cache['A'][i]))); dB.insert(0, np.sum(dZ[1], axis=1, keepdims=True))
        dW.insert(0, np.matmul(dZ[0], np.transpose(X))); dB.insert(0, np.sum(dZ[0], axis=1, keepdims=True))

        grads = {
            'dW': dW,
            'dB': dB
        }
        return grads

    def updateParam(self, X, Y, itrr=100, alpha=1e-2, beta1=0.9, beta2=0.999, espp=1e-8):
        '''
        Adam Optimizer:-
            Vdw = beta1*Vdw + (1-beta1)*dW; VdB = beta1*VdB + (1-beta1)*dB
            Sdw = beta2*Sdw + (1-beta2)*dW**2; SdB = beta2*SdB + (1-beta2)*dB**2
            Vdw, VdB /= (1 - beta1**(itteration)); Sdw, SdB /= (1 - beta2**(itteration))
            W -= alpha * Vdw/(sqrt(Sdw) + esp); B -= alpha * VdB/(sqrt(SdB) + esp)
        '''

        Vdw = [0. for i in range(len(self.wShape)-1)]
        VdB = [0. for i in range(len(self.wShape)-1)]
        Sdw = [0. for i in range(len(self.wShape)-1)]
        SdB = [0. for i in range(len(self.wShape)-1)]

        for itteration in range(1, itrr+1):
            Y_hat, cache = self.forwardPass(X)
            grads = self.backwardPass(cache, X, Y)
            for layer in range(len(self.wShape)-1):
                Vdw[layer] = (beta1*Vdw[layer] + (1-beta1)*grads['dW'][layer])/(1-beta1**itteration)
                VdB[layer] = (beta1*VdB[layer] + (1-beta1)*grads['dB'][layer])/(1-beta1**itteration)
                Sdw[layer] = (beta2*Sdw[layer] + (1-beta2)*np.power(grads['dW'][layer], 2))/(1-beta2**itteration)
                SdB[layer] = (beta2*SdB[layer] + (1-beta2)*np.power(grads['dB'][layer], 2))/(1-beta2**itteration)
                self.param['W'][layer] -= alpha * Vdw[layer]/(np.sqrt(Sdw[layer])+espp)
                self.param['B'][layer] -= alpha * VdB[layer]/(np.sqrt(SdB[layer])+espp)

    def accuracy(self, y_hat, y):
        return np.mean(y_hat==y)*100