#### Libraries
# Third Party Libraries
import numpy as np
from sklearn.model_selection import train_test_split
import theano
import theano.tensor as T
from theano.tensor.nnet import conv2d
from theano.tensor.signal import pool

"""
Explanation: LeNET5
We describe how to built a convolution neural network similar to that described in <u>Gradient Based Learning Applied to Document Recognition</u> by LeCunn et al.. We built all the necessary layers/components  using Theano and test various gradient descent methods on the MNIST dataset. 
A convolution neural network (CNN) is a type of feed forward artificial neural netowrk that that utilizes convolutions to exploit <b>local connectivity</b>, in addition <b>weight sharing</b>. We begin by describing a discrete convolution in 1-D, followed by defining local connectivity and weigth sharing. 
A 1-D discrete convolution is described as 
<h5 align=center>
$
\begin{align}
o[n] &= f[n]\ast g[n]\\ 
&= \sum\limits_{u\in \mathbb{R}} f[u]g[n-u]\\
&= \sum\limits_{u\in \mathbb{R}} f[n-u]g[u]\\
\end{align}
$
</h5>

<br><b>Local connectivity:</b> A CNN "ensures that the learnt "filters" produce the strongest response to a spatially local input pattern. Stacking many such layers leads to non-linear "filters" that become increasingly "global". This allows the network to first create good representations of small parts of the input, then assemble representations of larger areas from them. 
<img src="https://github.com/LukaszObara/LeNet5/blob/master/images/LocalConnectivity.png" style="width:auto;height:128px;">
<br><b>Shared wieghts:</b> Given a filter we run the same filter for all positions in the image. In other words, all the pixel positions “share” the same parameterization (weight vector and bias) and form a feature map. This means that all the neurons in a given convolutional layer detect exactly the same feature. Replicating units in this way allows for features to be detected regardless of their position in the visual field, thus constituting the property of translation invariance." 
Architecture
The architecture that we will be building consists of a <b>convolution layer</b> followed by a <b>pooling layer</b>, a convolution layer, a pooling layer, a <b>fully connected layer</b>, and finaly a predictive layer. The model can be depicted by:
<img src="architec2.png" style="width:auto;height:228px;">
The <b>convolution layer</b> is the workhorse of the the CNN and does most of the computation. "The conv layer's paramaters consist of a set of learnable filters, where every filter is small spatially, but extend through the full depth of the input volume. During the forward pass the filter slides accross the input computing a dot product between the filter and input at any given position. The network will learn filters that activate when a feature of interest is detected."
The <b>pooling layers</b> is a downsampling operation which acts by reducing the spatial size of the representation while attempting to retain important information. Pooling works by partiting the input into a set of non-overlapping rectangles and, for each sub-region outputs one of the following: average of the values within a given sub-region, taking the largest element, or by computing the $l_2$ norm. 
The <b>fully connectected layer</b> (FC) is a traditional Multi Layer Perceptron. The term “Fully Connected” implies that every neuron in the previous layer is connected to every neuron on the next layer. The inputs from the last conv/pool layer are flattened and passed through the FC-layer. The activations are computed with a matrix multiplication followed by a bias offset.
Theano Code
Importing necessary libraries
End of explanation
"""


