# Path: mozi/utils/theano_utils.py
# def shared_zeros(shape, dtype=floatX, name=None, **kwargs):
#     return sharedX(np.zeros(shape), dtype=dtype, name=name, **kwargs)
# 
# def alloc_zeros_matrix(*dims):
#     return T.alloc(np.cast[floatX](0.), *dims)
# 
# def shared_ones(shape, dtype=floatX, name=None, **kwargs):
#     return sharedX(np.ones(shape), dtype=dtype, name=name, **kwargs)
# 
# Path: mozi/layers/template.py
# class Template(object):
#     """
#     DESCRIPTION:
#         The interface to be implemented by any layer.
#     """
#     def __init__(self):
#         '''
#         FIELDS:
#             self.params: any params from the layer that needs to be updated
#                          by backpropagation can be put inside self.params
#             self.updates: use for updating any shared variables
#         '''
#         self.params = []
#         self.updates = []
# 
#     def _test_fprop(self, state_below):
#         '''
#         DESCRIPTION:
#             This is called during validating/testing of the model.
#         PARAM:
#             state_below: the input to layer
#         '''
#         return self._train_fprop(state_below)
# 
#     def _train_fprop(self, state_below):
#         '''
#         DESCRIPTION:
#             This is called during every training batch whereby the output from the
#             model will be used to update the parameters during error backpropagation.
#         PARAM:
#             state_below: the input to layer
#         '''
#         raise NotImplementedError()
# 
# 
#     def _layer_stats(self, state_below, layer_output):
#         """
#         DESCRIPTION:
#             Layer stats is used for debugging the layer by allowing user to peek
#             at the weight values or the layer output or any parameters of interest
#             during training. By computing the values of parameter of interest,
#             for example T.max(self.W) and put in the return list, the training will
#             print the maximum of the weight in the layer after every epoch.
#         PARAM:
#             state_below: the input to layer
#             layer_output: the output from the layer
#         RETURN:
#             A list of tuples of [('name_a', var_a), ('name_b', var_b)] whereby var is scalar
#         """
#         return []
# 
# Path: mozi/weight_init.py
# class OrthogonalWeight(WeightInitialization):
#     def __init__(self, scale=1.1):
#         self.scale = scale
# 
#     def __call__(self, dim, name='W', **kwargs):
#         ''' From Lasagne
#         '''
#         flat_shape = (dim[0], np.prod(dim[1:]))
#         a = np.random.normal(0.0, 1.0, flat_shape)
#         u, _, v = np.linalg.svd(a, full_matrices=False)
#         # pick the one with the correct shape
#         q = u if u.shape == flat_shape else v
#         q = q.reshape(dim)
#         return sharedX(name=name, value=self.scale * q[:dim[0],:dim[1]], borrow=True, **kwargs)
# 
# class GaussianWeight(WeightInitialization):
#     def __init__(self, mean=0, std=0.1):
#         self.mean = mean
#         self.std = std
# 
#     def __call__(self, dim, name='W', **kwargs):
#         W_values = np.random.normal(loc=self.mean, scale=self.std, size=dim)
#         return sharedX(name=name, value=W_values, borrow=True, **kwargs)
# 
# class Identity(WeightInitialization):
#     def __init__(self, scale=1):
#         self.scale = scale
# 
#     def __call__(self, dim, name='W', **kwargs):
#         if len(dim) != 2 or dim[0] != dim[1]:
#             raise Exception("Identity matrix initialization can only be used for 2D square matrices")
#         else:
#             return sharedX(self.scale * np.identity(dim[0]), **kwargs)


from mozi.utils.theano_utils import shared_zeros, alloc_zeros_matrix, shared_ones
from mozi.layers.template import Template
from mozi.weight_init import OrthogonalWeight, GaussianWeight, Identity
import theano.tensor as T
import theano



class LSTM(Template):

    def __init__(self, input_dim, output_dim, truncate_gradient=-1, return_sequences=True,
                weight_init=OrthogonalWeight(), inner_init=GaussianWeight(mean=0, std=0.1)):

        self.input_dim = input_dim
        self.output_dim = output_dim
        self.truncate_gradient = truncate_gradient
        self.return_sequences = return_sequences

        self.W_i = weight_init((self.input_dim, self.output_dim))
        self.U_i = inner_init((self.output_dim, self.output_dim))

