import theano
from theano import tensor
from theano.sandbox import neighbours

import tags


@tags.tags('nnet', 'max pooling', 'inverse')
def inverse_max_pooling(max_pool_out,output_shape,pooling_shape=(2,2),
                        ignore_borders = True):
    """
    Return a symbolic variable representing the inverse of a max pooling
    on a given tensor.

    Parameters
    ----------
    max_pool_out : 4D tensor
        A Theano variable representing the output of a max pooling
    output_shape : 4D shape
        The shape of the input before pooling
    pooling_shape : 2D shape
        The shape of the pooling windows
    ignore_borders : boolean
        Will pad borders with zeros if true
    """
    # flatten the input and repeat it 
    repeated_input = [max_pool_out.flatten()]*(pooling_shape[0]*pooling_shape[1])

    # concatenate the repeated vectors into
    # a 2D matrix in the format neibs2images wants
    stacked_conv_neibs = T.stack(*repeated_input).T
    
    # then get back a stretched version of the stacked neighbours
    stretch_unpooling_out = \
        neighbours.neibs2images(stacked_conv_neibs,
                                pooling_shape,
                                output_shape,
                                'ignore_borders' if ignore_borders else 'valid')
    return stretch_unpooling_out