import cv2
import numpy as np
import os
import time
from PIL import Image
from pltreader import PltReader
from concurrent.futures import ThreadPoolExecutor
from shapely import affinity
from scipy.optimize import linear_sum_assignment

def get_single_size_info(filepath, size_num, english_name):
    with open(filepath, 'r') as f:
        reader = PltReader(f)
        output = reader.get_output(tolerance=10)

        # 整幅输出
        full_sheet_path = english_name + '_full_sheet.png'
        cv2.imwrite(full_sheet_path, output.debug_full_sheet(scale_factor=72/1016))
        #cv2.imwrite(full_sheet_path, output.debug_full_sheet(scale_factor=0.1))
        print(f"已保存整幅裁片图像到: {full_sheet_path}")

        # 获取单片输出信息
        return output.get_single_size_info(size_num=size_num)

def get_principal_angle(polygon):
    '''计算主轴角度'''
    coords = np.array(polygon.exterior.coords)
    # 计算协方差矩阵
    cov = np.cov(coords.T)
    # 特征值分解
    eig_vals, eig_vecs = np.linalg.eig(cov)
    # 找到最大特征值对应的特征向量
    principal_vec = eig_vecs[:, np.argmax(eig_vals)]
    # 计算角度
    angle = np.degrees(np.arctan2(principal_vec[1], principal_vec[0]))
    return angle

def get_different_size_match_result(clusters):
    base_cluster = [piece for piece in clusters[0] if piece["parent"] == None]
    result = {piece["index"]:[] for piece in base_cluster}

    for size_index in range(1, len(clusters)):
        compare_cluster = [piece for piece in clusters[size_index] if piece["parent"] == None]
        assert len(base_cluster) == len(compare_cluster)

        # 初始化成本矩阵
        cost_matrix = np.zeros((len(base_cluster), len(base_cluster)))
        # 初始化旋转角度矩阵，表示每个裁片的最佳旋转角度
        rotation_matrix = np.zeros((len(base_cluster), len(base_cluster)))

        # 把每个裁片对应上
        for i in range(len(base_cluster)):
            for j in range(len(compare_cluster)):
                poly1 = base_cluster[i]["data"]
                poly2 = compare_cluster[j]["data"]

                # 质心对齐
                poly1 = affinity.translate(poly1, xoff=-poly1.centroid.x, yoff=-poly1.centroid.y)
                poly2 = affinity.translate(poly2, xoff=-poly2.centroid.x, yoff=-poly2.centroid.y)

                # 缩放到相同面积
                target_area = 10000
                scale_factor1 = (target_area / poly1.area) ** 0.5
                scale_factor2 = (target_area / poly2.area) ** 0.5
                # 绕原点缩放
                poly1 = affinity.scale(poly1, xfact=scale_factor1, yfact=scale_factor1, origin=(0, 0))
                poly2 = affinity.scale(poly2, xfact=scale_factor2, yfact=scale_factor2, origin=(0, 0))

                candidates = [0, 90, 180, 270]

                dist_min = float('inf')
                best_angle = 0
                for angle in candidates:
                    rotated_back = affinity.rotate(poly2, angle, origin='centroid')
                    dist = poly1.hausdorff_distance(rotated_back)
                    if dist < dist_min:
                        dist_min = dist
                        best_angle = angle

                cost_matrix[i, j] = dist_min
                rotation_matrix[i, j] = best_angle

        row_ind, col_ind = linear_sum_assignment(cost_matrix)
        for i, j in zip(row_ind, col_ind):
            result[base_cluster[i]['index']].append((compare_cluster[j]['index'], cost_matrix[i, j], rotation_matrix[i, j]))

    for base_index, matches in result.items():
        print(f"{base_index} <->", end=" ")
        for match in matches:
            print(f"{match[0]} (d={match[1]:.2f}  r={match[2]}°)", end="; ")
        print()
    print("\n")
    
    return result


