import cv2
import numpy as np
import pyautogui
from PIL import Image
import time
import os

current_directory = os.path.dirname(__file__)
pic = os.path.join(current_directory, "red.png")

# 捕捉敌人攻击特效颜色
class get_img_contours:
    def __init__(self, lower_hsv, upper_hsv, screenshot):
        self.lower_hsv = lower_hsv
        self.upper_hsv = upper_hsv
        self.screenshot = screenshot
        self.img = cv2.cvtColor(np.array(self.screenshot), cv2.COLOR_RGB2BGR)
        self.hsv_img = cv2.cvtColor(self.img, cv2.COLOR_BGR2HSV)

    def get_contours(self):
        mask = cv2.inRange(self.hsv_img, self.lower_hsv, self.upper_hsv)
        
        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)

        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        return contours
    
# 计算轮廓高度与宽度
def get_rectangle(contours):
    rectangle_info = []
    # 绘制矩形框
    for contour in contours.get_contours():
        x, y, w, h = cv2.boundingRect(contour)
        print(f"矩形框的宽: {w}, 高: {h}");
        rectangle_info.append((w, h))
    return rectangle_info

if __name__ == "__main__":

    time.sleep(2)

    width, height = pyautogui.size()

    # 传入自定义的黄色光效
    lower_hsv_yellow = np.array([20, 108, 236])
    upper_hsv_yellow = np.array([40, 235, 255])

    # 传入自定义的红色光效
    lower_hsv_red = np.array([0, 70, 240])
    upper_hsv_red = np.array([180, 160, 255])

    # 红光图片
    screenshot = Image.open(pic)

    # 创建类的实例
    contours_yellow = get_img_contours(lower_hsv_yellow, upper_hsv_yellow, screenshot)
    contours_red = get_img_contours(lower_hsv_red, upper_hsv_red, screenshot)

    w_h_list_yellow = get_rectangle(contours_yellow)
    w_h_list_red = get_rectangle(contours_red)

    # 查看捕捉结果
    if any(w > 0 and h > 0 for w, h in w_h_list_red):
        # 创建可调节大小的窗口
        cv2.namedWindow('light contours', cv2.WINDOW_NORMAL)
        # 设置窗口大小为图像的大小
        cv2.resizeWindow('light contours', screenshot.width, screenshot.height)
        # 在原始图像上绘制轮廓
        cv2.drawContours(contours_red.img, contours_red.get_contours(), -1, (0, 0, 255), 1)

        # 绘制矩形框
        for contour in contours_red.get_contours():
            x, y, w, h = cv2.boundingRect(contour)
            cv2.rectangle(contours_red.img, (x, y), (x + w, y + h), (255, 0, 0), 1)

        #  显示结果
        cv2.imshow('light contours', contours_red.img)
        cv2.waitKey(0)
        cv2.destroyAllWindows()
    else:
        print("未检测到红色轮廓")
        
    time.sleep(0.01)