#!/usr/bin/env python
from numpy import zeros, arange, meshgrid, array, matrix
from math import ceil, exp, pi, sqrt
from matplotlib.pyplot import imread, imshow, plot, xlabel, ylabel, show, \
        subplot, xlim, savefig, axis, ticklabel_format
from mpl_toolkits.mplot3d import Axes3D
from scipy.ndimage import convolve, convolve1d
from time import time
from sys import argv, exit

def Gauss(s):
    """Sample a two-dimensional Gaussian function of scale s."""
    size = int(ceil(3 * s))
    r = 2 * size + 1
    W = zeros((r, r))
    t = float(s) ** 2
    a = 1 / (2 * pi * t)

    # Sample the 2D Gaussian function
    for x in xrange(r):
        for y in xrange(r):
            W[x, y] = a * exp(-(((x - size) ** 2 + (y - size) ** 2) / (2 * t)))

    # Make sure that the sum of all kernel values is equal to one
    return W / W.sum()

def f_gauss(x, s):
    """Return the value of a 1D Gaussian function for a given x and scale."""
    return exp(-(x ** 2 / (2 * s ** 2))) / (sqrt(2 * pi) * s)

def f_gauss_der_1(x, s):
    """Return the first derivative of the 1D Gaussian function for a given x
    and scale."""
    return -x * exp(-(x ** 2 / (2 * s ** 2))) / (sqrt(2 * pi) * s ** 3)

def f_gauss_der_2(x, s):
    """Return the second derivative of the 1D Gaussian function for a given x
    and scale."""
    return (x ** 2 - s ** 2) * exp(-(x ** 2 / (2 * s ** 2))) \
            / (sqrt(2 * pi) * s ** 5)

def Gauss1(s, order=0):
    """Sample a one-dimensional Gaussian function of scale s."""
    f = [f_gauss, f_gauss_der_1, f_gauss_der_2][order]
    s = float(s)
    r = int(ceil(3 * s))
    size = 2 * r + 1
    W = zeros(size)

    # Sample the Gaussian function
    W = array([f(x - r, s) for x in xrange(size)])

    if not order:
        # Make sure that the sum of all kernel values is equal to one
        W /= W.sum()

    return W

def gD(F, s, iorder, jorder):
    """Create the Gaussian derivative convolution of image F."""
    Wy = Gauss1(s, iorder)
    Wx = Wy if jorder == iorder else Gauss1(s, jorder)
    G = convolve1d(F, Wy, axis=0, mode='nearest')

    return convolve1d(G, Wx, axis=1, mode='nearest')

def plot_kernel(W, ax):
    """Create a 3D plot of a kernel."""
    x = arange(W.shape[0])
    Y, X = meshgrid(x, x)
    ax.plot_surface(X, Y, array(W), rstride=1, cstride=1, linewidth=0, \
                    antialiased=True, cmap='jet')
    ax.set_xlabel('y')
    ax.set_ylabel('x')
    ax.set_zlabel('g(x, y)')

def exit_with_usage():
    """Print an error message with the program's usage and exit the program."""
    print 'Usage: python %s ( 2d SCALE | 1d SCALE IORDER JORDER' \
          ' | timer METHOD [ REPEAT ] | jet SCALE )' % argv[0]
    exit(1)

if __name__ == '__main__':
    if len(argv) < 2:
        exit_with_usage()

    F = imread('cameraman.png')

    if argv[1] == '2d':
        """Calculate and plot the convolution of the given scale."""
        if len(argv) < 3:
            exit_with_usage()

        s = float(argv[2])
        W = Gauss(s)
        G = convolve(F, W, mode='nearest')

        # Show the original image, kernel and convoluted image respectively
        subplot(131, title='Original image')
        imshow(F, cmap='gray')
        axis('off')
        plot_kernel(W, subplot(132, projection='3d', title='Kernel'))
        subplot(133, title='Convoluted image')
        imshow(G, cmap='gray')
        axis('off')
    elif argv[1] == '1d':
        """Calculate the gaussian kernel using derivatives of the specified
        order in both directions."""
        if len(argv) < 5:
            exit_with_usage()

        s = float(argv[2])
        iorder = int(argv[3])
        jorder = int(argv[4])

        G = gD(F, s, iorder, jorder)

        # Create a 2D weight function for plotting purposes only
        Fy = matrix([Gauss1(s, iorder)])
        Fx = Fy if jorder == iorder else matrix([Gauss1(s, jorder)])
        W = Fy.T * Fx

        # Show the original image, kernel and convoluted image respectively
        subplot(131, title='Original image')
        imshow(F, cmap='gray')
        axis('off')
        plot_kernel(W, subplot(132, projection='3d', title='Kernel'))
        subplot(133, title='Convoluted image')
        imshow(G, cmap='gray')
        axis('off')
    elif argv[1] == 'timer':
        """Time the performance of a 1D/2D convolution and plot the results."""
        if len(argv) < 3:
            exit_with_usage()

        method = argv[2]
        repeat = int(argv[3]) if len(argv) > 3 else 1

        # Time for multiple scales
        S = [1, 2, 3, 5, 7, 9, 11, 15, 19]
        times = []

        for s in S:
            t = 0

            # Average a number of timings to eliminate noise
            for k in xrange(repeat):
                start = time()

                if method == '1d':
                    W = Gauss1(s)
                    G = convolve1d(F, W, axis=0, mode='nearest')
                    convolve1d(G, W, axis=1, mode='nearest')
                elif method == '2d':
                    convolve(F, Gauss(s), mode='nearest')

                t += time() - start

            times.append(t / repeat)

        xlim(S[0], S[-1])
        xlabel('s')
        ylabel('time (s)')
        plot(S, times, 'o-')
    elif argv[1] == 'jet':
        """Show a 2d jet Gaussian plot of the image."""
        if len(argv) < 3:
            exit_with_usage()

        s = float(argv[2])
        subplot(331, title='Fs')
        imshow(gD(F, s, 0, 0), cmap='gray')
        axis('off')
        subplot(334, title='Fsx')
        imshow(gD(F, s, 1, 0), cmap='gray')
        axis('off')
        subplot(335, title='Fsy')
        imshow(gD(F, s, 0, 1), cmap='gray')
        axis('off')
        subplot(337, title='Fsxx')
        imshow(gD(F, s, 2, 0), cmap='gray')
        axis('off')
        subplot(338, title='Fsxy')
        imshow(gD(F, s, 1, 1), cmap='gray')
        axis('off')
        subplot(339, title='Fsyy')
        imshow(gD(F, s, 0, 2), cmap='gray')
        axis('off')
    else:
        exit_with_usage()

    show()
