"""autoencoders"""

import logging

import numpy as np

import theano.tensor as T
from theano.tensor.nnet import sigmoid

from . import Model, SpeedLayer

log = logging.getLogger(__name__)


class AELayer( SpeedLayer ):
    """An autoencoder layer"""

    def __init__(self, X, n_in, n_hidden, rng, thrng, dtype, clevel=0):
        """initialize an AE instance

        @param X: input (theano.tensor)
        @param n_in: input dimension (int)
        @param n_hidden: hiddden dimension (int)
        @param rng: random state (np.random.RandomState)
        @param thrng: random state (theano.tensor.shared_randomstreams.RandomStreams)
        @param dtype: dtype of weights (np.dtype)
        @param clevel: keeps ``1-corruption_level`` entries of the inputs the same
                         and zero-out randomly selected subset of size ``coruption_level``
        """

        super(AELayer, self).__init__([(n_in, n_hidden), (n_hidden,), (n_in,)],
                ["W", "b", "b_prime"], [0.0, 0.0, 0.0], [dtype, dtype, dtype])

        if type(rng) == np.random.RandomState:
            log.debug("AEW: rnd")
            thr = np.sqrt( 6. / (n_in+n_hidden) )
            self._weights_[0].set_value( np.asarray(rng.uniform(low=-thr, high=thr,
                size=(n_in, n_hidden)), dtype=dtype) )

        self.theano_rng = thrng
        self.input = X
        self.corruption_level = clevel
        self.tilde_input = self.corrupt( self.input, clevel )

        w, b, bp = self.get_params()
        self.encoder = sigmoid( T.dot(self.tilde_input, w) + b )
        self.decoder = sigmoid( T.dot(self.encoder, w.T) + bp )

    def activation(self):
        return self.encoder

    def corrupt(self, X, corrupt_level):
        """theano function that adds binomial noise to the given input

        first argument of theano.rng.binomial is the shape(size) of
        random numbers that it should produce
        second argument is the number of trials
        third argument is the probability of success of any trial
        this will produce an array of 0s and 1s where 1 has a probability of
        1 - ``corruption_level`` and 0 with ``corruption_level``
        """

        return self.theano_rng.binomial(size=X.shape, n=1, p=1-corrupt_level) * X
    
    def cost(self, l1, l2):
        l1_term = l1 * self.weight_norm("l1")
        l2_term = l2 * self.weight_norm("l2")
        
        return self.reconstruction_cost() + l1_term + l2_term

    def reconstruction_cost(self):
        z = self.decoder
        x = self.input
        # note : we sum over the size of a datapoint; if we are using minibatches,
        #        L will  be a vector, with one entry per example in minibatch
        L = -T.sum(x * T.log(z) + (1 - x) * T.log(1 - z), axis=1 )
        # note : L is now a vector, where each element is the cross-entropy cost
        #        of the reconstruction of the corresponding example of the
        #        minibatch. We need to compute the average of all these to get
        #        the cost of the minibatch
        return T.mean(L)

    def __str__(self):
        return "[(AE) %dx%d]" % self.get_weights()[0].shape

