# -*- coding: utf-8 -*-
"""
linear.py: A set of implementations of linear neural layers.

@author: vharisop
"""

from __future__ import division
import numpy
import theano
import theano.tensor as T

class LinearLayer(object):
    """
    A neuron that implements a linear neuron, with an optional non-linearity
    at the activation layer.
    """
    def __init__(self, rng, inputs, n_in, n_out, act_fun=None):
        """
        Creates a new LinearLayer with the specified dimensions.

        Arguments
        ---------
        rng : numpy.random.RandomState
            A random state, in case reproducibility is needed
        inputs : theano.tensor.matrix
            A symbolic tensor describing the inputs
        n_in : int
            The input's dimensionality
        n_out : int
            The output's dimensionality
        act_fun : theano.tensor
            An optional theano tensor describing an activation function.
            If no such parameter is provided, it is assumed to be the
            identity (linear activation)
        """
        self.inputs = inputs
        # values for weights
        W_val = numpy.asarray(
            rng.uniform(
                low=-numpy.sqrt(6. / (n_in + n_out)),
                high=numpy.sqrt(6. / (n_in + n_out)),
                size=(n_in, n_out)),
            dtype=theano.config.floatX)
        W = theano.shared(value=W_val, name='W', borrow=True)
        # values for coefficients
        B_val = numpy.asarray(
            rng.uniform(low=-0.5, high=0.5, size=(n_out,)),
            dtype=theano.config.floatX)
        B = theano.shared(value=B_val, name='B', borrow=True)

        # store generated coefficients and weights
        self.B, self.W = B, W
        # output of this layer (as well as y prediction function)
        next_layer = T.dot(self.inputs, self.W) + self.B
        self.output = next_layer if not act_fun else act_fun(next_layer)
        # parameters for gradient
        self.params = [self.W, self.B]


class SigmoidLayer(LinearLayer):
    """
    A linear layer with the logistic sigmoid activation function.
    """
    def __init__(self, rng, inputs, n_in, n_out):
        """
        Creates a new SigmoidLayer.
        Type `help(LinearLayer)` for information on the constructor
        arguments.
        """
        super(SigmoidLayer, self).__init__(
            rng, inputs, n_in, n_out, T.nnet.sigmoid)


class ReluLayer(LinearLayer):
    """
    A linear layer using the RelU = max(0, x) activation function.
    """
    def __init__(self, rng, inputs, n_in, n_out):
        """
        Creates a new ReluLayer.
        Type `help(LinearLayer)` for information on the constructor
        arguments.
        """
        super(ReluLayer, self).__init__(
            rng, inputs, n_in, n_out, T.nnet.relu)


class SoftmaxLayer(LinearLayer):
    """
    A neuron that implements a softmax layer, using the negative
    log-likelihood as its cost function.
    """
    def __init__(self, rng, inputs, n_in, n_out):
        """
        Creates a new SoftmaxLayer with the specified parameters.
        Type help(LinearLayer) for complete documentation regarding
        the constructor's parameters. The activation function is
        theano's nnet.softmax implementation.
        """
        act_fun = T.nnet.softmax
        super(SoftmaxLayer, self).__init__(rng, inputs, n_in, n_out, act_fun)
        self.pred_y = T.argmax(self.output, axis=1)


    def output_at(self, patterns, inputs=None):
        """
        Computes the response of this unit, given a set of input patterns.

        Arguments
        ---------
        patterns : numpy.array
            An array containing the input patterns
        inputs : theano.tensor.var.TensorVariable
            A symbolic tensor containing the input (optional)

        Returns
        -------
        numpy.array
            The response of this layer
        """
        response = theano.function(
            inputs=[self.inputs if not inputs else inputs],
            outputs=self.pred_y)
        return response(patterns)


    def validation_error(self, Y):
        """
        Computes a 0-1 loss over a set of patterns
        aimed for model validation.

        Arguments
        ---------
        Y : theano.tensor.var.TensorVariable
            A symbolic tensor of the validation outputs

        Returns
        -------
        theano.tensor.var.TensorVariable
            A symbolic tensor holding the validation loss
        """
        return T.sum(T.neq(self.pred_y, Y)) / T.shape(Y)


    def error(self, Y):
        """
        Computes the negative log likelihood, given a tensor Y describing
        the outputs of this layer.

        Arguments
        ---------
        Y : theano.tensor.var.TensorVariable
            A symbolic tensor describing the outputs of this layer.

        Returns
        -------
        float
            The negative log likelihood error
        """
        # T.log(self.output) -> numpy array
        # [T.arange(..), Y] -> slice that array keeping the Y'th
        # column and all the rows of the batch
        return -T.mean(
            T.log(self.output)[T.arange(Y.shape[0]), Y])
