import cv2
import argparse
import numpy as np
import matplotlib.pyplot as plt

desc = "Read in two images and blend them using Direct/Pyramid Blending."
epilog = "examples:"
epilog += "\n  python %(prog)s white.png black.png 6"
parser = argparse.ArgumentParser(description = desc, epilog = epilog, formatter_class = argparse.RawTextHelpFormatter)
parser.add_argument('leftImg', help = 'path of the left image')
parser.add_argument('rightImg', help = 'path of the right image')
parser.add_argument('levels', help = 'number of pyramid levels')
args = parser.parse_args()

def main():
    def constructGaussianPyramid(img, levels):
        image = img.copy()
        pyramid = [image]
        for _ in range(levels-1):
            image = cv2.pyrDown(image)
            pyramid.append(image)
        return pyramid
    
    def constructLaplacianPyramid(gaussian_pyramid):
        levels = len(gaussian_pyramid)
        pyramid = [gaussian_pyramid[levels-1]]
        for i in range(levels-1,0,-1):
            up = cv2.pyrUp(gaussian_pyramid[i])
            image = cv2.subtract(gaussian_pyramid[i-1],up)
            pyramid.insert(0,image)
        return pyramid
    
    leftImg = cv2.imread(args.leftImg, cv2.IMREAD_COLOR)
    rightImg = cv2.imread(args.rightImg, cv2.IMREAD_COLOR)
    levels = int(args.levels)
    
    # direct blending:
    
    direct_blending = np.concatenate((leftImg[:,0:leftImg.shape[1]//2], rightImg[:,rightImg.shape[1]//2:]), axis=1)
    
    # pyramid blending:
    
    leftGaussianPyramid = constructGaussianPyramid(leftImg, levels)
    rightGaussianPyramid = constructGaussianPyramid(rightImg, levels)
    #for i in range(len(rightGaussianPyramid)):
    #    cv2.namedWindow('level' + str(i))
    #    cv2.imshow('level' + str(i), rightGaussianPyramid[i])
    #cv2.waitKey(0)
    #cv2.destroyAllWindows()
    
    leftLaplacianPyramid = constructLaplacianPyramid(leftGaussianPyramid)
    rightLaplacianPyramid = constructLaplacianPyramid(rightGaussianPyramid)
    #for i in range(len(rightLaplacianPyramid)):
    #    cv2.namedWindow('level' + str(i))
    #    cv2.imshow('level' + str(i), rightLaplacianPyramid[i])
    #cv2.waitKey(0)
    #cv2.destroyAllWindows()
    
    joinedLaplacianPyramid = []
    for i in range(levels):
        leftHalf = leftLaplacianPyramid[i][:,0:leftLaplacianPyramid[i].shape[1]//2]
        rightHalf = rightLaplacianPyramid[i][:,rightLaplacianPyramid[i].shape[1]//2:]
        image = np.concatenate((leftHalf,rightHalf), axis=1)
        joinedLaplacianPyramid.append(image)
    #for i in range(len(joinedLaplacianPyramid)):
    #    cv2.namedWindow('level' + str(i))
    #    cv2.imshow('level' + str(i), joinedLaplacianPyramid[i])
    #cv2.waitKey(0)
    #cv2.destroyAllWindows()
    
    pyramid_blending = joinedLaplacianPyramid[-1]
    for i in range(levels-2,0,-1):
        pyramid_blending = cv2.pyrUp(pyramid_blending)
        pyramid_blending = cv2.add(pyramid_blending,joinedLaplacianPyramid[i])
    
    # display results:
    
    plt.subplot(221), plt.title('$leftImg$'), plt.xticks([]), plt.yticks([])
    plt.imshow(cv2.cvtColor(leftImg, cv2.COLOR_BGR2RGB))
    plt.subplot(222), plt.title('$rightImg$'), plt.xticks([]), plt.yticks([])
    plt.imshow(cv2.cvtColor(rightImg, cv2.COLOR_BGR2RGB))
    plt.subplot(223), plt.title('$Direct\ Blending$'), plt.xticks([]), plt.yticks([])
    plt.imshow(cv2.cvtColor(direct_blending, cv2.COLOR_BGR2RGB))
    plt.subplot(224), plt.title('$Pyramid\ Blending$'), plt.xticks([]), plt.yticks([])
    plt.imshow(cv2.cvtColor(pyramid_blending, cv2.COLOR_BGR2RGB))
    
    plt.gcf().canvas.set_window_title('debug')
    plt.show()
    
if __name__ == "__main__": main()