"""An extension of NN_api, used for gradient checking (to ensure that the gradients computed via backprop are correct).
We manually perturb each weight matrix element and bias vector element, which allows us to compute the loss function
with respect to said parameter. We compare this with the gradient calculated via back-prop. The program throws an error
if the two are not the same (within tolerance).

Ian Wu
04/11/2020
"""

import NN_api
import numpy as np


class GradientChecking(NN_api.Dense):
    """Class for performing gradient checking on the NN_api deep feedforward neural network. Extends the Dense class
    of NN_api."""

    def __init__(self):
        super().__init__()

    def gradient_check_first_pass(self, input_data, labels):
        """Perform standard forwards and backwards propagation, storing the gradient matrices for weights and biases"""
        outputs = self.forward_propagate(input_data)  # forward propagation
        # initialise the gradient matrix stores
        weight_gradients = []
        bias_gradients = []

        # now perform backprop, saving the gradient matrices in their respective stores
        current_layer = self.tail  # start at the tail (output layer L)
        activation = current_layer.get_activation()
        del_l = None
        if self.loss == 'MSE':
            del_l = (outputs - labels) * self.activation_prime(current_layer.get_cache(), activation)
        elif self.loss == 'binary_cross_entropy':
            outputs[outputs == 0] += 1e-15  # prevent instability of log(0)
            outputs[outputs == 1] -= 1e-15
            del_l = ((outputs - labels) / (outputs * (1 - outputs))) * \
                    self.activation_prime(current_layer.get_cache(), activation)
        # derivative of the loss wrt the affine transformed data of the output layer L
        d_cost_w = np.matmul(del_l, current_layer.get_prev_layer().get_data().T)  # derivative of cost wrt
        d_cost_b = np.sum(del_l, axis=1)  # derivative of cost wrt bias of output layer L
        weight_gradients.append(d_cost_w)  # store the weight grad matrix
        bias_gradients.append(d_cost_b)  # store the bias grad matrix

        current_layer = current_layer.get_prev_layer()
        while current_layer.get_state() != 'input':
            prev_layer = current_layer.get_prev_layer()
            activation = current_layer.get_activation()
            del_l = np.matmul(current_layer.get_next_layer().get_w().T, del_l) * \
                    self.activation_prime(current_layer.get_cache(), activation)
            # derivative of loss wrt affine transformation of lth layer
            d_cost_w = np.matmul(del_l, prev_layer.get_data().T)  # derivative of loss wrt lth layer weights
            d_cost_b = np.sum(del_l, axis=1)  # derivative of cost wrt bias of output layer L
            weight_gradients.append(d_cost_w)  # store the weight grad matrix
            bias_gradients.append(d_cost_b)  # store the bias grad matrix
            current_layer = prev_layer

        return weight_gradients, bias_gradients

    def check_forward_weights(self, input_data, labels, check_layer, epsilon, row, col, weight_grads, tolerance):
        """Manually compute a weight gradient matrix element
            Parameters:
                input_data: matrix of input data
                labels: vector of labels
                check_layer: the layer containing the element to perform gradient checking on
                epsilon: perturbation to use in gradient checking (< 1e-5)
                row: the row of the matrix containing the element to perform gradient checking on
                col: the column of the matrix containing the element to perform gradient checking on
                weight_grads: list of backprop-computed weight gradient matrices, generated by gradient_check_first_pass
                tolerance: tolerance for numerical error
        """
        self.head.set_data(input_data)  # pass input data into the network
        current_layer = self.head.get_next_layer()
        layer_index = 0
        store = []
        # step through the network and compute the transformations on the data
        while current_layer:
            layer_index += 1
            if layer_index == check_layer:
                current_weights = current_layer.get_w()
                current_weights[row][col] += 0.5 * epsilon  # locate the matrix element to perturb, and perturb it
                current_layer.set_w(current_weights)
            current_layer.compute_data()
            current_layer = current_layer.get_next_layer()
        store.append(self.tail.get_data())  # store the outputs of the pos. perturbed forward prop

        current_layer = self.head.get_next_layer()
        layer_index = 0
        while current_layer:
            layer_index += 1
            if layer_index == check_layer:
                current_weights = current_layer.get_w()
                current_weights[row][col] -= epsilon  # locate the matrix element to perturb, and perturb it
                # perturb in the opposite direction to before, double to offset the initial perturbation as well
                current_layer.set_w(current_weights)
            current_layer.compute_data()
            current_layer = current_layer.get_next_layer()
        store.append(self.tail.get_data())  # store the outputs of the neg. perturbed forward prop

        # reset the weights (remove the perturbations)
        current_layer = self.head.get_next_layer()
        layer_index = 0
        while current_layer:
            layer_index += 1
            if layer_index == check_layer:
                current_weights = current_layer.get_w()
                current_weights[row][col] += 0.5 * epsilon
                current_layer.set_w(current_weights)
                break

        p_loss = self.compute_loss(store[0], labels, loss=self.loss)  # convert our forward prop outputs to losses
        n_loss = self.compute_loss(store[1], labels, loss=self.loss)
        ep_grad = (p_loss - n_loss)/epsilon  # manually computed gradient
        list_ind = len(weight_grads) - check_layer
        unmodified_weight_grad = weight_grads[list_ind][row][col]
        # compare manually computed gradient with backprop-computed gradient. Throw error if it is greater than 1e-6
        assert(abs(unmodified_weight_grad - ep_grad) <= tolerance)

    def check_forward_biases(self, input_data, labels, check_layer, epsilon, row, bias_grads, tolerance):
        """Manually compute a bias gradient vector element
            Parameters:
                input_data: matrix of input data
                labels: vector of labels
                check_layer: the layer containing the element to perform gradient checking on
                epsilon: perturbation to use in gradient checking (< 1e-5)
                row: the element of the vector containing the element to perform gradient checking on
                bias_grads: list of backprop-computed bias gradient vectors, generated by gradient_check_first_pass
                tolerance: tolerance for numerical error
        """
        self.head.set_data(input_data)  # pass input data into the network
        current_layer = self.head.get_next_layer()
        layer_index = 0
        store = []
        # step through the network and compute the transformations on the data
        while current_layer:
            layer_index += 1
            if layer_index == check_layer:
                current_biases = current_layer.get_b()
                current_biases[row] += 0.5 * epsilon  # locate the vector element to perturb, and perturb it
                current_layer.set_b(current_biases)
            current_layer.compute_data()
            current_layer = current_layer.get_next_layer()
        store.append(self.tail.get_data())  # store the outputs of the pos. perturbed forward prop

        current_layer = self.head.get_next_layer()
        layer_index = 0
        while current_layer:
            layer_index += 1
            if layer_index == check_layer:
                current_biases = current_layer.get_b()
                current_biases[row] -= epsilon  # locate the vector element to perturb, and perturb it
                # perturb in the opposite direction to before, double to offset the initial perturbation as well
                current_layer.set_b(current_biases)
            current_layer.compute_data()
            current_layer = current_layer.get_next_layer()
        store.append(self.tail.get_data())  # store the outputs of the pos. perturbed forward prop

        # reset the biases (remove the perturbations)
        current_layer = self.head.get_next_layer()
        layer_index = 0
        while current_layer:
            layer_index += 1
            if layer_index == check_layer:
                current_biases = current_layer.get_b()
                current_biases[row] += 0.5 * epsilon
                current_layer.set_b(current_biases)
                break

        p_loss = self.compute_loss(store[0], labels, loss=self.loss)  # convert our forward prop outputs to losses
        n_loss = self.compute_loss(store[1], labels, loss=self.loss)
        ep_grad = (p_loss - n_loss)/epsilon  # manually computed gradient
        list_ind = len(bias_grads) - check_layer
        unmodified_bias_grad = bias_grads[list_ind][row]
        # compare manually computed gradient with backprop-computed gradient. Throw error if it is greater than 1e-6
        assert(abs(unmodified_bias_grad - ep_grad) <= tolerance)

    def gradient_check(self, input_data, labels, epsilon, tolerance):
        """Perform gradient checking"""
        # get the backwards-prop generated weight and bias gradient matrices, stored in lists
        weight_gradients, bias_gradients = self.gradient_check_first_pass(input_data, labels)
        current_layer = self.head.get_next_layer()
        layer_count = 0
        w_shapes = []
        b_shapes = []
        layers = []
        # get the number of layers in the network and the shapes of the different matrices
        while current_layer:
            layer_count += 1
            w_shapes.append(current_layer.get_w().shape)
            b_shapes.append(current_layer.get_b().shape)
            layers.append(layer_count)
            current_layer = current_layer.get_next_layer()
        # iterate through the layers and do gradient checking on all elements in the weight matrices and bias vectors
        for layer_number, weight_shapes in zip(layers, w_shapes):
            print('Checking all {} x {} weights and {} biases in Layer {}'.format(weight_shapes[0], weight_shapes[1],
                                                                                  weight_shapes[0], layer_number))
            for row in range(0, weight_shapes[0]):
                self.check_forward_biases(input_data=input_data, labels=labels, check_layer=layer_number,
                                       row=row, epsilon=epsilon, bias_grads=bias_gradients, tolerance=tolerance)
                for col in range(0, weight_shapes[1]):
                    self.check_forward_weights(input_data=input_data, labels=labels, check_layer=layer_number,
                                               row=row, col=col, epsilon=epsilon, weight_grads=weight_gradients,
                                               tolerance=tolerance)

        print('\n')
        print('££££££££££££££££££££££££££££££')
        print('Check Complete. Great Success!')
        print('££££££££££££££££££££££££££££££')


if __name__ == '__main__':

    training_data = np.load('data/fashion-train-imgs.npz')
    training_labels = np.load('data/fashion-train-labels.npz')
    training_data = np.reshape(training_data, (784, 12000))
    training_labels = np.reshape(training_labels, (-1, 1)).T

    # use only the first 12 data points (to speed up gradient checking)
    training_data = training_data[:, :12]
    training_labels = training_labels[:, :12]

    model = GradientChecking()  # initialise network
    model.add_layer('relu', 784, 4)  # build the network
    model.add_layer('sigmoid', 4, 1)
    model.compile(optimiser='gradient_descent', loss='MSE', metric='binary_accuracy',
                  learning_rate=0.001)
    # perform gradient checking
    model.gradient_check(input_data=training_data, labels=training_labels, epsilon=1e-5, tolerance=1e-5)






