
import numpy as np

class pool:
    """
    Wrapper for pool().
    """
    
    def __init__(self, pool_size, pooling_function):
        """
        Create a pooling function which will pool 'pool_size' values together using 'pooling_function'
        
        pool_size	=	A tuple 	(height, width)
        pooling_function	=	np.max, np.mean, ...
        
        """
        self._pool_size = pool_size
        self._pooling_function = pooling_function
        
    def apply_pooling(self, img):
        return apply_pooling(img, self._pool_size, self._pooling_function)


def apply_pooling(signal, pool_size, poolingfunc):
    """
    'Downsample' a 1D or 2D signal by taking pool_size samples together (average or max).
    
    poolingfunc = Might be one of: np.mean, np.max
    """

    num_split = np.array(signal.shape) / np.array(pool_size)

    
    # split by columns
    cols = np.split(signal, num_split[0], 0)
    pooled = poolingfunc(cols, 1)

    if len(signal.shape) > 1:
        rows = np.split(pooled, num_split[1], 1)
        pooled = poolingfunc(rows, 2)
        pooled = pooled.transpose()

    return pooled
