from numpy import *

__all__ = ['gabor_fn']

def gabor_fn(sigma_x,sigma_y,theta,lambda_,psi,gamma,width=None, height=None):
    '''
    Gabor function, adapted from Wikipedia's matlab code
    '''
    if width is None:
        sz_x = 6*sigma_x
        width = sz_x
    else:
        sz_x = width
    if sz_x%2==0:
        sz_x = sz_x+1

    if height is None:
        sz_y= 6*sigma_y
        height = sz_y
    else:
        sz_y = height
    if sz_y%2==0:
        sz_y = sz_y+1
     
    x,y = mgrid[-int(sz_x/2):int(sz_x/2)+1,int(-sz_y/2):int(sz_y/2)+1]
     
    # Rotation 
    x_theta = x*cos(theta)+y*sin(theta)
    y_theta = -x*sin(theta)+y*cos(theta)
     
    gb = exp(-.5*(x_theta**2/sigma_x**2+gamma**2*y_theta**2/sigma_y**2))*cos(2*pi/lambda_*x_theta+psi)
    return gb[0:width,0:height]

if __name__=='__main__':
    import pylab
    theta = pi/4
    gb = gabor_fn(100.,100.,theta,300.,0.,0.8,500,520)
    print gb.shape
    pylab.gray()
    pylab.imshow(gb)
    pylab.colorbar()
    pylab.show()