class ConvLayer(object):

    def __init__(self, input, filter_shape, image_shape, padding=(0, 0), 
                 stride=(1, 1), activation_fn=None):

        assert image_shape[1] == filter_shape[1]

        # rng = np.random.RandomState(seed)

        self.input = input
        self.filter_shape = filter_shape
        self.image_shape = image_shape
        self.activation_fn = activation_fn

        fan_in = np.prod(filter_shape[1:])
        fan_out = filter_shape[0]*np.prod(filter_shape[2:]) // 2
        W_bound = np.sqrt(6/(fan_in+fan_out))
        w = np.random.uniform(low=-W_bound, high=W_bound, size=filter_shape)
        b_vals = np.random.uniform(size=filter_shape[0])

        # Initiliaze weights with random variables
        self.W = theano.shared(name='weights',
                               value=w.astype(theano.config.floatX),
                               borrow=True)
        self.b = theano.shared(name='bias',
                               value=b_vals.astype(theano.config.floatX), 
                               borrow=True)

        conv_out = conv2d(input=input, filters=self.W, border_mode=padding,
                          subsample=stride, filter_shape=filter_shape, 
                          input_shape=image_shape)

        l_output = conv_out + self.b.dimshuffle(('x', 0, 'x', 'x'))
        self.output = (l_output if activation_fn is None 
                       else activation_fn(l_output))

        # Parameters of the model
        self.params = [self.W, self.b]

"""
Explanation: Convolution Layer
The convolution layer performs a convolution operation on the input. It needs to take into accout the following arguments:
<b>input:</b> The input feature map.
<br><b>filter_shape:</b> The shape of the filter.
<br><b>padding:</b> Inserts a 0-valued border of appropriate size around the input (image_shape). This gives us a little more control over the size of the feature map.  <br><b>stride:</b> The stride is the number of pixels by which we slide our filter over the input (image_shape).
<br><b>activation_fn:</b> A nonlinear function used to transform our input; generally representing the rate of action potential firing for a given pixel. The purpose of which is to ensure that the representation in the input is mapped to a different space in the output.
End of explanation
"""


class PoolingLayer(object):

    def __init__(self, input, pool_shape=(2, 2), ignore_border=True,
                 activation_fn=None):
        self.input = input
        self.pool_shape = pool_shape
        self.ignore_border = ignore_border

        l_output = pool.pool_2d(input=input, ds=pool_shape, 
                                ignore_border=self.ignore_border)

        self.output = (l_output if activation_fn is None 
                       else activation_fn(l_output))

"""
Explanation: Pooling layer
As described above the pooling layer is a downsampling operation, as such we decided to use MAX value in each sub-region. Hence, to construct this layer we need to take the following arguments into consideration:
<b>input:</b> The input feature map. 
<br><b>pool_shape:</b> The size of the window. 
<br><b>ignore_border:</b> If set to True then right and/or bottom row/column will be ignored if it cannot be fully covered by the window (pool_shape).
<br><b>activation_fn:</b> A nonlinear function used to transform our input; generally representing the rate of action potential firing for a given pixel. The purpose of which is to ensure that the representation in the input is mapped to a different space in the output.
End of explanation
"""


class FC(object):

    def __init__(self, input, n_in, n_out, W=None, b=None, seed=35,
                 activation_fn=None):

        # rng = np.random.RandomState(seed)

        self.input = input

        if W is None:
            W_values = np.random.uniform(low=-np.sqrt(6./(n_in+n_out)),
                                   high=np.sqrt(6./(n_in+n_out)),
                                   size=(n_out, n_in)).astype(theano.config.floatX)

            if activation_fn == theano.tensor.nnet.sigmoid:
                W_values *= 4

            W = theano.shared(name='Weights', value=W_values, borrow=True)


        if b is None:
            b_values = np.zeros(n_out, dtype=theano.config.floatX)
            b = theano.shared(name='bias', value=b_values, borrow=True)

        self.W = W
        self.b = b

        l_output = (T.dot(self.W, input.T)).T + self.b
        self.output = (l_output if activation_fn is None 
                       else activation_fn(l_output))

        self.params = [self.W, self.b]

"""
Explanation: Fully Connected Layer
The Fully Connected layer is a traditional Multi Layer Perceptron hence we need to take into account the following arguments:
<b>input:</b> The input feature map
<br><b>n_in:</b> The number of input neurons 
<br><b>n_out:</b> The number of output neurons.
<br><b>W:</b> (Optional) Load pretrained weights
<br><b>b:</b> (Optional) Load pretrained biases.
<br><b>activation_fn:</b> A nonlinear function used to transform our input; generally representing the rate of action potential firing for a given pixel. The purpose of which is to ensure that the representation in the input is mapped to a different space in the output.
End of explanation
"""


