#!/usr/bin/env python
# -*- coding: utf-8 -*-

from abc import ABCMeta, abstractmethod
import collections

class Dataset(object):
    """
    Abstract dataset class.
    """
    __metaclass__ = ABCMeta

    @abstractmethod
    def train_set(self):
        """
        :rtype: list of tuple
        """

    def valid_set(self):
        """
        :rtype: list of tuple
        """

    def test_set(self):
        """
        :rtype: list of tuple
        """

    def train_size(self):
        """
        Return size of training data. (optional)
        :rtype: number
        """
        train_set = self.train_set()
        if isinstance(train_set, collections.Iterable):
            return len(list(train_set))
        else:
            return None


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np

from controllers import TrainingController
from deepy.core.env import FLOATX
from deepy.core import graph

import logging as loggers
logging = loggers.getLogger(__name__)

class LearningRateAnnealer(TrainingController):
    """
    Learning rate annealer.
    """

    def __init__(self, patience=3, anneal_times=4):
        """
        :type trainer: deepy.trainers.base.NeuralTrainer
        """
        self._iter = 0
        self._annealed_iter = 0
        self._patience = patience
        self._anneal_times = anneal_times
        self._annealed_times = 0
        self._learning_rate = 0
        if type(self._learning_rate) == float:
            raise Exception("use shared_scalar to wrap the value in the config.")

    def bind(self, trainer):
        super(LearningRateAnnealer, self).bind(trainer)
        self._learning_rate = self._trainer.config.learning_rate
        self._iter = 0
        self._annealed_iter = 0

    def invoke(self):
        """
        Run it, return whether to end training.
        """
        self._iter += 1
        if self._iter - max(self._trainer.best_iter, self._annealed_iter) >= self._patience:
            if self._annealed_times >= self._anneal_times:
                logging.info("ending")
                self._trainer.exit()
            else:
                self._trainer.set_params(*self._trainer.best_params)
                self._learning_rate.set_value(self._learning_rate.get_value() * 0.5)
                self._annealed_times += 1
                self._annealed_iter = self._iter
                logging.info("annealed learning rate to %f" % self._learning_rate.get_value())

    @staticmethod
    def learning_rate(value=0.01):
        """
        Wrap learning rate.
        """
        return graph.shared(value, name="learning_rate")

class ScheduledLearningRateAnnealer(TrainingController):
    """
    Anneal learning rate according to pre-scripted schedule.
    """

    def __init__(self, start_halving_at=5, end_at=10, halving_interval=1, rollback=False):
        logging.info("iteration to start halving learning rate: %d" % start_halving_at)
        self.epoch_start_halving = start_halving_at
        self.end_at = end_at
        self._halving_interval = halving_interval
        self._rollback = rollback
        self._last_halving_epoch = 0
        self._learning_rate = None

    def bind(self, trainer):
        super(ScheduledLearningRateAnnealer, self).bind(trainer)
        self._learning_rate = self._trainer.config.learning_rate
        self._last_halving_epoch = 0

    def invoke(self):
        epoch = self._trainer.epoch()
        if epoch >= self.epoch_start_halving and epoch >= self._last_halving_epoch + self._halving_interval:
            if self._rollback:
                self._trainer.set_params(*self._trainer.best_params)
            self._learning_rate.set_value(self._learning_rate.get_value() * 0.5)
            logging.info("halving learning rate to %f" % self._learning_rate.get_value())
            self._trainer.network.train_logger.record("set learning rate to %f" % self._learning_rate.get_value())
            self._last_halving_epoch = epoch
        if epoch >= self.end_at:
            logging.info("ending")
            self._trainer.exit()

class ExponentialLearningRateAnnealer(TrainingController):
    """
    Exponentially decay learning rate after each update.
    """

    def __init__(self, decay_factor=1.000004, min_lr=.000001, debug=False):
        logging.info("exponentially decay learning rate with decay factor = %f" % decay_factor)
        self.decay_factor = np.array(decay_factor, dtype=FLOATX)
        self.min_lr = np.array(min_lr, dtype=FLOATX)
        self.debug = debug
        self._learning_rate = self._trainer.config.learning_rate
        if type(self._learning_rate) == float:
            raise Exception("use shared_scalar to wrap the value in the config.")
        self._trainer.network.training_callbacks.append(self.update_callback)

    def update_callback(self):
        if self._learning_rate.get_value() > self.min_lr:
            self._learning_rate.set_value(self._learning_rate.get_value() / self.decay_factor)

    def invoke(self):
        if self.debug:
            logging.info("learning rate: %.8f" % self._learning_rate.get_value())

class SimpleScheduler(TrainingController):

    """
    Simple scheduler with maximum patience.
    """

    def __init__(self, end_at=10):
        """
        :type trainer: deepy.trainers.base.NeuralTrainer
        """
        self._iter = 0
        self._patience = end_at

    def invoke(self):
        """
        Run it, return whether to end training.
        """
        self._iter += 1
        logging.info("{} epochs left to run".format(self._patience - self._iter))
        if self._iter >= self._patience:
            self._trainer.exit()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import theano.tensor as T
from deepy import NeuralLayer, AutoEncoder, Dense
from deepy import GaussianInitializer, global_theano_rand

class ReparameterizationLayer(NeuralLayer):
    """
    Reparameterization layer in a Variational encoder.
    Only binary output cost function is supported now.
    The prior value is recorded after the computation graph created.
    """

    def __init__(self, size, sample=True):
        """
        :param size: the size of latent variable
        :param sample: whether to get a clean latent variable
        """
        super(ReparameterizationLayer, self).__init__("VariationalEncoder")
        self.size = size
        self.output_dim = size
        self.sample = sample
        self._prior = None

    def prepare(self):
        self._mu_encoder = Dense(self.size, 'linear', init=GaussianInitializer(), random_bias=True).init(
            self.input_dim)
        self._log_sigma_encoder = Dense(self.size, 'linear', init=GaussianInitializer(), random_bias=True).init(
            self.input_dim)
        self.register_inner_layers(self._mu_encoder, self._log_sigma_encoder)

    def compute_tensor(self, x):
        # Compute p(z|x)
        mu = self._mu_encoder.compute_tensor(x)
        log_sigma = 0.5 * self._log_sigma_encoder.compute_tensor(x)
        self._prior = 0.5* T.sum(1 + 2*log_sigma - mu**2 - T.exp(2*log_sigma))
        # Reparameterization
        eps = global_theano_rand.normal((x.shape[0], self.size))

        if not self.sample:
            z = mu
        else:
            z = mu + T.exp(log_sigma) * eps
        return z

    def prior(self):
        """
        Get the prior value.
        """
        return self._prior

