import cv2
import numpy as np
import sys
import math
import argparse
 
# function for 2d convolution (used as is from previous project)
def conv2(img, kernel, padding_style):
   
    h, w = kernel.shape
   
    yO, xO, cO = img.shape
 
    paddedImg = np.zeros([yO + math.floor(h/2) + math.floor(h/2), xO + math.floor(w/2) + math.floor(w/2), cO], dtype='float32')
    paddedImg[math.floor(h/2) : yO + math.floor(h/2),math.floor(w/2) : xO + math.floor(w/2),:] = img[:,:,:]
   
    # appply padding
    if(padding_style == 'zero-padding'):
        paddedImg = paddedImg
    elif(padding_style == 'wrap-around'):
        i = math.floor(w/2) - 1
        j = xO + math.floor(w/2) - 1
        while(i >= 0):
            paddedImg[:,i,:] = paddedImg[:,j,:]
            j -= 1
            i -= 1
        i = math.floor(h/2) - 1
        j = yO + math.floor(w/2) - 1
        while(i >= 0):
            paddedImg[i,:,:] = paddedImg[j,:,:]
            j -= 1
            i -= 1
        i = math.floor(w/2) + xO
        j = math.floor(w/2)
        while(i < paddedImg.shape[1]):
            paddedImg[:,i,:] = paddedImg[:,j,:]
            j += 1
            i += 1
        i = math.floor(h/2) + yO
        j = math.floor(h/2)
        while(i < paddedImg.shape[0]):
            paddedImg[i,:,:] = paddedImg[j,:,:]
            j += 1
            i += 1
    elif(padding_style == 'copy-edge'):
        i = math.floor(w/2) - 1
        j = math.floor(w/2)
        while(i >= 0):
            paddedImg[:,i,:] = paddedImg[:,j,:]
            i -= 1
        i = math.floor(h/2) - 1
        j = math.floor(h/2)
        while(i >= 0):
            paddedImg[i,:,:] = paddedImg[j,:,:]
            i -= 1
        i = math.floor(w/2) + xO
        j = xO + math.floor(w/2) - 1
        while(i < paddedImg.shape[1]):
            paddedImg[:,i,:] = paddedImg[:,j,:]
            i += 1
        i = math.floor(h/2) + yO
        j = yO + math.floor(h/2) - 1
        while(i < paddedImg.shape[0]):
            paddedImg[i,:,:] = paddedImg[j,:,:]
            i += 1
    elif(padding_style == 'reflect-across-edge'):
        i = math.floor(w/2) - 1
        j = math.floor(w/2)
        while(i >= 0):
            paddedImg[:,i,:] = paddedImg[:,j,:]
            j += 1
            i -= 1
        i = math.floor(h/2) - 1
        j = math.floor(h/2)
        while(i >= 0):
            paddedImg[i,:,:] = paddedImg[j,:,:]
            j += 1
            i -= 1
        i = math.floor(w/2) + xO
        j = xO + math.floor(w/2) - 1
        while(i < paddedImg.shape[1]):
            paddedImg[:,i,:] = paddedImg[:,j,:]
            j -= 1
            i += 1
        i = math.floor(h/2) + yO
        j = yO + math.floor(h/2) - 1
        while(i < paddedImg.shape[0]):
            paddedImg[i,:,:] = paddedImg[j,:,:]
            j -= 1
            i += 1
    img = paddedImg
    y, x, c = img.shape
    b, g, r = cv2.split(img)
    channels = [b, g, r]
   
    if(all((r==b).flatten()) and all((g==b).flatten())):
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        c = 1
 
    # apply convlution by point wise multiplication in frequency domain
    if(c == 1):
        imgF = np.fft.fft2(img)
        imgFS = np.fft.fftshift(imgF)
        kernelF = np.fft.fft2(kernel, [y,x])
        kernelFS = np.fft.fftshift(kernelF)
        imgF = imgFS * kernelFS
        imgishift = np.fft.ifftshift(imgF)
        img = np.fft.ifft2(imgishift)
        img = np.abs(img)
 
    elif(c == 3):
        bF = np.fft.fft2(b)
        bFS = np.fft.fftshift(bF)
        kernelF = np.fft.fft2(kernel, [y,x])
        kernelFS = np.fft.fftshift(kernelF)
        bF = bFS * kernelFS
        bishift = np.fft.ifftshift(bF)
        b = np.fft.ifft2(bishift)
        b = np.abs(b)
 
        gF = np.fft.fft2(g)
        gFS = np.fft.fftshift(gF)
        gF = gFS * kernelFS
        gishift = np.fft.ifftshift(gF)
        g = np.fft.ifft2(gishift)
        g = np.abs(g)
 
        rF = np.fft.fft2(r)
        rFS = np.fft.fftshift(rF)
        rF = rFS * kernelFS
        rishift = np.fft.ifftshift(rF)
        r = np.fft.ifft2(rishift)
        r = np.abs(r)
        img = cv2.merge((b, g, r))
 
    img = img[h-1:h-1+yO, w-1:w-1+xO]
   
    return img
 
