import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import scipy.signal as signal
import numpy as n

barImg=mpimg.imread('bar.png')
#extract grey values
barImg = barImg[:,:,3]

imgplot = plt.imshow(barImg, cmap=cm.Greys_r)

img=mpimg.imread('stinkbug.png')
#extract grey values
bugImg = img[:,:,0]


imgplot = plt.imshow(bugImg, cmap=cm.Greys_r)

def gaussian2D(x, y, sigma):
    return (1.0/(1*math.pi*(sigma**2)))*math.exp(-(1.0/(2*(sigma**2)))*(x**2 + y**2))

"""make matrix from function"""
def receptiveFieldMatrix(func):
    h = 30
    g = zeros((h,h))
    for xi in range(0,h):
        for yi in range(0,h):
            x = xi-h/2
            y = yi-h/2
            g[xi, yi] = func(x,y);
    return g

def plotFilter(fun):
    g = receptiveFieldMatrix(fun) 
    plt.imshow(g, cmap=cm.Greys_r)

plotFilter(lambda x,y:gaussian2D(x,y,4))

Img_barGaussian = signal.convolve(barImg,receptiveFieldMatrix(lambda x,y: gaussian2D(x,y,5)), mode='same')

imgplot = plt.imshow(Img_barGaussian, cmap=cm.Greys_r)

Img_bugGaussian = signal.convolve(bugImg,receptiveFieldMatrix(lambda x,y: gaussian2D(x,y,3)), mode='same')

imgplot = plt.imshow(Img_bugGaussian, cmap=cm.Greys_r)

def mexicanHat(x,y,sigma1,sigma2): 
    return gaussian2D(x,y,sigma1) - gaussian2D(x,y,sigma2)

plotFilter(lambda x,y: mexicanHat(x,y,3,4))

Img_barHat = signal.convolve(barImg,receptiveFieldMatrix(lambda x,y:mexicanHat(x,y,3,4)), mode='same')

imgplot = plt.imshow(Img_barHat, cmap=cm.Greys_r)

Img_bugHat = signal.convolve(bugImg,receptiveFieldMatrix(lambda x,y: mexicanHat(x,y,2,3)), mode='same')

imgplot = plt.imshow(Img_bugHat, cmap=cm.Greys_r)

def oddGabor2D(x,y,sigma,orientation):
    return math.sin(x + orientation*y) * math.exp(-(x**2 + y**2)/(2*sigma))

def evenGabor2D(x,y, sigma, orientation):
    return math.cos(x + orientation*y) * math.exp(-(x**2 + y**2)/(2*sigma))

plotFilter(lambda x,y: oddGabor2D(x,y,7,1))

Img_barOddGabor = signal.convolve(barImg,receptiveFieldMatrix(lambda x,y: oddGabor2D(x,y,5,1)), mode='same')

imgplot = plt.imshow(Img_barOddGabor, cmap=cm.Greys_r)

Img_bugOddGabor = signal.convolve(bugImg,receptiveFieldMatrix(lambda x,y: oddGabor2D(x,y,5,1)), mode='same')

imgplot = plt.imshow(Img_bugOddGabor, cmap=cm.Greys_r)

Img_bugOddGaborEdge = signal.convolve(Img_bugHat,receptiveFieldMatrix(lambda x,y: oddGabor2D(x,y,5,1)), mode='same')

imgplot = plt.imshow(Img_bugOddGaborEdge, cmap=cm.Greys_r)

plotFilter(lambda x,y: evenGabor2D(x,y,7,1))

Img_barEvenGabor = signal.convolve(barImg,receptiveFieldMatrix(lambda x,y: evenGabor2D(x,y,5,1)), mode='same')

imgplot = plt.imshow(Img_barEvenGabor, cmap=cm.Greys_r)

Img_bugEvenGabor = signal.convolve(bugImg,receptiveFieldMatrix(lambda x,y: evenGabor2D(x,y,5,1)), mode='same')

imgplot = plt.imshow(Img_bugEvenGabor, cmap=cm.Greys_r)

def edgeEnergy(x,y,sigma, orientation):
    g1= oddGabor2D(x,y,sigma,orientation)
    g2= evenGabor2D(x,y,sigma,orientation)
    return(g1**2+g2**2)


plotFilter(lambda x,y:edgeEnergy(x,y,50,0))

Img_barEdgeEnergy = signal.convolve(barImg,receptiveFieldMatrix(lambda x,y: edgeEnergy(x,y,100,1)), mode='same')
imgplot = plt.imshow(Img_barEdgeEnergy, cmap=cm.Greys_r)

Img_bugEdgeEnergy = signal.convolve(bugImg,receptiveFieldMatrix(lambda x,y: edgeEnergy(x,y,10,1)), mode='same')
imgplot = plt.imshow(Img_bugEdgeEnergy, cmap=cm.Greys_r)