#!/usr/bin/env python
from matplotlib.pyplot import imread, imshow, subplot, show, axis
from numpy import arctan2, zeros, append, pi
from numpy.linalg import norm
from scipy.ndimage import convolve1d
from gauss import Gauss1, gD

def in_image(p, F):
    """Check if given pixel coordinates p are within the bound of image F."""
    return 0 <= p[0] < F.shape[0] and 0 <= p[1] < F.shape[1]

def canny(F, s, Tl=None, Th=None):
    """Apply the Canny Edge Detection algorithm with Gauss scale s to an
    image F. Optionally specify a low and high threshold (Tl and Th) for
    hysteresis thresholding."""
    # Noise reduction by a Gaussian filter
    F = gD(F, 1.4, 0, 0)

    # Find intensity gradient
    W = Gauss1(s, 1)
    Gx = convolve1d(F, W, axis=1, mode='nearest')
    Gy = convolve1d(F, W, axis=0, mode='nearest')
    G = zeros(F.shape)
    A = zeros(F.shape, dtype=int)

    for y in xrange(F.shape[0]):
        for x in xrange(F.shape[1]):
            p = (y, x)

            # Gradient norm and rounded angle
            G[p] = norm(append(Gy[p], Gx[p]))
            A[p] = (4 - int(round(4 * arctan2(Gy[p], Gx[p]) / pi))) % 4

    # Non-maximum suppression
    E = zeros(F.shape)
    compare = [((-1, 0), (1, 0)), ((-1, 1), (1, -1)), \
               ((0, 1), (0, -1)), ((1, 1), (-1, -1))]

    for y in xrange(F.shape[0]):
        for x in xrange(F.shape[1]):
            g = G[y, x]
            a, b = compare[A[y, x]]
            na = (y + a[1], x + a[0])
            nb = (y + b[1], x + b[0])

            if (not in_image(na, G) or g > G[na]) \
                    and (not in_image(nb, G) or g > G[nb]):
                E[y, x] = g

    # Only execute hysteresis thresholding if both thresholds are specified
    if Tl is None or Th is None:
        return E

    # Hysteresis thresholding
    # Scale the thresholds to the color range of the image
    Tl *= (E.max() - E.min()) / 255
    Th *= (E.max() - E.min()) / 255
    T = zeros(F.shape, dtype=bool)

    # Clear image borders
    for y in xrange(F.shape[0]):
        E[x, 0] = E[y, F.shape[1] - 1] = 0

    for x in xrange(x, F.shape[1] - 1):
        E[0, x] = E[F.shape[0] - 1, x] = 0

    def follow_nb(y, x):
        """Follow the neighbouring pixels of an edge pixel in E recursively.
        Add pixels with a value higher than Tl to the bitmap."""
        if T[y, x]:
            return

        T[y, x] = True

        for ny in xrange(-1, 2):
            for nx in xrange(-1, 2):
                if (ny or nx) and E[ny, nx] > Tl:
                    follow_nb(ny, nx)

    # Follow edges that have a starting value above Th
    for y in xrange(F.shape[0]):
        for x in xrange(F.shape[1]):
            if E[y, x] > Th:
                follow_nb(y, x)

    return E, T

if __name__ == '__main__':
    from sys import argv, exit

    if len(argv) < 2 or len(argv) == 3:
        print 'Usage: python %s SCALE [ TL TH ]' % argv[0]
        exit(1)

    F = imread('cameraman.png')
    #F = imread('flags.png')

    s = float(argv[1])

    if len(argv) > 3:
        # Execute with tracing edges
        E, T = canny(F, s, float(argv[2]), float(argv[3]))

        subplot(131, title='Original image')
        imshow(F, cmap='gray')
        axis('off')
        subplot(132, title='Gradient magnitudes')
        imshow(E, cmap='gray')
        axis('off')
        subplot(133, title='Thresholds applied')
        imshow(T, cmap='gray')
        axis('off')
    else:
        # Execute until nn-maximum suppression
        E = canny(F, s)

        subplot(121, title='Original image')
        imshow(F, cmap='gray')
        axis('off')
        subplot(122, title='Gradient magnitudes')
        imshow(E, cmap='gray')
        axis('off')

    show()
