import cv2
import math
import argparse
import numpy as np
import scipy.misc as misc
import matplotlib.pyplot as plt

desc = "Generate a hybrid image using NumPy's Fourier transform."
epilog = "examples:"
epilog += "\n  python %(prog)s abc.png text.png 25"
parser = argparse.ArgumentParser(description = desc, epilog = epilog, formatter_class = argparse.RawTextHelpFormatter)
parser.add_argument('lowPass', help = 'path of the image to apply low-pass filter to')
parser.add_argument('highPass', help = 'path of the image to apply high-pass filter to')
parser.add_argument('sigma', help = 'parameter sigma of a Gaussian filter kernel')
args = parser.parse_args()

def main():
    def G(rows,cols):
        sigma = int(args.sigma)
        
        def Gaussian(x,y):
            return math.exp(-((x - rows//2)**2 + (y - cols//2)**2)/(2*sigma*sigma))
        
        return np.array([[Gaussian(x,y) for y in range(cols)] for x in range(rows)])
    
    # load grayscale images:
    
    lowPass = cv2.imread(args.lowPass, cv2.IMREAD_GRAYSCALE) # i_1
    highPass = cv2.imread(args.highPass, cv2.IMREAD_GRAYSCALE) # i_2
    
    # compute Fourier transforms:
    
    I1 = np.fft.fftshift(np.fft.fft2(lowPass))
    I2 = np.fft.fftshift(np.fft.fft2(highPass))
    
    # get Gaussian filter kernels:
    
    G1 = G(I1.shape[0],I1.shape[1])
    G2 = G(I2.shape[0],I2.shape[1])
    
    # apply Gaussian filter kernels:
    
    I1G1 = np.fft.ifft2(np.fft.ifftshift(I1*G1))
    I2G2 = np.fft.ifft2(np.fft.ifftshift(I2*(1 - G2)))
    
    # generate hybrid image:
    
    hybrid = I1G1 + I2G2
    
    # display results:
    
    plt.subplot(241), plt.title('$lowPass = i_1$'), plt.xticks([]), plt.yticks([])
    plt.imshow(lowPass, cmap = plt.get_cmap('gray'))
    plt.subplot(242), plt.title('$I_1$'), plt.xticks([]), plt.yticks([])
    plt.imshow(np.log(np.abs(I1)), cmap = plt.get_cmap('gray'))
    plt.subplot(243), plt.title('$G_1$'), plt.xticks([]), plt.yticks([])
    plt.imshow(G1, cmap = plt.get_cmap('gray'))
    plt.subplot(244), plt.title('$I_1\cdot G_1$'), plt.xticks([]), plt.yticks([])
    plt.imshow(np.real(I1G1), cmap = plt.get_cmap('gray'))
    
    plt.subplot(245), plt.title('$highPass = i_2$'), plt.xticks([]), plt.yticks([])
    plt.imshow(highPass, cmap = plt.get_cmap('gray'))
    plt.subplot(246), plt.title('$I_2$'), plt.xticks([]), plt.yticks([])
    plt.imshow(np.log(np.abs(I2)), cmap = plt.get_cmap('gray'))
    plt.subplot(247), plt.title('$1 - G_2$'), plt.xticks([]), plt.yticks([])
    plt.imshow(1 - G2, cmap = plt.get_cmap('gray'))
    plt.subplot(248), plt.title('$I_2\cdot (1 - G_2)$'), plt.xticks([]), plt.yticks([])
    plt.imshow(np.real(I2G2), cmap = plt.get_cmap('gray'))
    
    cv2.imshow('hybrid', cv2.convertScaleAbs(np.real(hybrid)))
    
    plt.gcf().canvas.set_window_title('debug')
    plt.show()
    
    cv2.destroyAllWindows()
    
    # save results:
    
    #misc.imsave('lowPass.png', np.real(I1G1))
    #misc.imsave('highPass.png', np.real(I2G2))
    #cv2.imwrite('hybrid.png', cv2.convertScaleAbs(np.real(hybrid)))
    
if __name__ == "__main__": main()