# function to resize using nearest neighbour interpolation
def resize(img, newY, newX):
   
    res = np.zeros([newY, newX, img.shape[2]], dtype = 'uint8')
    i = 0
    factY = newY/img.shape[0]
    factX = newX/img.shape[1]
    while(i < newY):
        j = 0
        while(j < newX):
            res[i][j][:] = img[math.floor(i/factY)][math.floor(j/factX)][:]
            j += 1
        i += 1
    # kernel = gaussian_kernel(3,3)
    # res = conv2(np.uint8(res), kernel, 'reflect-across-edge')
    if(len(res.shape) == 2):
        res = cv2.cvtColor(np.uint8(res),cv2.COLOR_GRAY2RGB)
    return np.uint8(res)
 
 
# function to generate gaussian kernel
def gaussian_kernel(size, size_y):
    size = int(size)
    int(size_y)
    x, y = np.mgrid[-size:size+1, -size_y:size_y+1]
    g = np.exp(-(x**2/float(size)+y**2/float(size_y)))
    return g / g.sum()
 
# function for pyramid down (like cv2.pyrDown)
def pyrDown(img):
    k = int((4 * 2 * 2 / 6) +0.5)
    kernel = gaussian_kernel(k,k)
    imgBlur = conv2(np.uint8(img), kernel, 'copy-edge')
    if(len(imgBlur.shape) == 2):
        imgBlur = cv2.cvtColor(np.uint8(imgBlur),cv2.COLOR_GRAY2RGB)
    y, x, c = imgBlur.shape
    # imgDS = cv2.resize(imgBlur, (round(y/2), round(x/2)), interpolation = cv2.INTER_NEAREST)
    imgDS = resize(imgBlur, round(y/2), round(x/2))
    return np.uint8(imgDS)
 
# function for pyramid up (like cv2.pyrUp)
def pyrUp(img):
    y, x, c = img.shape
 
    k = int((4 * 2 * 2 / 6) +0.5)
    kernel = gaussian_kernel(k,k)
    # imgUS = cv2.resize(img, (y*2, x*2), interpolation = cv2.INTER_NEAREST)
    imgUS = resize(img, y*2, x*2)
    imgUS = conv2(np.uint8(imgUS), kernel, 'copy-edge')
    if(len(imgUS.shape) == 2):
        imgUS = cv2.cvtColor(np.uint8(imgUS),cv2.COLOR_GRAY2RGB)
    return np.uint8(imgUS)
 
# function to compute gaussian and laplacian pyramids of given image
def computePyr(img, number_of_layers):
    G = img.copy()
    gPyr = [G]
    max_num_of_layers = 0
    for i in range(number_of_layers - 1):
        G = pyrDown(G)
        gPyr.append(G)
        max_num_of_layers += 1
        if(G.shape[0] == 1 or G.shape[1] == 1):
            break
 
    lPyr = [gPyr[len(gPyr) - 1]]
    L = lPyr[0]
    for i in range(max_num_of_layers, 0, -1):
        GE = pyrUp(gPyr[i])
        L = cv2.subtract(gPyr[i-1],GE)
        lPyr.append(L)
       
    return gPyr, lPyr
 
# mouse callback function to draw rectangle
def draw_rect(event,x,y,flags,param):
    global ix,iy,drawing,mode, img, imgC, xMask, yMask
 
    if event == cv2.EVENT_LBUTTONDOWN:
        img = imgC.copy()
        drawing = True
        ix,iy = x,y
 
    elif event == cv2.EVENT_MOUSEMOVE:
        xMask,yMask = x,y
        if drawing == True:
            img = imgC.copy()
            cv2.rectangle(img,(ix,iy),(x,y),(0,255,0),2)
 
    elif event == cv2.EVENT_LBUTTONUP:
        xMask,yMask = x,y
        drawing = False
        cv2.rectangle(img,(ix,iy),(x,y),(0,255,0),2)
 