def elu(x, alpha=1.0):
    return T.switch(x > 0, x, T.exp(x)-1)

def l2_reg(x, lmbd=0.05):
    """
    L_2 regularization 
    """

    l2 = 0
    for elements in x:
        l2 += T.sum(elements[0]**2)

    return lmbd / 2 * l2 

"""
Explanation: Building The Model
Activation Function & Regularization
We decided to use the Exponential Linear Unit (ELU) as described in <u>Fast and accurate deep network learning by exponetial lienar units</u> by Clevert et al., since this non-linear function displayed quicker convergence and yielded a higher validation score relative to a sigmoid, tanh, or ReLu function. Furthermore, to reduce overfitting we used $l_2$ regularization which acts by penalizing our model by adding a complexity term that would yields a higher loss for more complex models. Both the activation function and regularization method or defined below.
End of explanation
"""


X = T.tensor4(name='X', dtype=theano.config.floatX) 
Y = T.imatrix(name='Y')
y = T.ivector(name='y')
lr = T.scalar(name='learning_rate', dtype=theano.config.floatX)

nkerns = [8, 32]
batch_size = 256
act_f = elu

conv_layer1 = ConvLayer(input=X, 
                        filter_shape=(nkerns[0], 1, 3, 3), 
                        image_shape=(batch_size, 1, 28, 28),
                        activation_fn=None)
pool_layer1 = PoolingLayer(input=conv_layer1.output,
                           activation_fn=act_f)
conv_layer2 = ConvLayer(input=pool_layer1.output,
                        filter_shape=(nkerns[1], nkerns[0], 5, 5), 
                        image_shape=(batch_size, nkerns[0], 13, 13),
                        activation_fn=None)
pool_layer2 = PoolingLayer(input=conv_layer2.output,
                           activation_fn=act_f)

# outputs from convolution network need to be flattend before being 
# passed through to the the fully-connected layer
fc_layer_input = pool_layer2.output.flatten(2) 

fc_layer1 = FC(input=fc_layer_input,
                n_in=nkerns[1] * 4 * 4,
                n_out=512,
                activation_fn=act_f)
fc_layer2 = FC(input=fc_layer1.output,
               n_in=512,
               n_out=10,
               activation_fn=act_f)

"""
Explanation: Convolutional Neural Network
End of explanation
"""


params = fc_layer2.params + fc_layer1.params\
       + conv_layer2.params + conv_layer1.params
    
cost_input = T.nnet.nnet.softmax(fc_layer2.output)
cost = T.mean(T.nnet.nnet.categorical_crossentropy(cost_input, Y)) \
     + l2_reg(params)

grads = T.grad(cost, params)

"""
Explanation: Cost Function
Since we will be performing classification we will be using the categorical cross entropy function as our cost function which is also the function which we seek to minimize. The cross entropy function is defined to be:
<h5 align="center">
$
\begin{equation}
H(p, q) = - \sum\limits_{x} p(x)\log\big(q(x)\big)
\end{equation}
$
</h5>

where $p(x)$ is the true distribution (the true label) and $q(x)$ is coding distribution (the predicted value). Due to $\log(\cdot) \ \in \ (0, \infty)$ we need to have $q(x) \in (0, 1)$ and this is accomplished by passing our input into a softmax, a generalization of the logistic function that "squashes" a K-dimensional vector $\mathbf{z}$  of arbitrary real values to a K-dimensional vector $\sigma(\mathbf{z})$ of real values in the range $(0, 1)$ that add up to 1 and is defined to be 
<h5 align="center">
$
\begin{equation}
\sigma(\mathbf{z}) = \frac{\exp(z_j)}{\sum_{k=1}^K \exp(z_k)} \text{ for } j =1\ldots K
\end{equation}
$
</h5>

In order to minimize the cross entropy function we need to take its gradient with respect to the paramaters in the network, which are those present in both fully connected layers and in both convolution layers. 
The code below show how to accomplish this using Theano. We begin by defining the paramaters in the network. We then pass the output of the final fully connected layer into a softmax function to generate a probaility distribution which is followed by passing our softmax values into the categorical cros sentropy function with $l_2$ regularization. Finally, we call tensor.grad() to get the gradient with respect to params.
End of explanation
"""


