import numpy as np
import gzip
import cPickle as pickle
import cv2

class Layer:
    """
    Abstract class for the Layer structure
    """

    def __init__(self, name, input_shape):
        """
        The constructor for the layer
        :param name: The name of the layer
        :param input_shape: A tuple representing the input shape of the layer
        """
        self.name = name
        self.is_trainable = True
        self.input_shape = input_shape
        self.cache = []
        self.opt_cache = []

    def set_trainable(self, val):
        """
        Set whether the layer is trainable or not
        :param val: Boolean value
        """
        self.is_trainable = val

    def make_cache(self, inputs, outputs):
        return None

    def get_outputs(self, inputs):
        return None

    def forward_pass(self, inputs):
        """
        Makes a forward pass through the layer
        :param inputs: The input to be passed forward
        :return: The output of the forward pass
        """
        outputs = self.get_outputs(inputs)
        self.make_cache(inputs, outputs)
        return outputs

    def update_params(self, top_delta, l2, alpha, beta1, beta2, epsilon):
        return None

    def get_delta(self, top_delta):
        return None

    def backward_pass(self, top_delta, l2, alpha, beta1, beta2, epsilon):
        """
        Makes a backward pass through the layer
        :param top_delta: The delta of the next layer
        :param l2: The L2 norm weighting for regularization
        :param alpha: The alpha parameter for the ADAM optimizer
        :param beta1: The beta parameter for the ADAM optimizer
        :param beta2: The beta2 parameter for the ADAM optimizer
        :param epsilon: The epsilon parameter for the ADAM optimizer
        :return: The delta generated by the backward pass
        """
        if self.is_trainable:
            self.update_params(top_delta, l2, alpha, beta1, beta2, epsilon)
        return self.get_delta(top_delta)

    def backward_pass_rel(self, top_rel, index=-1, epsilon=1e-4):
        """
        Backward passing the relevance
        :param top_rel: The relevance of the next layer
        :param index: The index of the node to be backpropagated. If -1, then do all
        :param epsilon: The smoothing factor
        :return: The relevance of this layer
        """
        return top_rel


class FCLayer(Layer):
    """
    A layer representing a linear transformation of vectors
    """

    def __init__(self, name, input_shape, n_units):
        Layer.__init__(self, name, input_shape)
        self.n_units = n_units
        stddev = np.sqrt(2.0 / (self.input_shape[0] + self.n_units))
        self.weights = np.random.normal(0.0, stddev, size=(self.input_shape[0], self.n_units))
        self.bias = np.zeros((self.n_units,))
        self.output_shape = (self.n_units,)
        self.opt_cache = [np.zeros(self.weights.shape), np.zeros(self.weights.shape), np.zeros(self.n_units),
                          np.zeros(self.n_units), 0]
        self.rel_cache = []

    def get_outputs(self, inputs):
        return np.dot(inputs, self.weights) + self.bias

    def make_cache(self, inputs, outputs):
        self.cache = [inputs]
        self.rel_cache = [inputs, self.weights, outputs]

    def get_delta(self, top_delta):
        return np.dot(top_delta, self.weights.T)

    def backward_pass_rel(self, top_rel, index=-1, epsilon=1e-4):
        if index == -1:
            return np.sum((self.rel_cache[0].T.dot(top_rel) * self.rel_cache[1]) / (self.rel_cache[2] + epsilon),
                          axis=1,
                          keepdims=True).T
        else:
            return np.sum(
                (self.rel_cache[0].T.dot(top_rel) * self.rel_cache[1][:, index: index + 1]) / (
                self.rel_cache[2][:, index] + epsilon),
                axis=1,
                keepdims=True).T

    def update_params(self, top_delta, l2, alpha=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8):
        grad_w = np.dot(self.cache[0].T, top_delta) + l2 * self.weights
        grad_b = top_delta.sum(axis=0) + l2 * self.bias
        self.opt_cache[0] = beta1 * self.opt_cache[0] + (1 - beta1) * grad_w
        self.opt_cache[2] = beta1 * self.opt_cache[2] + (1 - beta1) * grad_b
        self.opt_cache[1] = beta2 * self.opt_cache[1] + (1 - beta2) * (grad_w ** 2)
        self.opt_cache[3] = beta2 * self.opt_cache[3] + (1 - beta2) * (grad_b ** 2)
        self.opt_cache[4] += 1
        m_w = self.opt_cache[0] / (1 - beta1 ** self.opt_cache[4])
        m_b = self.opt_cache[2] / (1 - beta1 ** self.opt_cache[4])
        v_w = self.opt_cache[1] / (1 - beta2 ** self.opt_cache[4])
        v_b = self.opt_cache[3] / (1 - beta2 ** self.opt_cache[4])
        self.weights -= alpha * m_w / (np.sqrt(v_w) + epsilon)
        self.bias -= alpha * m_b / (np.sqrt(v_b) + epsilon)


