# -*- coding:utf-8 -*-
# @Author len
# @Create 2023/10/30 15:14


import math
import cv2
import numpy as np
import os
import matplotlib.pyplot as plt


def detect_shape(contour, epsilon_factor=0.04):
    epsilon = epsilon_factor * cv2.arcLength(contour, True)
    approx = cv2.approxPolyDP(contour, epsilon, True)

    if len(approx) == 3:
        return "三角形"
    elif len(approx) == 4:
        # # 计算四边形的内角
        # angles = []
        # for i in range(4):
        #     pt1 = approx[i]
        #     pt2 = approx[(i + 1) % 4]
        #     pt3 = approx[(i + 2) % 4]
        #     angle = np.degrees(np.arctan2(pt3[0][1] - pt2[0][1], pt3[0][0] - pt2[0][0]) -
        #                        np.arctan2(pt1[0][1] - pt2[0][1], pt1[0][0] - pt2[0][0]))
        #     angles.append(angle)
        # for angle in angles:
        #     print(angle)
        # # 检查角度是否接近90度
        # if all(85 <= abs(angle) <= 95 for angle in angles):
            return "矩形"
        # else:
        #     return "菱形"
    elif len(approx) == 10:
        return "五角星"
    else:
        area = cv2.contourArea(contour)
        perimeter = cv2.arcLength(contour, True)
        radius = perimeter / (2 * np.pi)
        circle_area = np.pi * (radius ** 2)

        if abs(circle_area - area) < area * 0.2:
            return "圆形"
        else:
            return "未知形状"

def process_image(image_path, epsilon_factor=0.04):
    img = cv2.imread(image_path)
    # img = equalize_image(img)
    h, w, _ = img.shape

    center_color = img[h // 2, w // 2]

    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    center_color_hsv = cv2.cvtColor(np.uint8([[center_color]]), cv2.COLOR_BGR2HSV)[0][0]

    lower_bound = np.array([center_color_hsv[0] - 10, 50, 50])
    upper_bound = np.array([center_color_hsv[0] + 10, 255, 255])
    mask = cv2.inRange(hsv, lower_bound, upper_bound)

    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if not contours:
        return img, "No contours found!"

    largest_contour = max(contours, key=cv2.contourArea)

    shape = detect_shape(largest_contour, epsilon_factor)
    cv2.drawContours(img, [largest_contour], -1, (0, 0, 0), 2)

    return img, shape

# 均衡化图像
def equalize_image(img):
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    hsv[:,:,2] = cv2.equalizeHist(hsv[:,:,2])
    img_equalized = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
    return img_equalized


def process_directory(directory_path, epsilon_factor=0.04):
    filenames = []
    images = []  # 用于存放处理过的图片
    labels = []  # 用于存放识别结果

    for file_name in os.listdir(directory_path):
        file_path = os.path.join(directory_path, file_name)
        if os.path.isfile(file_path) and file_name.lower().endswith(('.png', '.jpg', '.jpeg')):
            img, shape = process_image(file_path, epsilon_factor)
            images.append(img)
            labels.append(shape)

            # 保存文件名
            filenames.append(file_name)

    num_images = len(images)
    num_cols = 6  # 每行显示3张图片
    num_rows = math.ceil(num_images / num_cols)  # 计算需要的行数

    # 使用matplotlib显示图片
    plt.figure(figsize=(15, 10))

    # 设置字体以显示中文
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False

    for i, (img, label, filename) in enumerate(zip(images, labels, filenames)):
        plt.subplot(num_rows, num_cols, i + 1)
        plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        # plt.title(f"{filename}:{label}")
        plt.title(f"{label}")
        plt.axis('off')

    # 添加大标题
    # plt.suptitle(directory_path, fontsize=16, y=1.15)

    plt.tight_layout()
    plt.show()


# 使用方法
directory_path = r'D:\Code\Python\Embedded_game\polygon\exp16\crops\polygon'  # 替换为你的图片文件夹路径
process_directory(directory_path)