def sgd(l_rate, parameters, grads): 
    
    updates = []
    for param, grad in zip(parameters, grads):
        updates.append((param, param - l_rate * grad))

    return updates

"""
Explanation: Stochastic Gradient Descent Methods (SGD)
We present three (3) different SGD algorithm used to minimize our cost function. For more option we encourage the reader to see <a href="https://github.com/LukaszObara/Theano_gradient_descent">Theano_gradient_descent</a>. 
Vanilla SGD
This is the simplest form of update, it changes the parameters along the negative gradient direction. Vanilla SGD suffers from a slow convergence and displays a difficutly escaping from a saddle point. 
The update rule is given by:
<h5 align="center">
$
x_{t+1} = x_{t} - \eta \nabla C
$
</h5>
End of explanation
"""


def momentum(l_rate, parameters, grads, momentum=0.9):
    
    def update_rule(param, velocity, df):
        v_next = momentum * velocity - l_rate * df
        updates = (param, param+v_next), (velocity, v_next)

        return updates

    assert momentum <=1 and momentum >= 0

    velocities = [theano.shared(name='v_%s' % param,
                                value=param.get_value() * 0., 
                                broadcastable=param.broadcastable) 
                  for param in parameters]

    updates = []
    for p, v, g in zip(parameters, velocities, grads):
        param_updates, vel_updates = update_rule(p, v, g)
        updates.append(param_updates)
        updates.append(vel_updates)

    return updates

"""
Explanation: Momentum
Momentum is a varient of Vanilla SGD that helps accelerate SGD in the relevant direction and dampens oscillations. It does this by adding a fraction $\gamma$ of the update vector of the past time step to the current update vector. Nevertheless, like Vanilla SGD it still has difficulty escaping a saddle point. 
The update rule is given by:
<h5 align="center">
$\begin{align}
v_{t+1} &= \gamma v_t + \eta \nabla C\\
x_{t+1} &=  x_t- v_{t+1}
\end{align}$
</h5>

where $v_t$ can be considered as the velocity.
End of explanation
"""


def rmsprop(l_rate, d_rate=0.9, epsilon=1e-6, parameters=None, grads=None):

    one = T.constant(1.0)

    def update_rule(param, cache, df):
        cache_val = d_rate * cache + (one-d_rate) * df**2
        x = l_rate * df / (T.sqrt(cache_val) + epsilon)
        updates = (param, param-x), (cache, cache_val)

        return updates

    caches = [theano.shared(name='c_{}'.format(param),
                                value=param.get_value() * 0., 
                                broadcastable=param.broadcastable) 
              for param in parameters]

    updates = []
    for p, c, g in zip(parameters, caches, grads):
        param_updates, cache_updates = update_rule(p, c, g)
        updates.append(param_updates)
        updates.append(cache_updates)

    return updates