class Sigmoid(Layer):
    """
    A layer for the sigmoid activation
    """

    def __init__(self, name, input_shape):
        Layer.__init__(self, name, input_shape)
        self.output_shape = self.input_shape

    def get_outputs(self, inputs):
        return 1.0 / (1 + np.exp(-inputs))

    def make_cache(self, inputs, outputs):
        self.cache = [outputs]

    def get_delta(self, top_delta):
        return top_delta * self.cache[0] * (1 - self.cache[0])


class Tanh(Layer):
    """
    A layer for the tanh activation
    """

    def __init__(self, name, input_shape):
        Layer.__init__(self, name, input_shape)
        self.output_shape = self.input_shape

    def get_outputs(self, inputs):
        return np.tanh(inputs)

    def make_cache(self, inputs, outputs):
        self.cache = [outputs]

    def get_delta(self, top_delta):
        return top_delta * (1 - self.cache[0] ** 2)


class Relu(Layer):
    """A layer for the Relu activation"""

    def __init__(self, name, input_shape):
        Layer.__init__(self, name, input_shape)
        self.output_shape = self.input_shape

    def get_outputs(self, inputs):
        return inputs * (inputs > 0)

    def make_cache(self, inputs, outputs):
        self.cache = [(inputs > 0)]

    def get_delta(self, top_delta):
        return top_delta * self.cache[0]


class Softmax(Layer):
    """
    A layer for a softmax classifier
    """

    def __init__(self, name, input_shape):
        Layer.__init__(self, name, input_shape)
        self.output_shape = self.input_shape

    def get_outputs(self, inputs):
        x = np.exp(inputs - np.max(inputs, axis=1, keepdims=True))
        return x / x.sum(axis=1, keepdims=True)

    def make_cache(self, inputs, outputs):
        self.cache = [outputs]

    def get_delta(self, truth):
        return (self.cache[0] - truth) / self.cache[0].shape[0]

    def get_cost(self, truth):
        return np.mean(np.sum(-truth * np.log(self.cache[0] + 1e-20), axis=1), axis=0)

    def backward_pass(self, truth, l2, alpha, beta1, beta2, epsilon):
        return self.get_cost(truth), self.get_delta(truth)


def forward(layer_list, input):
    """
    Forward passes input through a list of layers
    :param layer_list: A list of layers of the network
    :param input: The input
    :return: The output of the network
    """
    x = layer_list[0].forward_pass(input)
    for i in layer_list[1:]:
        x = i.forward_pass(x)
    return x


def backward(layer_list, truth, l2, alpha=0.001, beta1=0.9, beta2=0.999, epsilon=1e-8):
    """
    Backward passes through the network, and trains it with the ADAM optimizer
    :param layer_list: A list of the layers of the network
    :param truth: The truth value
    :param l2: The weighting for the L2 regularization
    :param alpha: The params of the ADAM optimizer
    :param beta1: The params of the ADAM optimizer
    :param beta2: The params of the ADAM optimizer
    :param epsilon: The params of the ADAM optimizer
    :return: The cost of the pass
    """
    cost, delta = layer_list[-1].backward_pass(truth, l2, alpha, beta1, beta2, epsilon)
    for i in layer_list[::-1][1:]:
        delta = i.backward_pass(delta, l2, alpha, beta1, beta2, epsilon)
    return cost


def get_rel(layer_list, input, epsilon=1e-4):
    """
    Gets the relevance map for an input
    :param layer_list: The list of layers of the network
    :param input: Must be a single instance, but a 2d matrix
    :param epsilon: The epsilon parameter
    :return: The relevance map
    """
    output = forward(layer_list, input)
    maps = []
    for ind, i in enumerate(output[0]):
        temp = np.asarray(i).reshape((1, 1))
        temp = layer_list[-2].backward_pass_rel(temp, ind, epsilon)
        for layer in layer_list[::-1][3:]:
            temp = layer.backward_pass_rel(temp, -1, epsilon)
        maps.append(temp)
    return maps


Model = [FCLayer("FC1", (784,), 128), Relu("Relu", (128,)), FCLayer("FC2", (128,), 10), Softmax("Softmax1", (10,))]
Data = pickle.load(gzip.open("/home/tanmaya/mnist_one_hot.pkl.gz", "rb"))
TrainX, TrainY = Data[0]
iters = 100
batch_size = 500
for i in xrange(iters):
    cost = []
    for j in xrange(50000 // batch_size):
        _ = forward(Model, TrainX[j * batch_size: (j + 1) * batch_size])
        cost.append(backward(Model, TrainY[j * batch_size: (j + 1) * batch_size], 0.0001))
    print "Training cost is {}".format(np.mean(cost))

indices = np.arange(50000)
np.random.shuffle(indices)
for i in indices[:5]:
    maps = get_rel(Model, TrainX[i: i + 1])
    for j, m in enumerate(maps):
        m = (m - np.min(m)) * 255.0 / (np.max(m) - np.min(m))
        cv2.imwrite("Examples/map{}_{}.jpg".format(np.nonzero(TrainY[i])[0], j), np.asarray(m, dtype=np.uint8).reshape(28, 28))
