#!/usr/bin/python2


"""Demonstration of image convolution in Python

Yes, scipy.ndimage does a better job at it.
This is just a demonstration of proper optimization of numpy routines
"""


import numpy


def apply_filter_mask(image, mask):
    """Applies some kind of convolution filter

    Arguments:
    image -- Array-like object (preferably numpy.array). First and second
             dimension: Height and width. Third dimension: Color channels
    mask -- Array-like object of the filter coefficients. Two-dimensional.
            For example ((1, 2, 1), (2, 4, 2), (1, 2, 1)) for a Gauss filter

    Return value:
    A new image object of the same shape and type as image. Borders are
    initialized to 0
    """
    # First, ensure we don't get any surprises with the argument types
    image, mask = (numpy.asarray(arr) for arr in (image, mask))
    # Alright. Let's assume we have the mask mentioned in the description.
    # Over the course of this function we will need every pixel in image
    # multiplied by 1, 2, and 4. So let's do this only once in the beginning.
    #
    # Also, because we are multiplying a huge array with a scalar, it is worth
    # checking whether we actually need to do a multiplication. For 0 we skip
    # processing altogether.
    # This can be extended to deal with cases where a bit shift is sufficient
    # but you get the idea.
    parts = {factor: image if factor == 1 else image * factor for factor
             in set(mask.flat) if factor}
    # We don't need to care about the number of color channels.
    # Numpy's broadcast and vector operation rules will do it automatically
    height, width = image.shape[0:2]
    # For the iteration through the mask, we need the factor and the relative
    # offsets compared to the center pixel in the mask
    maskheight, maskwidth = mask.shape
    verticalborder, horizontalborder = (length / 2 for length
                                        in (maskheight, maskwidth))
    # Now build the output array. As an optimization, we could copy the first
    # parts entry, saving the zero initialization and an addition but let's
    # keep it simple for now
    result = numpy.zeros_like(image)
    # Here we just leave the borders at 0 but in general, when we need special
    # treatment for border pixels, it is faster to do it separately outside
    # the loop for the center frame
    center_frame = result[verticalborder:-verticalborder,
                          horizontalborder:-horizontalborder]
    # Let's skip being fancy and use a simple nested for loop for the rest.
    # It's only a few iterations, anyway
    for verticaloff, maskline in enumerate(mask):
        vertical_end = height - maskheight + verticaloff + 1
        for horizontaloff, factor in enumerate(maskline):
            if not factor:
                continue
            part = parts[factor]
            horizontal_end = width - maskwidth + horizontaloff + 1
            center_frame += part[verticaloff:vertical_end,
                                 horizontaloff:horizontal_end]
    # Alright. Nearly done. Now we just need to normalize the result. This is
    # an integer division and therefore extremely slow, especially on ARM CPUs.
    # So, let's go the extra mile and replace it with a bit shift if we can
    normfactor = mask.sum()
    if normfactor == 1: # maybe we are dealing with a normalized mask of floats?
        return result
    # In Python3 we could use math.log2. Oh well, let's do it the ugly way
    try:
        binnorm = bin(normfactor)
    except TypeError: # again, probably a mask of floats
        center_frame /= normfactor
    else:
        if binnorm.count('1') > 1: # no power of two
            center_frame /= normfactor
        else:
            normshift = binnorm.count('0') - 1
            center_frame >>= normshift
    return result


def main():
    image = numpy.random.random_integers(0, 10, 25).reshape(5, 5)
    mask = numpy.array(((1, 2, 1), (2, 4, 2), (1, 2, 1)))
    filtered = apply_filter_mask(image, mask)
    print "%s\n*\n%s\n=\n%s\n" % (image, mask, filtered)


if __name__ == '__main__':
    main()