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

# In[ ]:


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


# In[ ]:


class Dropout:

    def __init__(self):
        
        # 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))        
        
    
    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, train = 1):        
        """
        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(train, 1), self.__dropit(weight, drop), self.__dont_dropit(weight, drop))
        return result