
"""
-- Theano implementation of a single de-noising autoencoder, trained
on reconstruction error
-- Pieces of this code borrow from code snippets on deeplearning.net

NOTES:
-- Imports tt_algebra module, which performs all symbolic algebra
required to generate gradient updates
-- Function compilation to execute the SGD alegbra and train the network
 occurs in train_AE, which returns the learned parameters as well as the
mean loss for each epoch
-- Scoring function (score_AE) takes a layer parameter dictionary (as returned
by train_AE) as well as the layer options used and data to be scored
"""

import theano
from theano.tensor.shared_randomstreams import RandomStreams
import theano.tensor as T
import numpy as np
import sys

import weight_initializers as WI
import tt_algebra as ALG
from activation_ops import ActivationConstructors as AC
from data_utils import shared_dataset, dim_checker


################
### TRAINING ###
################


def train_AE(training_data, input_layer_ops):
    """

    :param input_layer_ops:
    :param training_data:
    :return:
    """

    try:
        # parse out options
        batch_size = input_layer_ops['batch_size']
        n_epochs = input_layer_ops['n_epoch']
        n_visible, n_hidden = input_layer_ops['layer_dim']
        weight_op = input_layer_ops['weight_init']

    except KeyError:

        print "Something in the global options class was mis-specified."
        print "Please check the AEOps class or your input dictionary."

    else:

        # set random number generators
        numpy_rng = np.random.RandomState(2 ** 30)
        theano_rng = RandomStreams(numpy_rng.randint(2 ** 30))

        # compute number of minibatches for training, validation and testing
        n_train_batches = dim_checker(training_data, 0) / batch_size

        # set training data as a theano shared variable
        training_shared = shared_dataset(training_data)

        # allocate symbolic variables for the data
        index = T.lscalar()  # index to a [mini]batch
        x = T.fmatrix('x')

        # initialize encoding weight matrix, decoding weight matrix
        initial_W = WI.WeightInits(n_visible, n_hidden, numpy_rng).weight_init(weight_op)
        W = theano.shared(value=initial_W, name='W', borrow=True)
        W_prime = W.T

        # initialize bias vectors
        b = theano.shared(value=np.zeros(n_hidden, dtype=theano.config.floatX), borrow=True)
        b_prime = theano.shared(value=np.zeros(n_visible, dtype=theano.config.floatX), borrow=True)

        # generate symbolic representations of SGD updates using mSGD
        cost, updates = ALG.AE_mSGD(x, [W, W_prime, b, b_prime], input_layer_ops, theano_rng)

        # compile a theano function that takes an index value as an argument
        # and updates the parameter values using the algebraic expressions
        # generated by mSGD
        train_dae = theano.function(
            [index],
            cost,
            updates=updates,
            givens={x: training_shared[index * batch_size: (index + 1) * batch_size]}
        )

        mean_loss = []
        # for each training epoch
        for epoch in xrange(n_epochs):
            l = []
            # compute gradient updates for each mini-batch
            for batch_index in xrange(n_train_batches):
                l.append(train_dae(batch_index))  # record loss for each batch
            # record the mean loss for the epoch
            epoch_loss = np.mean(l)
            mean_loss.append(epoch_loss)
            print "Loss for epoch " + str(epoch + 1) + ": " + str(epoch_loss)

        # once training procedure is complete, fetch weights
        out_params = {'W': W.get_value(),
                      'b': b.get_value()}

        out_values = {'parameters': out_params, 'loss': mean_loss}

        return out_values


def score_AE(unscored_x, input_layer_ops, layer_params):
    """

    :param unscored_x:
    :param input_layer_ops:
    :param layer_params:
    :return:
    """

    hidden_activation_op = input_layer_ops['activation']

    W_raw = layer_params['W']
    b_raw = layer_params['b']

    W = T.fmatrix('W')
    X = T.fmatrix('X')
    b = T.fvector('b')

    # compute symbolic hidden representation
    linear_hidden_expr = T.dot(X, W) + b
    nonlinear_hidden = AC(hidden_activation_op).activation(linear_hidden_expr)

    # compile function and score
    score_f = theano.function([X, W, b], nonlinear_hidden)
    scored = score_f(unscored_x, W_raw, b_raw)

    return scored
