import numpy as np

# --------------------------------------------------------------------------------
def scale_depth(depth, img):
    img_sc = np.zeros(img.shape).astype(img.dtype)
    for i in range(img.shape[-1]):
        img_sc[:,:,i] = depth * img[:,:,i]

    return img_sc
# --------------------------------------------------------------------------------


# ---------------------------------------------------------------------------------------------------------------
def blend_fg_bg(bg_img, fg_img, bg_depth, fg_depth):
    # Blends subject and background using lightwraping technique
    depth = bg_depth

    SEG_THRESH = 200
    # ksize = 5
    # kernel = cv2.getStructuringElement(cv2.MORPH_ERODE, (ksize, ksize))
    # depth_erode = cv2.erode(fg_depth.copy(), kernel, iterations=1)

    # Define background border around subject
    # border = fg_depth - depth_erode

    # bg_border = bg_img.copy()
    # bg_border[border < 50] = 0

    # bg_border = cv2.GaussianBlur(bg_border, (5, 5), 1)

    # Add weighted border to to subject
    # added_fg_img = cv2.addWeighted(bg_border, 0.3, fg_img, 0.7, 0)

    fg_matte = fg_depth.astype(np.float64) / 255.
    inv_fg_matte = 1 - fg_matte

    # Scale foreground and background by depth values and then add.
    blended_fg = scale_depth(fg_matte, fg_img)
    blended_bg = scale_depth(inv_fg_matte, bg_img)

    blended = blended_bg + blended_fg
    blended_img = blended.astype(np.uint8)

    # Get total depth by combining both foreground and background depth/seg maps.
    if depth is not None:
        depth[fg_depth > SEG_THRESH] = fg_depth[fg_depth > SEG_THRESH]

    return blended_img, depth
# ---------------------------------------------------------------------------------------------------------------