import cv2
import numpy as np

# return -1 if user press 'q'. return 1 if user press 'Enter'.
def display_image(window_title, image):
    # 创建可调整大小的窗口
    cv2.namedWindow(window_title, cv2.WINDOW_NORMAL)
    # 显示图像
    cv2.imshow(window_title, image)
    
    while True:
        # 检查窗口是否被关闭
        if cv2.getWindowProperty(window_title, cv2.WND_PROP_VISIBLE) < 1:
            return -1

        key = cv2.waitKey(1) & 0xFF
        if key == ord("q"):
            return -1
        # 'Enter' key is detected!
        if key == 13:
            return 1


class PointSelector(object):
    """
    ---------------------------------------------------
    | A simple gui point selector.                    |
    | Usage:                                          |
    |                                                 |
    | 1. call the `loop` method to show the image.    |
    | 2. click on the image to select key points,     |
    |    press `d` to delete the last points.         |
    | 3. press `q` to quit, press `Enter` to confirm. |
    ---------------------------------------------------
    """

    POINT_COLOR = (0, 0, 255)
    FILL_COLOR = (0, 255, 255)

    def __init__(self, image, title="PointSelector"):
        self.original_image = image.copy()  # 保存原始图像
        self.image = image
        self.title = title
        self.keypoints = []
        self.window_width = image.shape[1]
        self.window_height = image.shape[0]
        self.scale = 1.0  # 缩放比例
        self.last_window_size = (self.window_width, self.window_height)

    def draw_image(self):
        """
        Display the selected keypoints and draw the convex hull.
        """
        # 检查窗口大小是否改变
        current_width = cv2.getWindowImageRect(self.title)[2] if cv2.getWindowProperty(self.title, cv2.WND_PROP_VISIBLE) >= 1 else self.window_width
        current_height = cv2.getWindowImageRect(self.title)[3] if cv2.getWindowProperty(self.title, cv2.WND_PROP_VISIBLE) >= 1 else self.window_height
        
        if current_width != self.last_window_size[0] or current_height != self.last_window_size[1]:
            self.window_width = current_width
            self.window_height = current_height
            self.last_window_size = (current_width, current_height)
            # 计算新的缩放比例
            self.scale = min(current_width / self.original_image.shape[1], 
                           current_height / self.original_image.shape[0])

        # 基于当前缩放比例调整点坐标
        scaled_keypoints = [
            (int(x * self.scale), int(y * self.scale)) 
            for x, y in self.keypoints
        ]
        
        # 创建缩放后的图像副本
        scaled_image = cv2.resize(self.original_image, 
                                 (self.window_width, self.window_height))
        
        # 绘制选中的关键点
        for i, pt in enumerate(scaled_keypoints):
            cv2.circle(scaled_image, pt, 6, self.POINT_COLOR, -1)
            cv2.putText(scaled_image, str(i), (pt[0], pt[1] - 15),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, self.POINT_COLOR, 2)

        # 如果有两个点，绘制连接线
        if len(scaled_keypoints) == 2:
            p1, p2 = scaled_keypoints
            cv2.line(scaled_image, p1, p2, self.POINT_COLOR, 2)

        # 如果有两个以上的点，绘制凸包
        if len(scaled_keypoints) > 2:
            mask = self.create_mask_from_pixels(scaled_keypoints,
                                               (self.window_height, self.window_width))
            scaled_image = self.draw_mask_on_image(scaled_image, mask)

        cv2.imshow(self.title, scaled_image)

    def onclick(self, event, x, y, flags, param):
        """
        点击点(x, y)会将该点添加到列表并重新绘制图像。
        考虑窗口缩放，将点击坐标转换回原始图像坐标。
        """
        if event == cv2.EVENT_LBUTTONDOWN:
            # 将点击坐标转换回原始图像坐标
            orig_x = int(x / self.scale)
            orig_y = int(y / self.scale)
            print(f"click ({orig_x}, {orig_y}) (scaled: ({x}, {y}))")
            self.keypoints.append((orig_x, orig_y))
            self.draw_image()

    def loop(self):
        """
        Press "q" will exit the gui and return False
        press "d" will delete the last selected point.
        Press "Enter" will exit the gui and return True.
        """
        # 创建可调整大小的窗口
        cv2.namedWindow(self.title, cv2.WINDOW_NORMAL)
        # 设置窗口初始大小为图像大小
        cv2.resizeWindow(self.title, self.image.shape[1], self.image.shape[0])
        cv2.setMouseCallback(self.title, self.onclick, param=())
        self.draw_image()

        while True:
            # 检查窗口是否被关闭
            if cv2.getWindowProperty(self.title, cv2.WND_PROP_VISIBLE) < 1:
                return False

            key = cv2.waitKey(1) & 0xFF

            # 按q键返回False
            if key == ord("q"):
                return False

            # 按d键删除最后一个点
            if key == ord("d"):
                if len(self.keypoints) > 0:
                    x, y = self.keypoints.pop()
                    print(f"Delete ({x}, {y})")
                    self.draw_image()

            # 按Enter键确认
            if key == 13:
                return True

    def create_mask_from_pixels(self, pixels, image_shape):
        """
        Create mask from the convex hull of a list of pixels.
        """
        pixels = np.int32(pixels).reshape(-1, 2)
        hull = cv2.convexHull(pixels)
        mask = np.zeros(image_shape[:2], np.int8)
        cv2.fillConvexPoly(mask, hull, 1, lineType=8, shift=0)
        mask = mask.astype(bool)
        return mask

    def draw_mask_on_image(self, image, mask):
        """
        Paint the region defined by a given mask on an image.
        """
        new_image = np.zeros_like(image)
        new_image[:, :] = self.FILL_COLOR
        mask = np.array(mask, dtype=np.uint8)
        new_mask = cv2.bitwise_and(new_image, new_image, mask=mask)
        cv2.addWeighted(image, 1.0, new_mask, 0.5, 0.0, image)
        return image