# mouse callback function to draw ellipse
def draw_ellipse(event,x,y,flags,param):
    global ix,iy,drawing,mode, img, imgC, xMask, yMask
 
    if event == cv2.EVENT_LBUTTONDOWN:
        img = imgC.copy()
        drawing = True
        ix,iy = x,y
    elif event == cv2.EVENT_MOUSEMOVE:
        xMask,yMask = x,y
        if drawing == True:
            img = imgC.copy()
            cv2.ellipse(img,(int((ix+xMask)/2),int((iy+yMask)/2)),(int((-ix+xMask)/2),int((-iy+yMask)/2)),0,0,360,(0,255,0),2)
 
    elif event == cv2.EVENT_LBUTTONUP:
        xMask,yMask = x,y
        drawing = False
        cv2.ellipse(img,(int((ix+xMask)/2),int((iy+yMask)/2)),(int((-ix+xMask)/2),int((-iy+yMask)/2)),0,0,360,(0,255,0),2)
 
# function to create foreground mask using mouse callbacks
def drawMask(fg, shape):
    global ix,iy,drawing,mode, img, imgC, xMask, yMask
    cv2.namedWindow('Drag Mouse to draw ' + shape + ' mask, press ENTER to continue')
    if(shape == 'rect'):
        cv2.setMouseCallback('Drag Mouse to draw ' + shape + ' mask, press ENTER to continue',draw_rect)
    else:
        cv2.setMouseCallback('Drag Mouse to draw ' + shape + ' mask, press ENTER to continue',draw_ellipse)
    while(1):
        cv2.imshow('Drag Mouse to draw ' + shape + ' mask, press ENTER to continue',img)
        k = cv2.waitKey(1) & 0xFF
        if k == 13:
            break
 
    cv2.destroyAllWindows()
    # r = cv2.selectROI(fg, False)
    fg_mask = np.zeros([fg.shape[0],fg.shape[1], fg.shape[2]], dtype = 'uint8')
    if(shape == 'rect'):
        fg_mask[iy:yMask,ix:xMask, :] = 255
    else:
        cv2.ellipse(fg_mask,(int((ix+xMask)/2),int((iy+yMask)/2)),(int((-ix+xMask)/2),int((-iy+yMask)/2)),0,0,360,(255,255,255),-1)
    return fg_mask
 
# function to blend 2 images
def blend(bg, fg, number_of_layers, shape):
    fg_mask = drawMask(fg, shape)
    g_bg, l_bg = computePyr(bg, number_of_layers)
    g_fg, l_fg = computePyr(fg, number_of_layers)
    g_mask, l_mask = computePyr(fg_mask, number_of_layers)
 
    LS = []
    i = len(g_mask) - 1
    for l_b, l_f in zip(l_bg, l_fg):
        y, x, c = l_b.shape
        ls = (l_f * (g_mask[i]/255)) + (l_b * (1 - (g_mask[i]/255)))
        LS.append(np.uint8(ls))
        i -= 1
    # now reconstruct
    ls_ = LS[0]
    for i in range(1,len(LS)):
        ls_ = pyrUp(ls_)
        ls_ = cv2.add(np.uint8(ls_), np.uint8(LS[i]))
 
    return ls_
       
 
# driver code
 
parser = argparse.ArgumentParser(description='Pyramid Blending')
requiredArgs = parser.add_argument_group('required arguments')
requiredArgs.add_argument('-f', action='store', dest='fg_path',help='Enter foreground path (relative path)', required=True)
requiredArgs.add_argument('-b', action='store', dest='bg_path', type=str, help='Enter background path (relative path)', required=True)
requiredArgs.add_argument('-n', action='store', dest='number_of_layers', type=str, help='Enter number of layers in pyramid', required=True)
requiredArgs.add_argument('-s', action='store', dest='shape', type=str, help='Enter shape for creating mask (rect/ellipse)', required=True)
args = parser.parse_args()
bg_path = args.bg_path
fg_path = args.fg_path
number_of_layers = int(args.number_of_layers)
shape = args.shape
fg = cv2.imread(fg_path)
bg = cv2.imread(bg_path)
print(fg.shape)
print(bg.shape)
img = fg.copy()
imgC = fg.copy()
drawing = False
ix,iy = -1,-1
xMask,yMask = -1,-1
result = blend(bg, fg, number_of_layers, shape)
cv2.imshow("Result", result)
cv2.waitKey(0)
cv2.destroyAllWindows()