import time
import cv2
import math
import numpy as np


def draw_protect_mask(sourceImage, landmark_137):
    h, w, _ = sourceImage.shape
    inpaint_mask = np.zeros((h, w), dtype=np.uint8)
    cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[129:137]), (255,))
    cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[121:129]), (255,))
    cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[22:48]), (255,))
    cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[88:104]), (255,))
    cv2.fillConvexPoly(inpaint_mask, cv2.convexHull(landmark_137[105:121]), (255,))
    return inpaint_mask


def localtranslationwarpfastwithstrength(srcimg, kpt137, startx, starty, endx, endy, radius, strength):
    ddradius = float(radius * radius)
    mask_keep = draw_protect_mask(srcimg, kpt137)
    # mask_keep = cv2.imread('/home/colo/Pictures/test/102.png')
    # mask_keep = cv2.cvtColor(mask_keep, cv2.COLOR_BGR2GRAY)
    # copyimg = np.zeros(srcimg.shape, np.uint8)
    # copyimg = srcimg.copy()
    maskimg = np.zeros(srcimg.shape[:2], np.uint8)
    cv2.circle(maskimg, (startx, starty), math.ceil(radius), (255, 255, 255), -1)
    # cv2.imshow('maskimg_before', maskimg)
    # cv2.imshow('maskimg', maskimg)
    # cv2.imshow('mask_keep', mask_keep)
    # cv2.waitKey()
    maskimg = maskimg * (1-mask_keep/255).astype(np.uint8)

    k0 = 100 / strength  # 计算公式中的|m-c|^2
    ddmc_x = (endx - startx) * (endx - startx)
    ddmc_y = (endy - starty) * (endy - starty)
    h, w, c = srcimg.shape
    mapx = np.vstack([np.arange(w).astype(np.float32).reshape(1, -1)] * h)
    mapy = np.hstack([np.arange(h).astype(np.float32).reshape(-1, 1)] * w)
    distance_x = (mapx - startx) * (mapx - startx)
    distance_y = (mapy - starty) * (mapy - starty)
    distance = distance_x + distance_y
    k1 = np.sqrt(distance)
    ratio_x = (ddradius - distance_x) / (ddradius - distance_x + k0 * ddmc_x)
    ratio_y = (ddradius - distance_y) / (ddradius - distance_y + k0 * ddmc_y)
    ratio_x = ratio_x * ratio_x
    ratio_y = ratio_y * ratio_y
    ux = mapx - ratio_x * (endx - startx) * (1 - k1/radius)
    uy = mapy - ratio_y * (endy - starty) * (1 - k1/radius)
    np.copyto(ux, mapx, where=maskimg == 0)
    np.copyto(uy, mapy, where=maskimg == 0)
    ux = ux.astype(np.float32)
    uy = uy.astype(np.float32)
    copyimg = cv2.remap(srcimg, ux, uy, interpolation=cv2.INTER_LINEAR)
    return copyimg

image = cv2.imread("/home/colo/Pictures/for_hn/102.jpg")
processed_image = image.copy()
startx_left, starty_left, endx_left, endy_left = 170, 123, 190, 74
# startx_right, starty_right, endx_right, endy_right = 287, 275, 192, 233
radius = 60
strength = 100  # 瘦左边脸
t0 = time.time()
processed_image = localtranslationwarpfastwithstrength(processed_image, startx_left, starty_left, endx_left, endy_left, radius, strength)  # 瘦右边脸
# processed_image = localtranslationwarpfastwithstrength(processed_image, startx_right, starty_right, endx_right, endy_right, radius, strength)
# cv2.imwrite("thin.jpg", processed_image)
print('costs', time.time() - t0)
# cv2.imshow('image', image)

# cv2.imshow('processed_image', processed_image)
# cv2.waitKey()