from math import *
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import numpy as np
from numpy import ma
from numpy.fft import *

fig1 = plt.figure(1)
sub1 = fig1.add_subplot(111)

fig2 = plt.figure(2)
sub2 = fig2.add_subplot(111)

test_num = 3

if test_num == 1:
##    square_bump = [0]*100 + [1]*20 + [0]*100
##    square_bumps = ([0]*10+[1]*5)*10+[0]*10
    X = np.arange(-7,7.05,0.1)

##    sinusoidal = [sin(3*x) for x in X]
    pulse = [cos(2*pi*3*x)*exp(-pi*x**2) for x in X]

    L = pulse

    sub1.plot(X, L, 'b-')

    F = fft(L)
    sub2.plot(X, list(map(lambda x: abs(x)**2,F)))

elif test_num == 2:
    src_img = mpimg.imread("Lenna.png")
    img_red = src_img[:,:,:]

    sub1.imshow(src_img)
    sub2.imshow(img_red, cmap='Greys_r')
##    sub2.set_cmap('Greys')

    fig3 = plt.figure(3)
    sub3 = fig3.add_subplot(111)

    F = fftn(src_img)
    F_red = fft2(src_img[:,:,0])
    F_green = fft2(src_img[:,:,1])
    F_blue = fft2(src_img[:,:,2])
##    fftshift(F)
    sub3.imshow(np.log10( np.abs( F ) ))#, cmap='Greys_r')

##    F_masked = F
##    xx, yy = np.meshgrid( *[1 if abs(x-256)<=128 else 0 for x in range(512)]*2 )
##    mask = xx*yy
##    mask = [[1]*512]*128 + [[1 if abs(x-256)>128 else 0 for x in range(512)]]*256 + [[1]*512]*128
##    mask = np.array(mask)#.reshape(512,512)
    foo = np.logical_and(np.arange(512)>=50, np.arange(512)<=462)
    bar = [1 if 100 < abs(k-256) < 200 else 0 for k in range(512)]
    mask = np.logical_and.outer(foo, foo)
##    mask = np.logical_and.outer( *[np.logical_or(np.arange(512) < 100, np.arange(512) > 412)]*2 )
    mask = np.reshape(np.repeat(mask, 3, 1), mask.shape+(3,))

    F_masked = np.multiply( 1-mask, F )#*(cos(pi/4)+sin(pi/4)*1j) + np.multiply( 1-mask, F )
##    F_masked = F * -1#(cos(pi/4)+sin(pi/4)*1j)/2 + F/2

    fig4 = plt.figure(4)
    sub4 = fig4.add_subplot(111)
    sub4.imshow(np.log10( np.abs( F_masked )+1 ), cmap='Greys_r')

    fig5 = plt.figure(5)
    sub5 = fig5.add_subplot(111)

    iF = ifftn(F_masked)
##    fftshift(F2)
    sub5.imshow( abs(iF), cmap='Greys_r')

elif test_num == 3:
    X = np.arange(-7,7.05,0.1)
    Y = [floor(x)-x for x in X]
##    Y = X
##    Y = [.5*x + cos(3*x) for x in X]
##    Y = [ceil(x)*2 + cos(x*pi)*(2*(floor(x)%2)-1) for x in X]
##    Y = [(x//pi)*2 + cos(x%pi + pi) for x in X]

    sub1.plot(X, Y, 'b-')

    F = fft(Y)
##    F2 = ifft(Y)
    sub2.plot(X, list(map(lambda x: abs(x),F)))
##    sub2.plot(X, list(map(lambda x: 125*abs(x),F2)))

plt.show()