class VariationalAutoEncoder(AutoEncoder):
    """
    Variational Auto Encoder.
    Only binary output cost function is supported now.
    """

    def __init__(self, input_dim, rep_dim, input_tensor=None, sample=True):
        """
        """
        super(VariationalAutoEncoder, self).__init__(input_dim, rep_dim)
        self.sample = sample
        self._setup_monitors = True

    def stack_reparameterization_layer(self, layer_size):
        """
        Perform reparameterization trick for latent variables.
        :param layer_size: the size of latent variable
        """
        self.rep_layer = ReparameterizationLayer(layer_size, sample=self.sample)
        self.stack_encoders(self.rep_layer)

    def _cost_func(self, y):
        logpxz  = - T.nnet.binary_crossentropy(y, self.input_variables[0]).sum()
        logp = logpxz + self.rep_layer.prior()
        # the lower bound is the mean value of logp
        cost = - logp
        if self._setup_monitors:
            self._setup_monitors = False
            self.training_monitors.append(("lower_bound", logp / y.shape[0]))
            self.testing_monitors.append(("lower_bound", logp / y.shape[0]))
        return cost


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import deepy.tensor as DT

from network import NeuralNetwork

class AutoEncoder(NeuralNetwork):
    """
    A class for defining simple auto encoders.
    Must call stack_encoding before stack_decoding.
    """
    def __init__(self, input_dim, rep_dim, input_tensor=None):
        """
        Create an auto-encoder.
        :param input_dim: dimension of the input layer
        :param rep_dim: dimension of the representation layer
        :param input_tensor:
        :return:
        """
        super(AutoEncoder, self).__init__(input_dim, input_tensor=input_tensor)

        self.rep_dim = rep_dim
        self.encoding_layes = []
        self.decoding_layers = []
        self.encoding_network = None
        self.decoding_network = None

    def _cost_func(self, y):
        return DT.costs.least_squares(y, self.input_variables[0])

    @property
    def cost(self):
        return self._cost_func(self.output)

    @property
    def test_cost(self):
        return self._cost_func(self.test_output)

    def stack_encoders(self, *layers):
        """
        Stack encoding layers, this must be done before stacking decoding layers.
        """
        self.stack(*layers)
        self.encoding_layes.extend(layers)

    def stack_decoders(self, *layers):
        """
        Stack decoding layers.
        """
        self.stack(*layers)
        self.decoding_layers.extend(layers)

    def encode(self, x):
        """
        Encode given input.
        """
        if not self.encoding_network:
            self.encoding_network = NeuralNetwork(self.input_dim, self.input_tensor)
            self.encoding_network.input_variables = self.input_variables
            for layer in self.encoding_layes:
                self.encoding_network.stack_layer(layer, no_setup=True)
        return self.encoding_network.compute(*x)

    def decode(self, x):
        """
        Decode given representation.
        """
        if not self.rep_dim:
            raise Exception("rep_dim must be set to decode.")
        if not self.decoding_network:
            self.decoding_network = NeuralNetwork(self.rep_dim)
            for layer in self.decoding_layers:
                self.decoding_network.stack_layer(layer, no_setup=True)
        return self.decoding_network.compute(x)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
