# Path: mozi/weight_init.py
# class XavierUniformWeight(WeightInitialization):
#     def __call__(self, dim, name='W', **kwargs):
#         fan_in, fan_out = get_fans(dim)
#         W_values = np.random.uniform(low = -4 * np.sqrt(6. / (fan_in + fan_out)),
#                                      high = 4 * np.sqrt(6. / (fan_in + fan_out)),
#                                      size = dim)
#         return sharedX(name=name, value=W_values, 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)
# 
# 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/utils/theano_utils.py
# def shared_zeros(shape, dtype=floatX, name=None, **kwargs):
#     return sharedX(np.zeros(shape), dtype=dtype, name=name, **kwargs)


import numpy as np
import theano
import theano.tensor as T
from theano.sandbox.cuda.fftconv import conv2d_fft
from theano.tensor.signal.pool import pool_2d
from mozi.weight_init import XavierUniformWeight, GaussianWeight
from mozi.layers.template import Template
from mozi.utils.theano_utils import shared_zeros
from theano.tensor.nnet import conv2d
# from theano.tensor.signal import downsample
# from theano.tensor.nnet.conv import conv2d
# from theano.tensor.signal.conv import conv2d

floatX = theano.config.floatX

class Convolution2D(Template):
    def __init__(self, input_channels, filters, kernel_size=(3,3), stride=(1,1),
                 W=None, b=None, weight_init=GaussianWeight(mean=0, std=0.1), border_mode='valid'):
        '''
        PARAM:
            border_mode: (from theano)
                valid: only apply filter to complete patches of the image. Generates
                output of shape: image_shape - filter_shape + 1
                full: zero-pads image to multiple of filter shape to generate output
                of shape: image_shape + filter_shape - 1
        '''
        self.input_channels = input_channels
        self.filters = filters
        self.kernel_size = kernel_size
        self.stride = stride
        self.border_mode = border_mode

        self.W_shape = (self.filters, self.input_channels) + self.kernel_size
        self.W = W
        if self.W is None:
            self.W = weight_init(self.W_shape, name='W')

        self.b = b
        if self.b is None:

# Path: mozi/weight_init.py
# class XavierUniformWeight(WeightInitialization):
#     def __call__(self, dim, name='W', **kwargs):
#         fan_in, fan_out = get_fans(dim)
#         W_values = np.random.uniform(low = -4 * np.sqrt(6. / (fan_in + fan_out)),
#                                      high = 4 * np.sqrt(6. / (fan_in + fan_out)),
#                                      size = dim)
#         return sharedX(name=name, value=W_values, 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)
# 
# 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/utils/theano_utils.py
# def shared_zeros(shape, dtype=floatX, name=None, **kwargs):
#     return sharedX(np.zeros(shape), dtype=dtype, name=name, **kwargs)


# Path: mozi/layers/convolution.py
import numpy as np
import theano
import theano.tensor as T
from theano.sandbox.cuda.fftconv import conv2d_fft
from theano.tensor.signal.pool import pool_2d
from mozi.weight_init import XavierUniformWeight, GaussianWeight
from mozi.layers.template import Template
from mozi.utils.theano_utils import shared_zeros
from theano.tensor.nnet import conv2d

# from theano.tensor.signal import downsample
# from theano.tensor.nnet.conv import conv2d
# from theano.tensor.signal.conv import conv2d

floatX = theano.config.floatX

class Convolution2D(Template):
    def __init__(self, input_channels, filters, kernel_size=(3,3), stride=(1,1),
                 W=None, b=None, weight_init=GaussianWeight(mean=0, std=0.1), border_mode='valid'):
        '''
        PARAM:
            border_mode: (from theano)
                valid: only apply filter to complete patches of the image. Generates
                output of shape: image_shape - filter_shape + 1
                full: zero-pads image to multiple of filter shape to generate output
                of shape: image_shape + filter_shape - 1
        '''
        self.input_channels = input_channels
        self.filters = filters
        self.kernel_size = kernel_size
        self.stride = stride
        self.border_mode = border_mode

        self.W_shape = (self.filters, self.input_channels) + self.kernel_size
        self.W = W
        if self.W is None:
            self.W = weight_init(self.W_shape, name='W')

        self.b = b
        if self.b is None:

            self.b = shared_zeros(shape=(self.filters,), name='b')