"""
network.py
"""

import random
import numpy as np
from numpy import linalg as LA

class Network(object):

    def __init__(self, sizes):
        """The list ``sizes`` contains the number of neurons in the
        respective layers of the network.  For example, if the list
        was [2, 3, 1] then it would be a three-layer network, with the
        first layer containing 2 neurons, the second layer 3 neurons,
        and the third layer 1 neuron.  The biases and weights for the
        network are initialized randomly, using a Gaussian
        distribution with mean 0, and variance 1.  Note that the first
        layer is assumed to be an input layer, and by convention we
        won't set any biases for those neurons, since biases are only
        ever used in computing the outputs from later layers."""
        self.num_layers = len(sizes)
        self.sizes = sizes
        self.biases = [np.random.randn(y, 1) for y in sizes[1:]]
        self.weights = [np.random.randn(y, x)
                        for x, y in zip(sizes[:-1], sizes[1:])]

    def SGD(self, training_data, epochs, mini_batch_size, learning_rate,
            test_data):
        """Train the neural network using mini-batch stochastic
        gradient descent.  The ``training_data`` is a list of tuples
        ``(x, y)`` representing the training inputs and the desired
        outputs.  """
        print("Initial test accuracy: {0}".format(self.one_label_accuracy(test_data)))
        n = len(training_data)
        for j in range(epochs):
            random.shuffle(list(training_data))
            mini_batches = [
                training_data[k:k+mini_batch_size]
                for k in range(0, n, mini_batch_size)]
            for mini_batch in mini_batches:
                self.update_mini_batch(mini_batch, learning_rate)
            print ("Epoch {0} test accuracy: {1}".format(j, self.one_label_accuracy(test_data)))

    def SGD_with_plot(self, training_data, epochs, mini_batch_size, learning_rate,
            test_data):
        """Train the neural network using mini-batch stochastic
        gradient descent.  The ``training_data`` is a list of tuples
        ``(x, y)`` representing the training inputs and the desired
        outputs.  """
        training_accuracy = []
        training_loss = []
        test_accuracy = []
        n = len(training_data)
        for j in range(epochs):
            random.shuffle(list(training_data))
            mini_batches = [
                training_data[k:k+mini_batch_size]
                for k in range(0, n, mini_batch_size)]
            for mini_batch in mini_batches:
                self.update_mini_batch(mini_batch, learning_rate)
            training_accuracy.append(self.one_hot_accuracy(training_data))
            training_loss.append(self.loss(training_data))
            test_accuracy.append(self.one_label_accuracy(test_data))
        return training_accuracy, training_loss, test_accuracy

    def SGD_with_gradient_norms(self, training_data, epochs, mini_batch_size, learning_rate,
            test_data):
        """Train the neural network using mini-batch stochastic
        gradient descent.  The ``training_data`` is a list of tuples
        ``(x, y)`` representing the training inputs and the desired
        outputs.  """
        gradient_norms = []
        print("Initial test accuracy: {0}".format(self.one_label_accuracy(test_data)))
        n = len(training_data)
        for j in range(epochs):
            random.shuffle(list(training_data))
            mini_batches = [
                training_data[k:k+mini_batch_size]
                for k in range(0, n, mini_batch_size)]
            for mini_batch in mini_batches:
                gradient_norms_epoch = self.update_mini_batch_with_gradient_norm(mini_batch, learning_rate)
            gradient_norms.append(gradient_norms_epoch)
            print ("Epoch {0} test accuracy: {1}".format(j, self.one_label_accuracy(test_data)))

        return gradient_norms

    def update_mini_batch(self, mini_batch, learning_rate):
        """Update the network's weights and biases by applying
        stochastic gradient descent using backpropagation to a single mini batch.
        The ``mini_batch`` is a list of tuples ``(x, y)``."""
        nabla_b = [np.zeros(b.shape) for b in self.biases]
        nabla_w = [np.zeros(w.shape) for w in self.weights]
        for x, y in mini_batch:
            delta_nabla_b, delta_nabla_w = self.backprop(x, y)
            nabla_b = [nb+dnb for nb, dnb in zip(nabla_b, delta_nabla_b)]
            nabla_w = [nw+dnw for nw, dnw in zip(nabla_w, delta_nabla_w)]
        self.weights = [w - (learning_rate / len(mini_batch)) * nw
                        for w, nw in zip(self.weights, nabla_w)]
        self.biases = [b - (learning_rate / len(mini_batch)) * nb
                       for b, nb in zip(self.biases, nabla_b)]

    def update_mini_batch_with_gradient_norm(self, mini_batch, learning_rate):
        """Update the network's weights and biases by applying
        stochastic gradient descent using backpropagation to a single mini batch.
        The ``mini_batch`` is a list of tuples ``(x, y)``."""
        nabla_b = [np.zeros(b.shape) for b in self.biases]
        nabla_w = [np.zeros(w.shape) for w in self.weights]
        for x, y in mini_batch:
            delta_nabla_b, delta_nabla_w = self.backprop(x, y)
            nabla_b = [nb+dnb for nb, dnb in zip(nabla_b, delta_nabla_b)]
            nabla_w = [nw+dnw for nw, dnw in zip(nabla_w, delta_nabla_w)]

        self.weights = [w - (learning_rate / len(mini_batch)) * nw
                        for w, nw in zip(self.weights, nabla_w)]
        self.biases = [b - (learning_rate / len(mini_batch)) * nb
                       for b, nb in zip(self.biases, nabla_b)]

        gradient_norms = [LA.norm(np.array(dl_dbi) / len(mini_batch)) for dl_dbi in nabla_b]
        return gradient_norms

    def backprop(self, x, y):
        zl, al = self.forward(x)
        deltaL = self.calc_output_error(zl, al, y)
        deltal = self.backward(zl, deltaL)

        db = deltal
        dw = self.calc_dw(al, deltal)

        return db, dw

    def one_label_accuracy(self, data):
        """Return accuracy of network on data with numeric labels"""
        output_results = [(np.argmax(self.network_output_before_softmax(x)), y)
         for (x, y) in data]
        return sum(int(x == y) for (x, y) in output_results)/float(len(data))

    def one_hot_accuracy(self,data):
        """Return accuracy of network on data with one-hot labels"""
        output_results = [(np.argmax(self.network_output_before_softmax(x)), np.argmax(y))
                          for (x, y) in data]
        return sum(int(x == y) for (x, y) in output_results) / float(len(data))


    def network_output_before_softmax(self, x):
        """Return the output of the network before softmax if ``x`` is input."""
        layer = 0
        for b, w in zip(self.biases, self.weights):
            if layer == len(self.weights) - 1:
                x = np.dot(w, x) + b
            else:
                x = sigmoid(np.dot(w, x)+b)
            layer += 1
        return x

    def loss(self, data):
        """Return the loss of the network on the data"""
        loss_list = []
        for (x, y) in data:
            net_output_before_softmax = self.network_output_before_softmax(x)
            net_output_after_softmax = self.output_softmax(net_output_before_softmax)
            loss_list.append(np.dot(-np.log(net_output_after_softmax).transpose(),y).flatten()[0])
        return sum(loss_list) / float(len(data))

    def output_softmax(self, output_activations):
        """Return output after softmax given output before softmax"""
        output_exp = np.exp(output_activations)
        return output_exp/output_exp.sum()

    def loss_derivative_wr_output_activations(self, output_activations, y):
        """Return derivative of loss with respect to the output activations before softmax"""
        return self.output_softmax(output_activations) - y

    def forward(self, x):
        zl = []
        a = x
        al = [a]
        for b, w in zip(self.biases, self.weights):
            z = np.dot(w, a) + b
            a = sigmoid(z)
            zl.append(z)
            al.append(a)
        return zl, al

    def backward(self, zl, deltaL):
        deltal = [deltaL]

        for i in reversed(range(0, len(zl)-1)):
            delta_i_plus_1 = np.array(deltal[0])
            w_i_plus_1 = np.array(self.weights[i+1])

            zi = zl[i]
            sigma_tag = sigmoid_derivative(zi)
            w_i_plus_1_T = np.transpose(w_i_plus_1)
            deltai = np.multiply(w_i_plus_1_T.dot(delta_i_plus_1), sigma_tag)
            deltal = [deltai] + deltal

        return deltal

    def calc_output_error(self, zl, al, y):
        aL = al[-1]
        gradient_of_l_wr_a = self.loss_derivative_wr_output_activations(aL, y)
        zL = zl[-1]
        sigma_tag = sigmoid_derivative(zL)
        deltaL = np.multiply(gradient_of_l_wr_a, sigma_tag)
        return deltaL

    def calc_dw(self, al, deltal):
        dw = []
        for i in range(0, self.num_layers - 1):  # num_of_weights = num_of_layers-1
            rows = self.weights[i].shape[0]
            cols = self.weights[i].shape[1]

            ai_minus1_mat = np.tile(np.array(al[i]).transpose(), (rows, 1))
            di_mat = np.tile(deltal[i], (1, cols))
            dwi = np.multiply(ai_minus1_mat, di_mat)
            dw.append(dwi)

        return dw


def sigmoid(z):
    """The sigmoid function."""
    return 1.0/(1.0+np.exp(-z))

def sigmoid_derivative(z):
    """Derivative of the sigmoid function."""
    return sigmoid(z)*(1-sigmoid(z))