"""
Explanation: RMSProp
The update rule that yielded the best results was RMSProp, which is an adaptive learning rate method. This allows the the learning rate to be adaptively tuned after each epoch, eliminating the need for any type of annealing (not discussed). In addition adaptive learning rates are well-behaved for a broader range of hyperparameter values than the methods presents above. In the case of RMSProp the adaptive learning is accomplished by dividing the learning rate by an exponentially decaying average of squared gradients. Relative to the methods discussed above it does not seem to suffer from being stuck in a saddle point. 
The update rule is given by:
<h5 align="center">
$\begin{align}
E[\nabla C^2]_{t} &= \delta*E[\nabla C^2]_{t-1} + (1-\delta)\nabla C^2_t\\
x_{t+1} &=  x_t- \eta\frac{1}{\sqrt{E[g^2]_t +\epsilon}}\nabla C_t
\end{align}$
</h5>

where $E[\nabla C^2]$ is the cache value that stores a weighted sum of previous cache values and gradients squared. It is then used to normalize the parameter update step, element-wise. $\delta$ is the decay rate.
End of explanation
"""


train = theano.function(inputs=[X, Y, lr], outputs=cost, 
                        updates=rmsprop(l_rate=lr, parameters=params, 
                                     grads=grads),
                        allow_input_downcast=True)

# Validation results
pred_result = cost_input.argmax(axis=1)
accu = theano.function(inputs=[X, y], outputs=T.sum(T.eq(pred_result, y)), 
                       allow_input_downcast=True)

pred = theano.function(inputs=[X], outputs=pred_result, 
                       allow_input_downcast=True)

"""
Explanation: Training and Validating
To train the network we need to define a training and validation function using theano.function(). Using RMSProp as our SGD method we create the necessary functions below.  All the functions are defined in a similar manner, both validation function and a prediction functions omit the update argument, since they are used for pushing the data forward.
End of explanation
"""


def train_model(training_data, validation_data, test_data=None,
                learning_rate=1e-4, epochs=100):

    print('---Training Model---')
    predicted_results = [] 

    total_values, total_val_values = len(training_data), len(validation_data)

    for epoch in range(epochs):
        print('Currently on epoch {}'.format(epoch+1))
        np.random.shuffle(training_data)

        mini_batches = [training_data[k: k+batch_size]
                        for k in range(0, total_values, batch_size)] 
        validation_batches = [validation_data[m: m+batch_size]
                              for m in range(0, total_val_values, batch_size)]

        training_cost, accuracy = 0, 0
        training_cost_list, accuracy_list = [], []

        for mini_batch in mini_batches:
            labels = mini_batch[:, 0]
            label_matrix = np.zeros(shape=(256, 10), dtype=theano.config.floatX)

            for i, label in enumerate(labels):
                vec = scalar_to_vec(int(label), 10)
                label_matrix[i] = vec

            digits = mini_batch[:, 1:]/255
            digits = digits.reshape(-1, 1, 28, 28)
            cost_ij = train(digits, label_matrix, learning_rate)
            training_cost += cost_ij

        for val_batch in validation_batches:
            labels = mini_batch[:, 0]
            label_matrix = np.zeros(shape=(256, 10), dtype=theano.config.floatX)

            for i, label in enumerate(labels):
                vec = scalar_to_vec(int(label), 10)
                label_matrix[i] = vec

            digits = mini_batch[:, 1:]/255
            digits = digits.reshape(-1, 1, 28, 28)
            accuracy += accu(digits, labels)

        training_cost_list.append(training_cost/total_values)
        accuracy_list.append(accuracy/total_val_values)

        print('The accuracy is: {}'.format(accuracy/total_val_values))
        print('The loss is: {}'.format(training_cost/total_values))
        print('--------------------------')

    if np.any(test_data):
        print('===================================')
        print('Using test data to predict results')
        total_values = len(test_data)

        mini_batches = [test_data[k: k+batch_size]
                        for k in range(0, total_values, batch_size)] 

        for mini_batch in mini_batches:
            digits = mini_batch[:, :]/255
            digits = digits.reshape(-1, 1, 28, 28)
            result = pred(digits)
            predicted_results = np.append(predicted_results, result)

        print('Done')

    return training_cost_list, accuracy_list, predicted_results

"""
Explanation: Training the Model
The last stage involves creating a function that shuffles and partitions into appropriate sized batches before being processed further during each epoch.
End of explanation
"""


