from skimage.color import rgb2gray
from helpers import gaussian
from main_2 import hybrid_image

def fourier(im):
   im = rgb2gray(im)
   return np.log(np.abs(np.fft.fftshift(np.fft.fft2(im))))

im1 = sk.img_as_float(skio.imread("img_2/tay_swift_out.jpg"))
im2 = sk.img_as_float(skio.imread("img_2/brook_tay_out.jpg"))

s1 = 5
s2 = 8

im1_ft = fourier(im1)
im2_ft = fourier(im2)

im1_hpf = im1 - gaussian(im1, s1)
im2_lpf = gaussian(im2, s2)

im1_hpf_ft = fourier(im1_hpf)
im2_lpf_ft = fourier(im2_lpf)

hybrid = hybrid_image(im1, im2, s1, s2)
hybrid_ft = fourier(hybrid)

# cheat a little bit to get all the FT colorscales to be the same
fts = [im1_ft, im2_ft, im1_hpf_ft, im2_lpf_ft, hybrid_ft]
ft_max = np.max(fts)
ft_min = np.max(fts)
for f in fts:
   f[0,0] = ft_max
   f[0,1] = ft_min

plt.figure(figsize=(8,20))
#plt.tight_layout()
ax00 = plt.subplot2grid((5,2), (0,0))
ax01 = plt.subplot2grid((5,2), (0,1))
ax00.set_title("Taylor Swift")
ax00.imshow(im1)
ax01.set_title("Brook Taylor")
ax01.imshow(im2)

ax10 = plt.subplot2grid((5,2), (1,0))
ax11 = plt.subplot2grid((5,2), (1,1))
ax10.set_title("FFT(Taylor Swift)")
ax10.imshow(im1_ft)
ax11.set_title("FFT(Brook Taylor)")
ax11.imshow(im2_ft)

ax20 = plt.subplot2grid((5,2), (2,0))
ax21 = plt.subplot2grid((5,2), (2,1))
ax20.set_title("hfp(Taylor Swift)")
ax20.imshow(im1_hpf)
ax21.set_title("lpf(Brook Taylor)")
ax21.imshow(im2_lpf)

ax30 = plt.subplot2grid((5,2), (3,0))
ax31 = plt.subplot2grid((5,2), (3,1))
ax30.set_title("FFT(hfp(Taylor Swift))")
ax30.imshow(im1_hpf_ft)
ax31.set_title("FFT(lpf(Brook Taylor))")
ax31.imshow(im2_lpf_ft)

ax4 = plt.subplot2grid((5,2), (4,0), colspan=2)
ax4.set_title("FFT(hybrid)")
ax4.imshow(hybrid_ft)