class AutoEncoder( Model ):
    """Denoising autoencoder

    A denoising autoencoders tries to reconstruct the input from a corrupted
    version of it by projecting it first in a latent space and reprojecting
    it afterwards back in the input space. Please refer to Vincent et al.,2008
    for more details. If x is the input then equation (1) computes a partially
    destroyed version of x by means of a stochastic mapping q_D. Equation (2)
    computes the projection of the input into the latent space. Equation (3)
    computes the reconstruction of the input, while equation (4) computes the
    reconstruction error.

    .. math::

       \tilde{x} ~ q_D(\tilde{x}|x)                                     (1)

       y = s(W \tilde{x} + b)                                           (2)

       x = s(W' y  + b')                                                (3)

       L(x,z) = -sum_{k=1}^d [x_k \log z_k + (1-x_k) \log( 1-z_k)]      (4)

    """

    def __init__(self, ins, hs, rng, theano_rng, wdtype, corruption_level):
        """ Initialize the dA class by specifying the number of visible units (the
        dimension d of the input ), the number of hidden units ( the dimension
        d' of the latent or hidden space ) and the corruption level.

        @param ins: input dimension (int)
        @param hs: hidden dimension (int)
        @param rng: random state (np.random.RandomState)
        @param theano_rng: random state (theano.tensor.shared_randomstreams.RandomStreams)
        @param wdtype: dtype of weights (np.dtype)
        @param corruption_level: keeps ``1-corruption_level`` entries of the inputs the same
                         and zero-out randomly selected subset of size ``coruption_level``

        """

        X = T.matrix("X")
        super(AutoEncoder, self).__init__( [AELayer(X, ins, hs, rng, theano_rng,
            wdtype, corruption_level)] )

    def cost(self, l1, l2):
        return self[0].cost(l1, l2)


    def update_params(self, train_batches, gradient_f, momentum, lrate):
        """step on the direction of the

        step on the direction of the gradient for
        a whole epoch and update the model params in place. 
        By definition speed is initialized to 0.
        new_speed = -rho * dE/dw + mu * speed
        new_weight = w + new_speed

        @param train_batches: indexes of batches (list)
        @param gradient_f: function that returns the list of gradients from the batch index.
        @param momentum: mu
        @param lrate: rho
        @return: none
        """

        for batch_i in train_batches:
            all_grads = gradient_f(batch_i)
            for layer in self:
                l_grads = map(lambda i: all_grads.pop(0),
                        range(len(layer.get_params())))

                layer.speed_update(l_grads, momentum, lrate)
                layer.weight_update()
                
class AEStack( Model ):
    """a stack of denoising autoencoders. Each layer is a denoising autoencoder.
    A denoising autoencoders tries to reconstruct the input from a corrupted
    version of it by projecting it first in a latent space and reprojecting
    it afterwards back in the input space. Please refer to Vincent et al.,2008
    for more details. If x is the input then equation (1) computes a partially
    destroyed version of x by means of a stochastic mapping q_D. Equation (2)
    computes the projection of the input into the latent space. Equation (3)
    computes the reconstruction of the input, while equation (4) computes the
    reconstruction error.

    .. math::

       \tilde{x} ~ q_D(\tilde{x}|x)                                     (1)

       y = s(W \tilde{x} + b)                                           (2)

       x = s(W' y  + b')                                                (3)

       L(x,z) = -sum_{k=1}^d [x_k \log z_k + (1-x_k) \log( 1-z_k)]      (4)
    
    this is a completely unsupervised model that concatenates
    autoencoders.
    """

    def __init__(self, ins, hs_lst, rng, theano_rng, wdtype, corruption_level):
        """initialize a stack of autoencoders


        @param ins: input dimension (int)
        @param hs_lst: hidden dimension (int)
        @param rng: random state (np.random.RandomState)
        @param theano_rng: random state (theano.tensor.shared_randomstreams.RandomStreams)
        @param wdtype: dtype of weights (np.dtype)
        @param corruption_level: keeps ``1-corruption_level`` entries of the
                         inputs the same and zero-out randomly selected subset
                         of size ``coruption_level``
        """

        X = T.matrix("X")
        layers = []
        for hs in hs_lst:
            ael = AELayer(X, ins, hs, rng, theano_rng,
                          wdtype, corruption_level)
            layers.append( ael )
            X = layers[-1].encoder
            ins = hs
        super(AEStack, self).__init__(layers)

    def update_params(self, train_batches, gradient_f, momentum, lrate, lidx):
        """step on the direction of the gradient

        step on the direction of the gradient for
        a whole epoch and update the model params in place. 
        By definition speed is initialized to 0.
        new_speed = -rho * dE/dw + mu * speed
        new_weight = w + new_speed

        @param train_batches: indexes of batches (list)
        @param gradient_f: function that returns the list of
                            gradients from the batch index.
        @param momentum: mu
        @param lrate: rho
        @return: none
        """

        for batch_i in train_batches:
            all_grads = gradient_f(batch_i)
            layer = self[lidx]
            l_grads = map(lambda i: all_grads.pop(0),
                    range(len(layer.get_params())))

            layer.speed_update(l_grads, momentum, lrate)
            layer.weight_update()



