from lib.include import *




#----------

# from https://stackoverflow.com/questions/42463172/how-to-perform-max-mean-pooling-on-a-2d-array-using-numpy
# https://gist.github.com/post2web/a92be14008646a3d10b4183c8d35375f
def as_stride(arr, sub_shape, stride):
    '''Get a strided sub-matrices view of an ndarray.
    See also skimage.util.shape.view_as_windows()
    '''
    s0,s1=arr.strides[:2]
    m1,n1=arr.shape[:2]
    m2,n2=sub_shape
    view_shape=(1+(m1-m2)//stride[0],1+(n1-n2)//stride[1],m2,n2)+arr.shape[2:]
    strides=(stride[0]*s0,stride[1]*s1,s0,s1)+arr.strides[2:]
    subs=np.lib.stride_tricks.as_strided(arr,view_shape,strides=strides)
    return subs



def np_pooling (data, kernel_size, stride=None, padding=False, method='max'):

    '''Overlapping pooling on 2D or 3D data.
    <mat>: ndarray, input array to pool.
    <ksize>: tuple of 2, kernel size in (ky, kx).
    <stride>: tuple of 2 or None, stride of pooling window.
              If None, same as <ksize> (non-overlapping pooling).
    <method>: str, 'max for max-pooling,
                   'mean' for mean-pooling.
    <pad>: bool, pad <mat> or not. If no pad, output has size
           (n-f)//s+1, n being <mat> size, f being kernel size, s stride.
           if pad, output has size ceil(n/s).
    Return <result>: pooled matrix.
    '''

    m, n  = data.shape[:2]
    ky,kx = kernel_size
    if stride is None:
        stride=(ky,kx)
    sy,sx=stride

    _ceil=lambda x,y: int(np.ceil(x/float(y)))

    if padding:
        ny=_ceil(m,sy)
        nx=_ceil(n,sx)
        size=((ny-1)*sy+ky, (nx-1)*sx+kx) + mat.shape[2:]
        mat_pad=np.full(size,np.nan)
        mat_pad[:m,:n,...]=data
    else:
        mat_pad=data[:(m-ky)//sy*sy+ky, :(n-kx)//sx*sx+kx, ...]

    view = as_stride(mat_pad,kernel_size,stride)

    if method=='max':
        pool = np.nanmax(view,axis=(2,3))

    elif method=='mean':
        pool = np.nanmean(view,axis=(2,3))

    else:
        raise NotImplementedError

    return pool

#----------

from lib.include import *




#----------

# from https://stackoverflow.com/questions/42463172/how-to-perform-max-mean-pooling-on-a-2d-array-using-numpy
# https://gist.github.com/post2web/a92be14008646a3d10b4183c8d35375f
def as_stride(arr, sub_shape, stride):
    '''Get a strided sub-matrices view of an ndarray.
    See also skimage.util.shape.view_as_windows()
    '''
    s0,s1=arr.strides[:2]
    m1,n1=arr.shape[:2]
    m2,n2=sub_shape
    view_shape=(1+(m1-m2)//stride[0],1+(n1-n2)//stride[1],m2,n2)+arr.shape[2:]
    strides=(stride[0]*s0,stride[1]*s1,s0,s1)+arr.strides[2:]
    subs=np.lib.stride_tricks.as_strided(arr,view_shape,strides=strides)
    return subs



def np_pooling (data, kernel_size, stride=None, padding=False, method='max'):

    '''Overlapping pooling on 2D or 3D data.
    <mat>: ndarray, input array to pool.
    <ksize>: tuple of 2, kernel size in (ky, kx).
    <stride>: tuple of 2 or None, stride of pooling window.
              If None, same as <ksize> (non-overlapping pooling).
    <method>: str, 'max for max-pooling,
                   'mean' for mean-pooling.
    <pad>: bool, pad <mat> or not. If no pad, output has size
           (n-f)//s+1, n being <mat> size, f being kernel size, s stride.
           if pad, output has size ceil(n/s).
    Return <result>: pooled matrix.
    '''

    m, n  = data.shape[:2]
    ky,kx = kernel_size
    if stride is None:
        stride=(ky,kx)
    sy,sx=stride

    _ceil=lambda x,y: int(np.ceil(x/float(y)))

    if padding:
        ny=_ceil(m,sy)
        nx=_ceil(n,sx)
        size=((ny-1)*sy+ky, (nx-1)*sx+kx) + mat.shape[2:]
        mat_pad=np.full(size,np.nan)
        mat_pad[:m,:n,...]=data
    else:
        mat_pad=data[:(m-ky)//sy*sy+ky, :(n-kx)//sx*sx+kx, ...]

    view = as_stride(mat_pad,kernel_size,stride)

    if method=='max':
        pool = np.nanmax(view,axis=(2,3))

    elif method=='mean':
        pool = np.nanmean(view,axis=(2,3))

    else:
        raise NotImplementedError

    return pool

#----------

#---------- # from https://stackoverflow.com/questions/42463172/how-to-perform-max-mean-pooling-on-a-2d-array-using-numpy # https://gist.github.com/post2web/a92be14008646a3d10b4183c8d35375f Get a strided sub-matrices view of an ndarray.     See also skimage.util.shape.view_as_windows() Overlapping pooling on 2D or 3D data.     <mat>: ndarray, input array to pool.     <ksize>: tuple of 2, kernel size in (ky, kx).     <stride>: tuple of 2 or None, stride of pooling window.               If None, same as <ksize> (non-overlapping pooling).     <method>: str, 'max for max-pooling,                    'mean' for mean-pooling.     <pad>: bool, pad <mat> or not. If no pad, output has size            (n-f)//s+1, n being <mat> size, f being kernel size, s stride.            if pad, output has size ceil(n/s).     Return <result>: pooled matrix. #----------