from math import fabs, sin, cos, radians
from cv2 import (
    imread,
    cvtColor,
    COLOR_BGR2GRAY,
    TM_CCOEFF_NORMED,
    matchTemplate,
    getRotationMatrix2D,
    warpAffine,
    rectangle,
    imshow,
    imwrite,
    waitKey,
)
import numpy as np
from create_picture import create_pic


class parts_match:
    def __init__(self, test_pic_path) -> None:
        self.test_img, self.test_img_gray = self.read_test_pic(test_pic_path)

    def read_test_pic(self, test_pic_path):
        test_img = imread(test_pic_path)
        test_img_gray = cvtColor(test_img, COLOR_BGR2GRAY)  # 彩色图扩展成灰度图
        return test_img, test_img_gray

    def match_pic(self, temp_img, threshold):
        test_img = self.test_img
        test_img_gray = self.test_img_gray
        w, h = temp_img.shape[::-1]
        res = matchTemplate(test_img_gray, temp_img, TM_CCOEFF_NORMED)
        loc = np.where(res >= threshold)
        for pt in zip(*loc[::-1]):
            rectangle(test_img, pt, (pt[0] + w, pt[1] + h), (0, 0, 255), 1)

    def run(self, threshold=0.5):
        test_img = self.test_img
        lst_rotation = [x * 15 for x in range(24)]

        def read_temp_pic(pic_path, rotation, filled_color=(255, 255, 255)):
            temp_img = imread(pic_path, 0)
            if rotation > 0:
                """
                使用OpenCV进行旋转后空白填充
                """
                # 计算旋转后的尺寸设为_new
                height, width = temp_img.shape[:2]
                height_new = int(
                    width * fabs(sin(radians(rotation)))
                    + height * fabs(cos(radians(rotation)))
                )
                width_new = int(
                    height * fabs(sin(radians(rotation)))
                    + width * fabs(cos(radians(rotation)))
                )
                mat_rotation = getRotationMatrix2D(
                    (width / 2, height / 2), rotation, 1
                )
                mat_rotation[0, 2] += (width_new - width) / 2
                mat_rotation[1, 2] += (height_new - height) / 2
                img_rotated = warpAffine(
                    temp_img,
                    mat_rotation,
                    (width_new, height_new),
                    borderValue=filled_color,
                )
                # imwrite("img_rotated_%d.jpg" % rotation, img_rotated)
                return img_rotated
            else:
                return temp_img

        for rotation in lst_rotation:
            temp_img = read_temp_pic("./Pictures/template/temp_A.png", rotation)
            self.match_pic(temp_img, threshold)

        imwrite("picture_matched.png", test_img)
        imshow("img", test_img)
        waitKey(0)


if __name__ == "__main__":
    test_pic_path = "./picture_test.png"
    # 初始化建图对象
    obj_create_pic = create_pic(5000, 3000)
    # 获取模拟柔性振动板振动后的工件分散位置和姿态坐标系
    lstdic_coordinate = obj_create_pic.get_coordinate()
    # 创建测试图
    obj_create_pic.create_base_canvas(lstdic_coordinate, test_pic_path)
    # 初始化匹配对象
    obj_parts_match = parts_match(test_pic_path)
    obj_parts_match.run(0.7)