# Path: mozi/utils/theano_utils.py
# def shared_zeros(shape, dtype=floatX, name=None, **kwargs):
#     return sharedX(np.zeros(shape), dtype=dtype, name=name, **kwargs)
# 
# def alloc_zeros_matrix(*dims):
#     return T.alloc(np.cast[floatX](0.), *dims)
# 
# def shared_ones(shape, dtype=floatX, name=None, **kwargs):
#     return sharedX(np.ones(shape), dtype=dtype, name=name, **kwargs)
# 
# Path: mozi/layers/template.py
# class Template(object):
#     """
#     DESCRIPTION:
#         The interface to be implemented by any layer.
#     """
#     def __init__(self):
#         '''
#         FIELDS:
#             self.params: any params from the layer that needs to be updated
#                          by backpropagation can be put inside self.params
#             self.updates: use for updating any shared variables
#         '''
#         self.params = []
#         self.updates = []
# 
#     def _test_fprop(self, state_below):
#         '''
#         DESCRIPTION:
#             This is called during validating/testing of the model.
#         PARAM:
#             state_below: the input to layer
#         '''
#         return self._train_fprop(state_below)
# 
#     def _train_fprop(self, state_below):
#         '''
#         DESCRIPTION:
#             This is called during every training batch whereby the output from the
#             model will be used to update the parameters during error backpropagation.
#         PARAM:
#             state_below: the input to layer
#         '''
#         raise NotImplementedError()
# 
# 
#     def _layer_stats(self, state_below, layer_output):
#         """
#         DESCRIPTION:
#             Layer stats is used for debugging the layer by allowing user to peek
#             at the weight values or the layer output or any parameters of interest
#             during training. By computing the values of parameter of interest,
#             for example T.max(self.W) and put in the return list, the training will
#             print the maximum of the weight in the layer after every epoch.
#         PARAM:
#             state_below: the input to layer
#             layer_output: the output from the layer
#         RETURN:
#             A list of tuples of [('name_a', var_a), ('name_b', var_b)] whereby var is scalar
#         """
#         return []
# 
# Path: mozi/weight_init.py
# class OrthogonalWeight(WeightInitialization):
#     def __init__(self, scale=1.1):
#         self.scale = scale
# 
#     def __call__(self, dim, name='W', **kwargs):
#         ''' From Lasagne
#         '''
#         flat_shape = (dim[0], np.prod(dim[1:]))
#         a = np.random.normal(0.0, 1.0, flat_shape)
#         u, _, v = np.linalg.svd(a, full_matrices=False)
#         # pick the one with the correct shape
#         q = u if u.shape == flat_shape else v
#         q = q.reshape(dim)
#         return sharedX(name=name, value=self.scale * q[:dim[0],:dim[1]], borrow=True, **kwargs)
# 
# class GaussianWeight(WeightInitialization):
#     def __init__(self, mean=0, std=0.1):
#         self.mean = mean
#         self.std = std
# 
#     def __call__(self, dim, name='W', **kwargs):
#         W_values = np.random.normal(loc=self.mean, scale=self.std, size=dim)
#         return sharedX(name=name, value=W_values, borrow=True, **kwargs)
# 
# class Identity(WeightInitialization):
#     def __init__(self, scale=1):
#         self.scale = scale
# 
#     def __call__(self, dim, name='W', **kwargs):
#         if len(dim) != 2 or dim[0] != dim[1]:
#             raise Exception("Identity matrix initialization can only be used for 2D square matrices")
#         else:
#             return sharedX(self.scale * np.identity(dim[0]), **kwargs)


# Path: mozi/layers/recurrent.py
from mozi.utils.theano_utils import shared_zeros, alloc_zeros_matrix, shared_ones
from mozi.layers.template import Template
from mozi.weight_init import OrthogonalWeight, GaussianWeight, Identity
import theano.tensor as T
import theano




class LSTM(Template):

    def __init__(self, input_dim, output_dim, truncate_gradient=-1, return_sequences=True,
                weight_init=OrthogonalWeight(), inner_init=GaussianWeight(mean=0, std=0.1)):

        self.input_dim = input_dim
        self.output_dim = output_dim
        self.truncate_gradient = truncate_gradient
        self.return_sequences = return_sequences

        self.W_i = weight_init((self.input_dim, self.output_dim))
        self.U_i = inner_init((self.output_dim, self.output_dim))

        self.b_i = shared_zeros((self.output_dim), name='b_i')