import numpy as np
from theano.sandbox.rng_mrg import MRG_RandomStreams as RandomStreams
import theano
import theano.tensor as T
from theano.ifelse import ifelse

class Dropout:

    def __init__(self, train=1):

        # Set seed for the random numbers
        np.random.seed(1234)
        rng = np.random.RandomState(1234)

        # Generate a theano RandomStreams
        self.srng = RandomStreams(rng.randint(999999))
        self.train = train


    def __dropit(self, weight, drop):
        """
        Drops some of the cells of 'weight' matrix based on a binomial representation with prob = drop

        # Arguments:

        :param weight : Parameter matrix of the model
        :type : theano.tensor

        :param drop : proportion to dropout from the 'weight'
        :type : float32

        # Returns: Dropped out matrix
        """
        retain_prob = 1 - drop
        mask = self.srng.binomial(n=1, p=retain_prob, size=weight.shape, dtype='floatX')

        return theano.tensor.cast(weight * mask, theano.config.floatX)

    def __dont_dropit(self, weight, drop):
        """
        Scale the weight parameters if train = 0

        # Arguments:

        :param weight : Parameter matrix of the model
        :type : theano.tensor

        :param drop : proportion to dropout from the 'weight'
        :type : float32

        # Returns: Scaled Dropped out matrix
        """
        return (1 - drop) * theano.tensor.cast(weight, theano.config.floatX)

    def dropout_layer(self, weight, drop):
        """
        Drops some of the cells of 'weight' matrix based on a binomial representation with prob = drop

        # Arguments:

        :param weight : Parameter matrix of the model
        :type : theano.tensor

        :param drop : proportion to dropout from the 'weight'
        :type : float32

        :param train : Flag to switch on/off the dropout
        :type : int32

        # Returns: Dropped out matrix if train = 1
        """
        result = theano.ifelse.ifelse(theano.tensor.eq(self.train, 1), self.__dropit(weight, drop), self.__dont_dropit(weight, drop))
        return result