def export_pieces_by_size_group(filepath, size_labels, output_dir="output", dpi=150, rotation=0):
    """
    按尺码和组号导出裁片PNG
    
    Args:
        filepath: PLT文件路径
        size_labels: 尺码标签列表，如 ["S", "M", "L", "XL", "2XL", "3XL", "4XL"]
        output_dir: 输出目录
        dpi: 输出分辨率 (DPI)
        rotation: 旋转角度，可选值:
            0   - 不旋转
            90  - 顺时针旋转90°
            -90 - 逆时针旋转90°
            180 - 旋转180°
    
    输出文件命名: S-1.png, M-1.png, L-1.png, ...
    """
    start_time = time.time()
    
    # 根据 DPI 计算 scale_factor (1016 绘图仪单位 = 1 英寸)
    scale_factor = dpi / 1016
    size_num = len(size_labels)
    
    print(f"==========开始读取PLT文件: {filepath}==========")
    with open(filepath, 'r') as f:
        reader = PltReader(f)
        output = reader.get_output(tolerance=10)
    
    # 获取尺码聚类信息
    clusters = output.get_single_size_info(size_num=size_num)
    
    # 获取不同尺码间的匹配关系
    print("==========开始匹配不同尺码的相同裁片==========")
    base_cluster = [piece for piece in clusters[0] if piece["parent"] == None]
    match_result = {piece["index"]: [piece["index"]] for piece in base_cluster}  # 基准尺码的index作为key
    
    for size_index in range(1, len(clusters)):
        compare_cluster = [piece for piece in clusters[size_index] if piece["parent"] == None]
        
        # 初始化成本矩阵
        cost_matrix = np.zeros((len(base_cluster), len(compare_cluster)))
        
        for i in range(len(base_cluster)):
            for j in range(len(compare_cluster)):
                poly1 = base_cluster[i]["data"]
                poly2 = compare_cluster[j]["data"]
                
                # 质心对齐
                poly1 = affinity.translate(poly1, xoff=-poly1.centroid.x, yoff=-poly1.centroid.y)
                poly2 = affinity.translate(poly2, xoff=-poly2.centroid.x, yoff=-poly2.centroid.y)
                
                # 缩放到相同面积
                target_area = 10000
                scale_factor1 = (target_area / poly1.area) ** 0.5
                scale_factor2 = (target_area / poly2.area) ** 0.5
                poly1 = affinity.scale(poly1, xfact=scale_factor1, yfact=scale_factor1, origin=(0, 0))
                poly2 = affinity.scale(poly2, xfact=scale_factor2, yfact=scale_factor2, origin=(0, 0))
                
                # 尝试不同旋转角度
                dist_min = float('inf')
                for angle in [0, 90, 180, 270]:
                    rotated = affinity.rotate(poly2, angle, origin='centroid')
                    dist = poly1.hausdorff_distance(rotated)
                    if dist < dist_min:
                        dist_min = dist
                
                cost_matrix[i, j] = dist_min
        
        # 匈牙利算法匹配
        row_ind, col_ind = linear_sum_assignment(cost_matrix)
        for i, j in zip(row_ind, col_ind):
            base_index = base_cluster[i]['index']
            matched_index = compare_cluster[j]['index']
            match_result[base_index].append(matched_index)
    
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 构建 index -> node 的映射（包含所有节点，不只是parent为None的）
    all_nodes = {node["index"]: node for node in output.nodes}
    
    # 导出图片
    rotation_desc = {0: "不旋转", 90: "顺时针90°", -90: "逆时针90°", 180: "旋转180°"}.get(rotation, f"旋转{rotation}°")
    print(f"==========开始导出裁片图片到: {output_dir}/ ({rotation_desc})==========")
    group_id = 1
    for base_index, matched_indices in match_result.items():
        print(f"组 {group_id}: ", end="")
        for size_idx, piece_index in enumerate(matched_indices):
            size_label = size_labels[size_idx]
            
            # 找到对应的node（包括其子节点）
            node = all_nodes[piece_index]
            nodes_to_draw = [node] + node['child']
            
            # 绘制图像
            img = output._draw_nodes(nodes_to_draw, scale_factor, show_id=False)
            
            # OpenCV 的 BGRA 转 Pillow 的 RGBA
            img_rgba = cv2.cvtColor(img, cv2.COLOR_BGRA2RGBA)
            pil_img = Image.fromarray(img_rgba)
            
            # 应用旋转
            if rotation == 90:
                pil_img = pil_img.rotate(-90, expand=True)  # Pillow的rotate是逆时针，所以用负值实现顺时针
            elif rotation == -90:
                pil_img = pil_img.rotate(90, expand=True)   # 逆时针90°
            elif rotation == 180:
                pil_img = pil_img.rotate(180, expand=True)
            
            # 保存文件（写入 DPI 元数据）
            filename = f"{size_label}-{group_id}.png"
            filepath_out = os.path.join(output_dir, filename)
            pil_img.save(filepath_out, dpi=(dpi, dpi))
            print(f"{filename}", end=" ")
        
        print()
        group_id += 1
    
    print(f"\n导出完成！共 {group_id - 1} 组裁片，每组 {len(size_labels)} 个尺码")
    print(f"文件保存在: {output_dir}/")
    
    # ========== 输出每个尺码的裁片中心坐标 ==========
    SAFETY_MARGIN = 10  # 安全边距 10cm
    
    print("\n" + "=" * 70)
    print(f"裁片中心坐标（单位: cm，原点: 左下角，安全边距: {SAFETY_MARGIN}cm，{rotation_desc}）")
    print("=" * 70)
    
    # PLT单位转厘米的系数 (1016单位 = 1英寸 = 2.54cm)
    plt_to_cm = 2.54 / 1016
    
    # 按尺码分组输出坐标
    for size_idx, size_label in enumerate(size_labels):
        # 收集该尺码下所有裁片
        size_pieces = []
        for base_index, matched_indices in match_result.items():
            piece_index = matched_indices[size_idx]
            node = all_nodes[piece_index]
            polygon = node["data"]
            size_pieces.append({
                "group_id": list(match_result.keys()).index(base_index) + 1,
                "polygon": polygon,
                "centroid": polygon.centroid
            })
        
        # 计算该尺码的边界框（用于重置原点）
        all_polygons = [p["polygon"] for p in size_pieces]
        from shapely.geometry import MultiPolygon as ShapelyMultiPolygon
        multi_poly = ShapelyMultiPolygon(all_polygons)
        min_x, min_y, max_x, max_y = multi_poly.bounds
        
        # 原始画布尺寸（厘米）
        orig_width_cm = (max_x - min_x) * plt_to_cm
        orig_height_cm = (max_y - min_y) * plt_to_cm
        
        # 根据旋转调整画布宽高
        if rotation in [90, -90]:
            width_cm, height_cm = orig_height_cm, orig_width_cm  # 宽高互换
        else:
            width_cm, height_cm = orig_width_cm, orig_height_cm
        
        # 含安全边距的画布尺寸
        total_width_cm = width_cm + 2 * SAFETY_MARGIN
        total_height_cm = height_cm + 2 * SAFETY_MARGIN
        
        print(f"\n【尺码 {size_label}】")
        print(f"  裁片区域尺寸: 宽 {width_cm:.2f} cm, 高 {height_cm:.2f} cm")
        print(f"  含安全边距尺寸: 宽 {total_width_cm:.2f} cm, 高 {total_height_cm:.2f} cm")
        print("-" * 70)
        print(f"  {'裁片名称':<10} {'原始中心坐标':<22} {'含安全边距坐标':<22} {'裁片尺寸'}")
        print("-" * 70)
        
        for piece in size_pieces:
            # 原始坐标（以左下角为原点，Y轴向上）
            orig_center_x = (piece["centroid"].x - min_x) * plt_to_cm
            orig_center_y = (piece["centroid"].y - min_y) * plt_to_cm
            
            # 原始裁片尺寸
            piece_bounds = piece["polygon"].bounds
            orig_piece_width = (piece_bounds[2] - piece_bounds[0]) * plt_to_cm
            orig_piece_height = (piece_bounds[3] - piece_bounds[1]) * plt_to_cm
            
            # 根据旋转变换坐标
            if rotation == 90:  # 顺时针90°
                center_x = orig_center_y
                center_y = orig_width_cm - orig_center_x
                piece_width, piece_height = orig_piece_height, orig_piece_width
            elif rotation == -90:  # 逆时针90°
                center_x = orig_height_cm - orig_center_y
                center_y = orig_center_x
                piece_width, piece_height = orig_piece_height, orig_piece_width
            elif rotation == 180:  # 旋转180°
                center_x = orig_width_cm - orig_center_x
                center_y = orig_height_cm - orig_center_y
                piece_width, piece_height = orig_piece_width, orig_piece_height
            else:  # 不旋转
                center_x, center_y = orig_center_x, orig_center_y
                piece_width, piece_height = orig_piece_width, orig_piece_height
            
            # 含安全边距的坐标（加上10cm偏移）
            safe_x = center_x + SAFETY_MARGIN
            safe_y = center_y + SAFETY_MARGIN
            
            # 裁片名称与PNG文件名一致
            piece_name = f"{size_label}-{piece['group_id']}"
            print(f"  {piece_name:<10} ({center_x:7.2f}, {center_y:7.2f}) cm  ({safe_x:7.2f}, {safe_y:7.2f}) cm  {piece_width:.2f} x {piece_height:.2f} cm")
    
    print("\n" + "=" * 70)
    
    # 计算并显示处理时间
    elapsed_time = time.time() - start_time
    if elapsed_time >= 60:
        minutes = int(elapsed_time // 60)
        seconds = elapsed_time % 60
        print(f"处理耗时: {minutes}分{seconds:.2f}秒")
    else:
        print(f"处理耗时: {elapsed_time:.2f}秒")


def get_contour_match_result(standard_clusters, rotated_clusters):

    for standard_cluster, rotated_cluster in zip(standard_clusters, rotated_clusters):
        # 仅保留父节点
        standard_cluster = [piece for piece in standard_cluster if piece["parent"] == None]
        rotated_cluster = [piece for piece in rotated_cluster if piece["parent"] == None]
        assert len(standard_cluster) == len(rotated_cluster)

        # 初始化成本矩阵
        cost_matrix = np.zeros((len(standard_cluster), len(standard_cluster)))
        # 初始化旋转角度矩阵，表示每个裁片的最佳旋转角度
        rotation_matrix = np.zeros((len(standard_cluster), len(standard_cluster)))

        # 把旋转后的每个裁片和标准的裁片对应上
        for i in range(len(standard_cluster)):
            for j in range(len(rotated_cluster)):
                poly1 = standard_cluster[i]["data"]
                poly2 = rotated_cluster[j]["data"]

                #poly1 = poly1.simplify(0)
                #poly2 = poly2.simplify(0)

                # 质心对齐
                poly1 = affinity.translate(poly1, xoff=-poly1.centroid.x, yoff=-poly1.centroid.y)
                poly2 = affinity.translate(poly2, xoff=-poly2.centroid.x, yoff=-poly2.centroid.y)

                # 计算主轴角度
                #angle1 = get_principal_angle(poly1)
                #angle2 = get_principal_angle(poly2)

                # 计算角度差
                #angle_diff = angle1 - angle2
                #candidates = [angle_diff, angle_diff + 90, angle_diff + 180, angle_diff + 270]
                candidates = [0, 90, 180, 270]

                dist_min = float('inf')
                best_angle = 0
                for angle in candidates:
                    rotated_back = affinity.rotate(poly2, angle, origin='centroid')
                    dist = poly1.hausdorff_distance(rotated_back)
                    if dist < dist_min:
                        dist_min = dist
                        best_angle = angle

                cost_matrix[i, j] = dist_min
                rotation_matrix[i, j] = best_angle

        row_ind, col_ind = linear_sum_assignment(cost_matrix)
        print("\n标准裁片与旋转后裁片的最佳匹配结果:")
        for i, j in zip(row_ind, col_ind):
            print(f"标准裁片ID: {standard_cluster[i]['index']}, 旋转后裁片ID: {rotated_cluster[j]['index']}, 距离: {cost_matrix[i, j]:.4f}, 最佳旋转角度: {rotation_matrix[i, j]}度")
        # 打印成本矩阵
        #print(cost_matrix)


def main():

    STANDARD_FILE_PATH = "默认方案-POLO衫长袖-A料-标准.plt"
    ROTATED_FILE_PATH = "默认方案-POLO衫长袖-A料-已旋转.plt"
    SIZE_NUM = 5
    
    # ========== 新功能：按尺码分组导出裁片 ==========
    # 尺码标签列表（根据你的实际情况修改）
    SIZE_LABELS = ["S", "M", "L", "XL", "2XL"]
    
    # 导出标准文件的裁片（按 S-1, M-1, L-1... 命名）
    export_pieces_by_size_group(
        filepath=STANDARD_FILE_PATH,
        size_labels=SIZE_LABELS,
        output_dir="standard_pieces",
        dpi=150,  # 输出分辨率，Photoshop 会正确识别
        rotation=90  # 顺时针旋转90°
    )
    
    # 如果需要导出旋转后文件的裁片，取消下面的注释
    # export_pieces_by_size_group(
    #     filepath=ROTATED_FILE_PATH,
    #     size_labels=SIZE_LABELS,
    #     output_dir="rotated_pieces",
    #     scale_factor=72/1016
    # )
    
    # ========== 以下是原有的匹配功能 ==========
    # print("==========开始读取轮廓数据==========")
    # with ThreadPoolExecutor(max_workers=2) as executor:
    #     standard_future = executor.submit(get_single_size_info, STANDARD_FILE_PATH, SIZE_NUM, "standard")
    #     rotated_future = executor.submit(get_single_size_info, ROTATED_FILE_PATH, SIZE_NUM, "rotated")
    #
    #     standard_clusters = standard_future.result()
    #     rotated_clusters = rotated_future.result()
    #
    # print("==========开始获取两个文件的轮廓匹配结果==========")
    # get_contour_match_result(standard_clusters, rotated_clusters)
    #
    # print("==========开始获取一个文件内不同尺寸的匹配结果==========")
    # print(f"{STANDARD_FILE_PATH}:")
    # get_different_size_match_result(standard_clusters)
    # print(f"{ROTATED_FILE_PATH}:")
    # get_different_size_match_result(rotated_clusters)


if __name__ == "__main__":
    main()