This code is copied from https://github.com/vsvinayak/mnist-helper.
It requires Scipy to perform convolve2d.
Default parameters are modified.
"""

import numpy as np
import math
from scipy.signal import convolve2d

from deepy.core import env

def create_2d_gaussian(dim, sigma):
    """
    This function creates a 2d gaussian kernel with the standard deviation
    denoted by sigma

    :param dim: integer denoting a side (1-d) of gaussian kernel
    :param sigma: floating point indicating the standard deviation

    :returns: a numpy 2d array
    """

    # check if the dimension is odd
    if dim % 2 == 0:
        raise ValueError("Kernel dimension should be odd")

    # initialize the kernel
    kernel = np.zeros((dim, dim), dtype=np.float16)

    # calculate the center point
    center = dim/2

    # calculate the variance
    variance = sigma ** 2

    # calculate the normalization coefficeint
    coeff = 1. / (2 * variance)

    # create the kernel
    for x in range(0, dim):
        for y in range(0, dim):
            x_val = abs(x - center)
            y_val = abs(y - center)
            numerator = x_val**2 + y_val**2
            denom = 2*variance

            kernel[x,y] = coeff * np.exp(-1. * numerator/denom)

    return kernel/sum(sum(kernel))

def elastic_distortion(image, kernel_dim=21, sigma=6, alpha=30, negated=True):
    """
    This method performs elastic transformations on an image by convolving
    with a gaussian kernel.
    :param image: a numpy nd array
    :kernel_dim: dimension(1-D) of the gaussian kernel
    :param sigma: standard deviation of the kernel
    :param alpha: a multiplicative factor for image after convolution
    :param negated: a flag indicating whether the image is negated or not
    :returns: a nd array transformed image
    """
    # check if the image is a negated one
    if not negated:
        image = 255-image

    # check if kernel dimesnion is odd
    if kernel_dim % 2 == 0:
        raise ValueError("Kernel dimension should be odd")

    # create an empty image
    result = np.zeros(image.shape)

    # create random displacement fields
    displacement_field_x = np.array([[env.numpy_rand.random_integers(-1, 1) for x in xrange(image.shape[0])] \
                            for y in xrange(image.shape[1])]) * alpha
    displacement_field_y = np.array([[env.numpy_rand.random_integers(-1, 1) for x in xrange(image.shape[0])] \
                            for y in xrange(image.shape[1])]) * alpha

    # create the gaussian kernel
    kernel = create_2d_gaussian(kernel_dim, sigma)

    # convolve the fields with the gaussian kernel
    displacement_field_x = convolve2d(displacement_field_x, kernel)
    displacement_field_y = convolve2d(displacement_field_y, kernel)

    # make the distortrd image by averaging each pixel value to the neighbouring
    # four pixels based on displacement fields

    for row in xrange(image.shape[1]):
        for col in xrange(image.shape[0]):
            low_ii = row + int(math.floor(displacement_field_x[row, col]))
            high_ii = row + int(math.ceil(displacement_field_x[row, col]))

            low_jj = col + int(math.floor(displacement_field_y[row, col]))
            high_jj = col + int(math.ceil(displacement_field_y[row, col]))

            if low_ii < 0 or low_jj < 0 or high_ii >= image.shape[1] -1 \
               or high_jj >= image.shape[0] - 1:
                continue

            res = image[low_ii, low_jj]/4 + image[low_ii, high_jj]/4 + \
                    image[high_ii, low_jj]/4 + image[high_ii, high_jj]/4

            result[row, col] = res

    return result


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import cPickle as pickle
import gzip
import logging as loggers
import os
from threading import Thread

import filelock
import numpy as np
import theano
import theano.tensor as T

import deepy
from deepy.utils.map_dict import MapDict
from deepy.layers.block import Block
from deepy.layers.layer import NeuralLayer
from deepy.utils import dim_to_var
from deepy.trainers.train_logger import TrainLogger

logging = loggers.getLogger("network")
save_logger = loggers.getLogger("saving")

DEEPY_MESSAGE = "deepy version = %s" % deepy.__version__

def save_network_params(params, path):
    lock = filelock.FileLock(path)
    with lock:
        if path.endswith('gz'):
            opener = gzip.open if path.lower().endswith('.gz') else open
            handle = opener(path, 'wb')
            pickle.dump(params, handle)
            handle.close()
        elif path.endswith('uncompressed.npz'):
            np.savez(path, *params)
        elif path.endswith('.npz'):
            np.savez_compressed(path, *params)
        else:
            raise Exception("File format of %s is not supported, use '.gz' or '.npz' or '.uncompressed.gz'" % path)

class NeuralNetwork(object):
    """
    The base class of neural networks.
    """

    def __init__(self, input_dim, input_tensor=None):
        logging.info(DEEPY_MESSAGE)
        self.input_dim = input_dim
        self.input_tensor = input_tensor
        self.parameter_count = 0

        self.parameters = []
        self.free_parameters = []

        self.training_updates = []
        self.updates = []

        self.input_variables = []
        self.target_variables = []

        self.training_callbacks = []
        self.testing_callbacks = []
        self.epoch_callbacks = []

        self.layers = []
        self._test_outputs = []
        self._test_output = None
        self._output_keys = None

        self._hidden_outputs = []
        self.training_monitors = []
        self.testing_monitors = []

        self.setup_variables()
        self.train_logger = TrainLogger()

    def stack_layer(self, layer, no_setup=False):
        """
        Stack a neural layer.
        :type layer: NeuralLayer
        :param no_setup: whether the layer is already initialized
        """
        if layer.name:
            layer.name += "%d" % (len(self.layers) + 1)
        if not self.layers:
            layer.init(self.input_dim, no_prepare=no_setup)
        else:
            layer.init(self.layers[-1].output_dim, no_prepare=no_setup)
        self._output = layer.compute_tensor(self._output)
        self._test_output = layer.compute_tensor(self._test_output)
        self._hidden_outputs.append(self._output)
        self.register_layer(layer)
        self.layers.append(layer)

    def register(self, *layers):
        """
        Register multiple layers as the components of the network.
        The parameter of those layers will be trained.
        But the output of the layer will not be stacked.
        """
        for layer in layers:
            self.register_layer(layer)

    def register_layer(self, layer):
        """
        Register the layer so that it's param will be trained.
        But the output of the layer will not be stacked.
        """
        if type(layer) == Block:
            layer.fix()
        self.parameter_count += layer.parameter_count
        self.parameters.extend(layer.parameters)
        self.free_parameters.extend(layer.free_parameters)
        self.training_monitors.extend(layer.training_monitors)
        self.testing_monitors.extend(layer.testing_monitors)
        self.updates.extend(layer.updates)
        self.training_updates.extend(layer.training_updates)
        self.input_variables.extend(layer.external_inputs)
        self.target_variables.extend(layer.external_targets)

        self.training_callbacks.extend(layer.training_callbacks)
        self.testing_callbacks.extend(layer.testing_callbacks)
        self.epoch_callbacks.extend(layer.epoch_callbacks)

    def first_layer(self):
        """
        Return first layer.
        """
        return self.layers[0] if self.layers else None

    def stack(self, *layers):
        """
        Stack layers.
        """
        for layer in layers:
            self.stack_layer(layer)
        return self

    def prepare_training(self):
        """
        This function will be called before training.
        """
        self.report()

    def monitor_layer_outputs(self):
        """
        Monitoring the outputs of each layer.
        Useful for troubleshooting convergence problems.
        """
        for layer, hidden in zip(self.layers, self._hidden_outputs):
            self.training_monitors.append(('mean(%s)' % (layer.name), abs(hidden).mean()))

    @property
    def all_parameters(self):
        """
        Return all parameters.
        """
        params = []
        params.extend(self.parameters)
        params.extend(self.free_parameters)

        return params

    def setup_variables(self):
        """
        Set up variables.
        """
        if self.input_tensor:
            if type(self.input_tensor) == int:
                x = dim_to_var(self.input_tensor, name="x")
            else:
                x = self.input_tensor
        else:
            x = T.matrix('x')
        self.input_variables.append(x)
        self._output = x
        self._test_output = x

    def _compile(self):
        if not hasattr(self, '_compute'):
            if isinstance(self._test_outputs, dict):
                self._output_keys, out_tensors = map(list, zip(*self._test_outputs.items()))
                if self._test_output:
                    self._output_keys.insert(0, "cost")
                    out_tensors.insert(0, self._test_output)
            elif isinstance(self._test_outputs, list) and self._test_outputs:
                out_tensors = self._test_outputs
                if self._test_output:
                    out_tensors.insert(0, self._test_output)
            else:
                out_tensors = self._test_output
            self._compute = theano.function(
                filter(lambda x: x not in self.target_variables, self.input_variables),
                out_tensors, updates=self.updates, allow_input_downcast=True)

    def compute(self, *x):
        """
        Return network output.
        """
        self._compile()
        outs = self._compute(*x)
        if self._output_keys:
            return MapDict(dict(zip(self._output_keys, outs)))
        else:
            return outs

    @property
    def output(self):
        """
        Return output variable.
        """
        return self._output

    @property
    def test_output(self):
        """
        Return output variable in test time.
        """
        return self._test_output

    @property
    def cost(self):
        """
        Return cost variable.
        """
        return T.constant(0)

    @property
    def test_cost(self):
        """
        Return cost variable in test time.
        """
        return self.cost

    def save_params(self, path, new_thread=False):
        """
        Save parameters to file.
        """
        save_logger.info(path)
        param_variables = self.all_parameters
        params = [p.get_value().copy() for p in param_variables]
        if new_thread:
            thread = Thread(target=save_network_params, args=(params, path))
            thread.start()
        else:
            save_network_params(params, path)
        self.train_logger.save(path)

    def load_params(self, path, exclude_free_params=False):
        """
        Load parameters from file.
        """
        if not os.path.exists(path): return;
        logging.info("loading parameters from %s" % path)
        # Decide which parameters to load
        if exclude_free_params:
            params_to_load = self.parameters
        else:
            params_to_load = self.all_parameters
        # Load parameters
        if path.endswith(".gz"):
            opener = gzip.open if path.lower().endswith('.gz') else open
            handle = opener(path, 'rb')
            saved_params = pickle.load(handle)
            handle.close()
            # Write parameters
            for target, source in zip(params_to_load, saved_params):
                logging.info('%s: setting value %s', target.name, source.shape)
                target.set_value(source)
        elif path.endswith(".npz"):
            arrs = np.load(path)
            # Write parameters
            for target, idx in zip(params_to_load, range(len(arrs.keys()))):
                source = arrs['arr_%d' % idx]
                logging.info('%s: setting value %s', target.name, source.shape)
                target.set_value(source)
        else:
            raise Exception("File format of %s is not supported, use '.gz' or '.npz' or '.uncompressed.gz'" % path)

        self.train_logger.load(path)

    def report(self):
        """
        Print network statistics.
        """
        logging.info("network inputs: %s", " ".join(map(str, self.input_variables)))
        logging.info("network targets: %s", " ".join(map(str, self.target_variables)))
        logging.info("network parameters: %s", " ".join(map(str, self.all_parameters)))
        logging.info("parameter count: %d", self.parameter_count)

    def epoch_callback(self):
        """
        Callback for each epoch.
        """
        for cb in self.epoch_callbacks:
            cb()

    def training_callback(self):
        """
        Callback for each training iteration.
        """
        for cb in self.training_callbacks:
            cb()

    def testing_callback(self):
        """
        Callback for each testing iteration.
        """
        for cb in self.training_callbacks:
            cb()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging as loggers

import numpy as np
import theano

from deepy.utils import UniformInitializer
from deepy.core.env import env
from deepy.core.tensor_conversion import neural_computation_prefer_tensor, convert_to_theano_var

logging = loggers.getLogger("deepy")

class NeuralLayer(object):

    def __init__(self, name=None):
        """
        Create a neural layer.
        """
        self.name = name if name else self.__class__.__name__
        self.input_dim = 0
        self.input_dims = [0]
        self.output_dim = 0
        self.output_dims= [0]

        self._linked_block = None

        self.initialized = False
        self.updates = []
        self.training_updates = []
        self.free_parameters = []
        self.parameters = []
        self.training_monitors = []
        self.testing_monitors = []
        self._registered_monitors = set()
        self._registered_updates = set()
        self._registered_training_updates = set()
        self.external_inputs = []
        self.external_targets = []
        self.parameter_count = 0
        self.epoch_callbacks = []
        self.training_callbacks = []
        self.testing_callbacks = []

    def init(self, input_dim=0, input_dims=None, no_prepare=False):
        """
        Initialize the layer.
        :param no_prepare: avoid calling preparation function
        """
        if self.initialized:
            return
        # configure input dimensions
        if input_dims:
            self.input_dims = input_dims
            self.input_dim = input_dims[0]
        else:
            self.input_dim = input_dim
            self.input_dims = [input_dims]
        # set default output dimension
        if self.output_dim == 0:
            self.output_dim = self.input_dim
        self.initialized = True
        # call prepare
        if not no_prepare:
            self.prepare()
        return self

    def compute(self, *inputs, **kwargs):
        """
        Compute based on NeuralVariable.
        :type inputs:  list of NeuralVariable
        :return: NeuralVariable
        """
        from deepy.core.neural_var import NeuralVariable
        from deepy.core.graph import graph
        if type(inputs[0]) != NeuralVariable:
            raise SystemError("The input of `compute` must be NeuralVar")

        dims = [t.dim() for t in inputs]
        if len(inputs) == 1:
            self.init(input_dim=dims[0])
        else:
            self.init(input_dims=dims)
        # Check block
        if self.parameters and not self._linked_block:
            self.belongs_to(graph.default_block())
        # convert kwargs
        train_kwargs, _, _ = convert_to_theano_var(kwargs)

        output = self.compute_tensor(*[t.tensor for t in inputs], **train_kwargs)

        if type(output) != list and type(output) != tuple:
            return NeuralVariable(output, dim=self.output_dim)
        else:
            return [NeuralVariable(*item) for item in zip(output, self.output_dims)]

    def prepare(self):
        """
        Prepare function will be called after connected.
        """

    @neural_computation_prefer_tensor
    def compute_tensor(self, *args, **kwargs):
        """
        Compute with tensors in Theano.
        """
        raise NotImplementedError("output function of '%s' is not implemented" % self.name)

    def belongs_to(self, block):
        """
        Let the given block or network manage the parameters of this layer.
        :param block: Block or NeuralNetwork
        :return: NeuralLayer
        """
        if self._linked_block:
            raise SystemError("The layer {} has already blonged to {}".format(self.name, self._linked_block.name))
        self._linked_block = block
        block.register_layer(self)
        return self

    def register(self, *layers):
        """
        Register inner layers.
        """
        self.register_inner_layers(*layers)

    def register_inner_layers(self, *layers):
        for layer in layers:
            self.register_parameters(*layer.parameters)
            self.register_updates(*layer.updates)
            self.register_training_updates(*layer.training_updates)
            self.training_monitors.extend(layer.training_monitors)
            self.testing_monitors.extend(layer.testing_monitors)

    def register_parameters(self, *parameters):
        """
        Register parameters.
        """
        for param in parameters:
            self.parameter_count += np.prod(param.get_value().shape)
        self.parameters.extend(parameters)

    def register_free_parameters(self, *free_parameters):
        """
        Register free parameters, which means their value will not be learned by trainer.
        """
        return self.free_parameters.extend(free_parameters)

    def register_updates(self, *updates):
        """
        Register updates that will be executed in each iteration.
        """
        for key, node in updates:
            if key not in self._registered_updates:
                self.updates.append((key, node))
                self._registered_updates.add(key)

    def register_training_updates(self, *updates):
        """
        Register updates that will only be executed in training phase.
        """
        for key, node in updates:
            if key not in self._registered_training_updates:
                self.training_updates.append((key, node))
                self._registered_training_updates.add(key)

    def register_monitors(self, *monitors):
        """
        Register monitors they should be tuple of name and Theano variable.
        """
        for key, node in monitors:
            if key not in self._registered_monitors:
                node *= 1.0 # Avoid CudaNdarray
                self.training_monitors.append((key, node))
                self.testing_monitors.append((key, node))
                self._registered_monitors.add(key)

    def register_external_inputs(self, *variables):
        """
        Register external input variables.
        """
        self.external_inputs.extend(variables)

    def register_external_targets(self, *variables):
        """
        Register extenal target variables.
        """
        self.external_targets.extend(variables)

    def register_training_callbacks(self, *callbacks):
        """
        Register callback for each iteration in the training.
        """
        self.training_callbacks.extend(callbacks)

    def register_testing_callbacks(self, *callbacks):
        """
        Register callback for each iteration in the testing.
        """
        self.testing_callbacks.extend(callbacks)

    def register_epoch_callbacks(self, *callbacks):
        """
        Register callback which will be called after epoch finished.
        """
        self.epoch_callbacks.extend(callbacks)

    def create_weight(self, input_n=1, output_n=1, label="W", initializer=None, shape=None):
        if not shape:
            shape = (input_n, output_n)

        if not initializer:
            initializer = env.default_initializer

        weight = theano.shared(initializer.sample(shape).astype(env.FLOATX), name='{}_{}'.format(self.name, label))

        logging.info('create param %s %s for %s', label, str(shape), self.name)
        return weight

    def create_bias(self, output_n=1, label="B", value=0., shape=None):
        if not shape:
            shape = (output_n, )
        bs =  np.ones(shape)
        bs *= value
        bias = theano.shared(bs.astype(env.FLOATX), name='{}_{}'.format(self.name, label))
        logging.info('create param %s %s for %s', label, str(shape), self.name)
        return bias

    def create_scalar(self, name="S", value=0, dtype=env.FLOATX):
        bs = np.array(0, dtype=dtype)
        bs += value
        v = theano.shared(bs, name='{}_{}'.format(self.name, name))

        logging.info('create scalar %s', name)
        return v

    def create_vector(self, n, name="V", dtype=env.FLOATX):
        bs =  np.zeros(n, dtype=dtype)
        v = theano.shared(bs, name='{}_{}'.format(self.name, name))

        logging.info('create vector %s: %d', name, n)
        return v

    def create_matrix(self, m, n, name="M"):

        matrix = theano.shared(np.zeros((m, n)).astype(env.FLOATX), name="{}_{}".format(self.name, name))

        logging.info('create matrix %s: %d x %d', name, m, n)
        return matrix

    def activation(self, name):
        from deepy.tensor.activations import get_activation
        return get_activation(name)

    def callback_forward_propagation(self):
        pass

    def set_name(self, name):
        """
        Set the name of this layer.
        This will be the key of saved parameters.
        """
        self.name = name


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import theano.tensor as T
import inspect

def wrap_core(func, train_config, *args):
    spec = inspect.getargspec(func)
    params = spec.args[-len(spec.defaults):]
    default_values = spec.defaults
    kwargs = dict(zip(params, default_values))
    if train_config:
        for param, default in zip(params, default_values):
            config_val = train_config.get(param, default)
            kwargs[param] = config_val
    return list(func(*args, **kwargs))

def multiple_l2_norm(tensors):
    """
    Get the L2 norm of multiple tensors.
    This function is taken from blocks.
    """
    # Another way for doing this, I don't know which one is fast
    # return T.sqrt(sum(T.sum(t ** 2) for t in tensors))
    flattened = [T.as_tensor_variable(t).flatten() for t in tensors]
    flattened = [(t if t.ndim > 0 else t.dimshuffle('x'))
                 for t in flattened]
    joined = T.join(0, *flattened)
    return T.sqrt(T.sqr(joined).sum())


#!/usr/bin/env python
# -*- coding: utf-8 -*-

try:
  from cPickle import dumps, loads
except ImportError:
  from pickle import dumps, loads

class StreamPickler(object):
    """
    Pickle massive objects without using too much memory.
    The code of this class is modified upon https://code.google.com/p/streaming-pickle/.
    """

    @staticmethod
    def dump(iterable_to_pickle, file_obj):
        """
        dump contents of an iterable iterable_to_pickle to file_obj, a file
        opened in write mode
        """
        for elt in iterable_to_pickle:
            StreamPickler.dump_one(elt, file_obj)

    @staticmethod
    def dump_one(elt_to_pickle, file_obj):
        """
        dumps one element to file_obj, a file opened in write mode
        """
        pickled_elt_str = dumps(elt_to_pickle)
        file_obj.write(pickled_elt_str)
        # record separator is a blank line
        # (since pickled_elt_str might contain its own newlines)
        file_obj.write('\n\n')

    @staticmethod
    def load(file_obj):
        """
        load contents from file_obj, returning a generator that yields one
        element at a time
        """
        cur_elt = []
        for line in file_obj:
          cur_elt.append(line)

          if line == '\n':
            pickled_elt_str = ''.join(cur_elt)
            cur_elt = []
            try:
                elt = loads(pickled_elt_str)
            except ValueError:
                continue

            yield elt


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from deepy.layers.layer import NeuralLayer

class Block(NeuralLayer):
    """
    Create a block, which contains the parameters of many connected layers.
    """

    _BLOCK_COUNT = 0

    def __init__(self, *layers, **kwargs):
        """
        Create a new parameter block with some layers.
        You can also specify a name through kwargs.
        """
        name = kwargs['name'] if 'name' in kwargs else "block_{}".format(self._BLOCK_COUNT + 1)
        super(Block, self).__init__(name)
        self._BLOCK_COUNT += 1
        self.layers = list(layers)
        self.fixed = False

    def fix(self):
        """
        Fix the block, register all the parameters of sub layers.
        :return:
        """
        if not self.fixed:
            for layer in self.layers:
                if not layer.initialized:
                    raise Exception("All sub layers in a block must be initialized when fixing it.")
                self.register_inner_layers(layer)
            self.fixed = True

    def add_parameters(self, *parameters):
        from deepy.core.neural_var import NeuralVariable
        for param in parameters:
            if isinstance(param, NeuralVariable):
                param = param.tensor
            self.parameters.append(param)

    def register(self, *layers):
        """
        Register many connected layers.
        :type layers: list of NeuralLayer
        """
        for layer in layers:
            self.register_layer(layer)

    def register_layer(self, layer):
        """
        Register one connected layer.
        :type layer: NeuralLayer
        """
        if self.fixed:
            raise Exception("After a block is fixed, no more layers can be registered.")
        self.layers.append(layer)

    def compute_tensor(self, x):
        return x

    def load_params(self, path, exclude_free_params=False):
        from deepy.core import graph
        """
        Load parameters to the block.
        """
        from deepy.core.comp_graph import ComputationalGraph
        model = graph.compile(blocks=[self])
        model.load_params(path, exclude_free_params=exclude_free_params)

    @property
    def all_parameters(self):
        return self.parameters


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from abc import ABCMeta, abstractmethod

import numpy as np
import theano.tensor as T

from deepy.core.tensor_conversion import neural_computation
from deepy.utils import XavierGlorotInitializer, OrthogonalInitializer, Scanner
from deepy.core import env
import deepy.tensor as DT
from . import NeuralLayer

OUTPUT_TYPES = ["sequence", "one"]
INPUT_TYPES = ["sequence", "one"]

class RecurrentLayer(NeuralLayer):
    __metaclass__ = ABCMeta

    def __init__(self, name, state_names, hidden_size=100, input_type="sequence", output_type="sequence",
                 inner_init=None, outer_init=None,
                 gate_activation='sigmoid', activation='tanh',
                 steps=None, backward=False, mask=None,
                 additional_input_dims=None):
        from deepy.core.neural_var import NeuralVariable
        from deepy.tensor.activations import get_activation
        super(RecurrentLayer, self).__init__(name)
        self.state_names = state_names
        self.main_state = state_names[0]
        self.hidden_size = hidden_size
        self._gate_activation = gate_activation
        self._activation = activation
        self.gate_activate = get_activation(self._gate_activation)
        self.activate = get_activation(self._activation)
        self._input_type = input_type
        self._output_type = output_type
        self.inner_init = inner_init if inner_init else OrthogonalInitializer()
        self.outer_init = outer_init if outer_init else XavierGlorotInitializer()
        self._steps = steps
        self._mask = mask.tensor if type(mask) == NeuralVariable else mask
        self._go_backwards = backward
        self.additional_input_dims = additional_input_dims if additional_input_dims else []

        if input_type not in INPUT_TYPES:
            raise Exception("Input type of {} is wrong: {}".format(name, input_type))
        if output_type not in OUTPUT_TYPES:
            raise Exception("Output type of {} is wrong: {}".format(name, output_type))

    @neural_computation
    def step(self, step_inputs):
        new_states = self.compute_new_state(step_inputs)

        # apply mask for each step if `output_type` is 'one'
        if step_inputs.get("mask"):
            mask = step_inputs["mask"].dimshuffle(0, 'x')
            for state_name in new_states:
                new_states[state_name] = new_states[state_name] * mask + step_inputs[state_name] * (1 - mask)

        return new_states

    @abstractmethod
    def compute_new_state(self, step_inputs):
        """
        :type step_inputs: dict
        :rtype: dict
        """

    @abstractmethod
    def merge_inputs(self, input_var, additional_inputs=None):
        """
        Merge inputs and return a map, which will be passed to core_step.
        :type input_var: T.var
        :param additional_inputs: list
        :rtype: dict
        """

    @abstractmethod
    def prepare(self):
        pass

    @neural_computation
    def compute_step(self, state, lstm_cell=None, input=None, additional_inputs=None):
        """
        Compute one step in the RNN.
        :return: one variable for RNN and GRU, multiple variables for LSTM
        """
        if not self.initialized:
            input_dim = None
            if input and hasattr(input.tag, 'last_dim'):
                input_dim = input.tag.last_dim
            self.init(input_dim)

        input_map = self.merge_inputs(input, additional_inputs=additional_inputs)
        input_map.update({"state": state, "lstm_cell": lstm_cell})
        output_map = self.compute_new_state(input_map)
        outputs = [output_map.pop("state")]
        outputs += output_map.values()
        for tensor in outputs:
            tensor.tag.last_dim = self.hidden_size
        if len(outputs) == 1:
            return outputs[0]
        else:
            return outputs

    @neural_computation
    def get_initial_states(self, input_var, init_state=None):
        """
        :type input_var: T.var
        :rtype: dict
        """
        initial_states = {}
        for state in self.state_names:
            if state != "state" or not init_state:
                if self._input_type == 'sequence' and input_var.ndim == 2:
                    init_state = T.alloc(np.cast[env.FLOATX](0.), self.hidden_size)
                else:
                    init_state = T.alloc(np.cast[env.FLOATX](0.), input_var.shape[0], self.hidden_size)
            initial_states[state] = init_state
        return initial_states

    @neural_computation
    def get_step_inputs(self, input_var, states=None, mask=None, additional_inputs=None):
        """
        :type input_var: T.var
        :rtype: dict
        """
        step_inputs = {}
        if self._input_type == "sequence":
            if not additional_inputs:
                additional_inputs = []
            if mask:
                step_inputs['mask'] = mask.dimshuffle(1, 0)
            step_inputs.update(self.merge_inputs(input_var, additional_inputs=additional_inputs))
        else:
            # step_inputs["mask"] = mask.dimshuffle((1,0)) if mask else None
            if additional_inputs:
                step_inputs.update(self.merge_inputs(None, additional_inputs=additional_inputs))
        if states:
            for name in self.state_names:
                step_inputs[name] = states[name]

        return step_inputs

    def compute(self, input_var, mask=None, additional_inputs=None, steps=None, backward=False, init_states=None, return_all_states=False):
        from deepy.core.neural_var import NeuralVariable
        if additional_inputs and not self.additional_input_dims:
            self.additional_input_dims = map(lambda var: var.dim(), additional_inputs)
        result_var = super(RecurrentLayer, self).compute(input_var,
                                                   mask=mask, additional_inputs=additional_inputs, steps=steps, backward=backward, init_states=init_states, return_all_states=return_all_states)
        if return_all_states:
            state_map = {}
            for k in result_var.tensor:
                state_map[k] = NeuralVariable(result_var.tensor[k], self.output_dim)
            return state_map
        else:
            return result_var

    def compute_tensor(self, input_var, mask=None, additional_inputs=None, steps=None, backward=False, init_states=None, return_all_states=False):
        # prepare parameters
        backward = backward if backward else self._go_backwards
        steps = steps if steps else self._steps
        mask = mask if mask else self._mask
        if mask and self._input_type == "one":
            raise Exception("Mask only works with sequence input")
        # get initial states
        init_state_map = self.get_initial_states(input_var)
        if init_states:
            for name, val in init_states.items():
                if name in init_state_map:
                    init_state_map[name] = val
        # get input sequence map
        if self._input_type == "sequence":
            # Move middle dimension to left-most position
            # (sequence, batch, value)
            if input_var.ndim == 3:
                input_var = input_var.dimshuffle((1,0,2))

            seq_map = self.get_step_inputs(input_var, mask=mask, additional_inputs=additional_inputs)
        else:
            init_state_map[self.main_state] = input_var
            seq_map = self.get_step_inputs(None, mask=mask, additional_inputs=additional_inputs)
        # scan
        retval_map, _ = Scanner(
            self.step,
            sequences=seq_map,
            outputs_info=init_state_map,
            n_steps=steps,
            go_backwards=backward
        ).compute()
        # return main states
        main_states = retval_map[self.main_state]
        if self._output_type == "one":
            if return_all_states:
                return_map = {}
                for name, val in retval_map.items():
                    return_map[name] = val[-1]
                return return_map
            else:
                return main_states[-1]
        elif self._output_type == "sequence":
            if return_all_states:
                return_map = {}
                for name, val in retval_map.items():
                    return_map[name] = val.dimshuffle((1,0,2))
                return return_map
            else:
                main_states = main_states.dimshuffle((1,0,2)) # ~ batch, time, size
                # if mask: # ~ batch, time
                #     main_states *= mask.dimshuffle((0, 1, 'x'))
                return main_states

class RNN(RecurrentLayer):

    def  __init__(self, hidden_size, **kwargs):
        kwargs["hidden_size"] = hidden_size
        super(RNN, self).__init__("RNN", ["state"], **kwargs)

    @neural_computation
    def compute_new_state(self, step_inputs):
        xh_t, h_tm1 = map(step_inputs.get, ["xh_t", "state"])
        if not xh_t:
            xh_t = 0

        h_t = self.activate(xh_t + T.dot(h_tm1, self.W_h) + self.b_h)

        return {"state": h_t}

    @neural_computation
    def merge_inputs(self, input_var, additional_inputs=None):
        if not additional_inputs:
            additional_inputs = []
        all_inputs = ([input_var] if input_var else []) + additional_inputs
        h_inputs = []
        for x, weights in zip(all_inputs, self.input_weights):
            wi, = weights
            h_inputs.append(T.dot(x, wi))
        merged_inputs = {
            "xh_t": sum(h_inputs)
        }
        return merged_inputs

    def prepare(self):
        self.output_dim = self.hidden_size

        self.W_h = self.create_weight(self.hidden_size, self.hidden_size, "h", initializer=self.outer_init)
        self.b_h = self.create_bias(self.hidden_size, "h")

        self.register_parameters(self.W_h, self.b_h)

        self.input_weights = []
        if self._input_type == "sequence":
            normal_input_dims = [self.input_dim]
        else:
            normal_input_dims = []

        all_input_dims = normal_input_dims + self.additional_input_dims
        for i, input_dim in enumerate(all_input_dims):
            wi = self.create_weight(input_dim, self.hidden_size, "wi_{}".format(i+1), initializer=self.outer_init)
            weights = [wi]
            self.input_weights.append(weights)
            self.register_parameters(*weights)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

from theano.tensor.var import TensorVariable
from deepy.utils.map_dict import MapDict

def convert_to_theano_var(obj):
    """
    Convert neural vars to theano vars.
    :param obj: NeuralVariable or list or dict or tuple
    :return: theano var, test var, tensor found, neural var found
    """
    from deepy.core.neural_var import NeuralVariable
    if type(obj) == tuple:
        return tuple(convert_to_theano_var(list(obj)))
    if type(obj) == list:
        unpacked_list = map(convert_to_theano_var, obj)
        normal_list = []
        test_list = []
        theano_var_found = False
        neural_var_found = False
        for normal_var, tensor_found, neural_found in unpacked_list:
            normal_list.append(normal_var)
            if tensor_found: theano_var_found = True
            if neural_found: neural_var_found = True
        return normal_list, theano_var_found, neural_var_found
    elif type(obj) == dict:
        normal_map = {}
        theano_var_found = False
        neural_var_found = False
        for key in obj:
            normal_var, tensor_found, neural_found = convert_to_theano_var(obj[key])
            normal_map[key] = normal_var
            if tensor_found: theano_var_found = True
            if neural_found: neural_var_found = True
        return normal_map, theano_var_found, neural_var_found
    elif type(obj) == MapDict:
        normal_map = {}
        theano_var_found = False
        neural_var_found = False
        for key in obj:
            normal_var, tensor_found, neural_found = convert_to_theano_var(obj[key])
            normal_map[key] = normal_var
            if tensor_found: theano_var_found = True
            if neural_found: neural_var_found = True
        return MapDict(normal_map), theano_var_found, neural_var_found
    elif type(obj) == NeuralVariable:
        theano_tensor = obj.tensor
        theano_tensor.tag.last_dim = obj.dim()
        return theano_tensor, False, True
    elif type(obj) == TensorVariable:
        return obj, True, False
    elif type(obj) == slice:
        normal_args = []
        theano_var_found = False
        neural_var_found = False
        for arg in [obj.start, obj.stop, obj.step]:
            normal_var, tensor_found, neural_found = convert_to_theano_var(arg)
            normal_args.append(normal_var)
            if tensor_found: theano_var_found = True
            if neural_found: neural_var_found = True
        return slice(*normal_args), theano_var_found, neural_var_found
    else:
        return obj, False, False

def convert_to_neural_var(obj):
    """
    Convert object and a test object into neural var.
    :param obj: tensor or list or dict or tuple
    :param test_obj: NeuralVar or list or dict or tuple
    :return:
    """
    from theano.tensor.var import TensorVariable
    from deepy.core.neural_var import NeuralVariable
    if type(obj) == list:
        return [convert_to_neural_var(item) for item in obj]
    elif type(obj) == tuple:
        return tuple(convert_to_neural_var(list(obj)))
    elif type(obj) == dict:
        merged_map = {}
        for key in obj:
            merged_map[key] = convert_to_neural_var(obj[key])
        return merged_map
    elif type(obj) == MapDict:
        merged_map = {}
        for key in obj:
            merged_map[key] = convert_to_neural_var(obj[key])
        return MapDict(merged_map)
    elif type(obj) == TensorVariable:
        deepy_var = NeuralVariable(obj)
        if hasattr(obj, 'tag') and hasattr(obj.tag, 'last_dim'):
            deepy_var.output_dim = obj.tag.last_dim
        return deepy_var
    else:
        return obj

def neural_computation(original_func, prefer_tensor=False):
    """
    An annotation to enable theano-based fucntions to be called with NeuralVar.
    :param original_func:
    :param prefer_tensor: a switch to return tensors when no inputs
    :return:
    """

    def wrapper(*args, **kwargs):
        normal_args, tensor_found_in_args, neural_found_in_args = convert_to_theano_var(args)
        normal_kwargs, tensor_found_in_kwargs, neural_found_in_kwargs = convert_to_theano_var(kwargs)

        tensor_found = tensor_found_in_args or tensor_found_in_kwargs
        neural_found = neural_found_in_args or neural_found_in_kwargs

        if tensor_found and neural_found:
            raise Exception("Theano tensor variables can not be used together with neural variables.")

        normal_result = original_func(*normal_args, **normal_kwargs)

        if tensor_found or (not neural_found and prefer_tensor):
            # No neural variables are inputted, so output tensors
            return normal_result
        else:
            # Output neural variables, auto set output_dim
            result_var = convert_to_neural_var(normal_result)
            if (isinstance(normal_result, TensorVariable) and
                    hasattr(normal_result.tag, "test_value") and
                    hasattr(normal_result.tag.test_value, "shape") and
                    normal_result.tag.test_value.shape):
                result_var.output_dim = normal_result.tag.test_value.shape[-1]
            return result_var
    return wrapper

def neural_computation_prefer_tensor(original_func):
    return neural_computation(original_func, prefer_tensor=True)


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import numpy as np
from theano import tensor as T

from deepy.core import env

def onehot_tensor(i_matrix, vocab_size):
    """
    # batch x time
    """
    dim0, dim1 = i_matrix.shape
    i_vector = i_matrix.reshape((-1,))
    hot_matrix = T.extra_ops.to_one_hot(i_vector, vocab_size).reshape((dim0, dim1, vocab_size))
    return hot_matrix

def onehot(size, eye):
    return np.eye(1, size, eye, dtype=env.FLOATX)[0]