import theano
from theano import tensor as T
from theano.tensor import nnet as NN
import numpy as N

class DenoisingAA(T.RModule):

    def __init__(self, input = None, regularize = True, tie_weights = True):
        super(DenoisingAA, self).__init__()

        # MODEL CONFIGURATION
        self.regularize = regularize
        self.tie_weights = tie_weights

        # ACQUIRE/MAKE INPUT
        if not input:
            input = T.matrix('input')
        self.input = theano.External(input)

        # HYPER-PARAMETERS
        self.lr = theano.Member(T.scalar())

        # PARAMETERS
        self.w1 = theano.Member(T.matrix())
        if not tie_weights:
            self.w2 = theano.Member(T.matrix())
        else:
            self.w2 = self.w1.T
        self.b1 = theano.Member(T.vector())
        self.b2 = theano.Member(T.vector())


        # REGULARIZATION COST
        self.regularization = self.build_regularization()


        ### NOISELESS ###

        # HIDDEN LAYER
        self.hidden_activation = T.dot(self.input, self.w1) + self.b1
        self.hidden = self.hid_activation_function(self.hidden_activation)

        # RECONSTRUCTION LAYER
        self.output_activation = T.dot(self.hidden, self.w2) + self.b2
        self.output = self.out_activation_function(self.output_activation)

        # RECONSTRUCTION COST
        self.reconstruction_costs = self.build_reconstruction_costs(self.output)
        self.reconstruction_cost = T.mean(self.reconstruction_costs)

        # TOTAL COST
        self.cost = self.reconstruction_cost
        if self.regularize:
            self.cost = self.cost + self.regularization


        ### WITH NOISE ###
        self.corrupted_input = self.build_corrupted_input()

        # HIDDEN LAYER
        self.nhidden_activation = T.dot(self.corrupted_input, self.w1) + self.b1
        self.nhidden = self.hid_activation_function(self.nhidden_activation)

        # RECONSTRUCTION LAYER
        self.noutput_activation = T.dot(self.nhidden, self.w2) + self.b2
        self.noutput = self.out_activation_function(self.noutput_activation)

        # RECONSTRUCTION COST
        self.nreconstruction_costs = self.build_reconstruction_costs(self.noutput)
        self.nreconstruction_cost = T.mean(self.nreconstruction_costs)

        # TOTAL COST
        self.ncost = self.nreconstruction_cost
        if self.regularize:
            self.ncost = self.ncost + self.regularization


        # GRADIENTS AND UPDATES
        if self.tie_weights:
            self.params = self.w1, self.b1, self.b2
        else:
            self.params = self.w1, self.w2, self.b1, self.b2
        gradients = T.grad(self.ncost, self.params)
        updates = dict((p, p - self.lr * g) for p, g in zip(self.params, gradients))

        # INTERFACE METHODS
        self.update = theano.Method(self.input, self.ncost, updates)
        self.compute_cost = theano.Method(self.input, self.cost)
        self.noisify = theano.Method(self.input, self.corrupted_input)
        self.reconstruction = theano.Method(self.input, self.output)
        self.representation = theano.Method(self.input, self.hidden)
        self.reconstruction_through_noise = theano.Method(self.input, [self.corrupted_input, self.noutput])

    def _instance_initialize(self, obj, input_size = None, hidden_size = None, seed = None, **init):
        if (input_size is None) ^ (hidden_size is None):
            raise ValueError("Must specify hidden_size and target_size or neither.")
        super(DenoisingAA, self)._instance_initialize(obj, **init)
        if seed is not None:
            R = N.random.RandomState(seed)
        else:
            R = N.random
        if input_size is not None:
            sz = (input_size, hidden_size)
            inf = 1/N.sqrt(input_size)
            hif = 1/N.sqrt(hidden_size)
            obj.w1 = R.uniform(size = sz, low = -inf, high = inf)
            if not self.tie_weights:
                obj.w2 = R.uniform(size = list(reversed(sz)), low = -inf, high = inf)
            obj.b1 = N.zeros(hidden_size)
            obj.b2 = N.zeros(input_size)
        if seed is not None:
            obj.seed(seed)
        obj.__hide__ = ['params']

    def build_regularization(self):
        return T.zero() # no regularization!


class SigmoidXEDenoisingAA(DenoisingAA):

    def build_corrupted_input(self):
        self.noise_level = theano.Member(T.scalar())
        return self.random.binomial(T.shape(self.input), 1, 1 - self.noise_level) * self.input

    def hid_activation_function(self, activation):
        return NN.sigmoid(activation)

    def out_activation_function(self, activation):
        return NN.sigmoid(activation)

    def build_reconstruction_costs(self, output):
        reconstruction_cost_matrix = -(self.input * T.log(output) + (1 - self.input) * T.log(1 - output))
        return T.sum(reconstruction_cost_matrix, axis=1)

    def build_regularization(self):
        self.l2_coef = theano.Member(T.scalar())
        if self.tie_weights:
            return self.l2_coef * T.sum(self.w1 * self.w1)
        else:
            return self.l2_coef * (T.sum(self.w1 * self.w1) + T.sum(self.w2 * self.w2))

    def _instance_initialize(self, obj, input_size = None, hidden_size = None, seed = None, **init):
        init.setdefault('noise_level', 0)
        init.setdefault('l2_coef', 0)
        super(SigmoidXEDenoisingAA, self)._instance_initialize(obj, input_size, hidden_size, seed, **init)