'''
参考资料：
https://www.gnu.org/software/hp2xx/hp2xx.html
https://zhuanlan.zhihu.com/p/622090369
'''

import os
import typing
import string
import re
import warnings
from enum import Enum
import shapely
import shapely.ops
from shapely.geometry import LineString, Polygon, MultiPolygon, MultiLineString
import numpy as np
import cv2
import itertools
from shapely.validation import make_valid


def my_warning_handler(message, category, filename, lineno, file=None, line=None):
    print(f"PLT警告: {message}")
warnings.showwarning = my_warning_handler

def assert_warning(condition, message):
    if not condition:
        warnings.warn(message, UserWarning)

class PltCommand(Enum):
    IN = 0
    PU = 1
    PD = 2
    PA = 3

AbstractSyntaxTree = typing.List[typing.Dict]


class VirtualMachineOutput(object):
    '''虚拟机执行后的结果'''
    MIN_FILTER_AREA_FACTOR = 0.0068 ** 2
    MAX_FILTER_AREA_FACTOR = 0.8 ** 2

    def __init__(self, polygons :typing.List[Polygon], dispensable: typing.List[LineString]) -> None:
        for poly in polygons:
            assert isinstance(poly, Polygon)
        for line in dispensable:
            assert isinstance(line, LineString)

        # 过滤结果

        # 计算画布大小
        min_x, min_y, max_x, max_y = MultiPolygon(polygons).bounds

        filter_polygons = []
        for polygon in polygons:
            if (max_x - min_x) * (max_y - min_y) * VirtualMachineOutput.MIN_FILTER_AREA_FACTOR < polygon.area < (max_x - min_x) * (max_y - min_y) *  VirtualMachineOutput.MAX_FILTER_AREA_FACTOR:
                filter_polygons.append(polygon)
            else:
                boundary = polygon.boundary
                if isinstance(boundary, LineString):
                    dispensable.append(boundary)
                elif isinstance(boundary, MultiLineString):
                    dispensable.extend(list(boundary.geoms))

        print(f"过滤了{len(polygons) - len(filter_polygons)}个多边形")
        polygons = filter_polygons

        # 计算多边形的包含关系，构建多边形树
        self.nodes = [{"data":polygon, "child":[], "parent":None, "index":index} for index, polygon in enumerate(polygons)]
        self.dispensable = dispensable

        # 找每个node的最佳parent
        for node1 in self.nodes:
            best_parent = None
            for node2 in self.nodes:
                if node1 == node2:
                    continue

                if node1["index"] == 8 and node2["index"] == 9:
                    pass

                # 判断是否包含
                #if node2["data"].contains(node1["data"]):
                if node2["data"].area > node1["data"].area:
                    intersection = node2["data"].intersection(node1["data"])
                    if intersection.area > node1["data"].area * 0.99:
                        if best_parent is None:
                            best_parent = node2
                        else:
                            if best_parent["data"].contains(node2["data"]):
                                best_parent = node2

            node1["parent"] = best_parent

        # 填充child
        self.tree = []
        for node in self.nodes:
            if node["parent"] is not None:
                node["parent"]["child"].append(node)
            else:
                self.tree.append(node)

        # 构建广度优先遍历
        self.bfs = self.tree.copy()
        index = 0
        while index < len(self.bfs):
            self.bfs.extend(self.bfs[index]["child"])
            index += 1

    def _draw_nodes(self, nodes: typing.List, scale_factor: float, show_id: bool = True) -> np.ndarray:
        # 计算画布大小
        min_x, min_y, max_x, max_y = MultiPolygon([node["data"] for node in nodes]).bounds

        max_x = int(np.ceil(max_x * scale_factor))
        max_y = int(np.ceil(max_y * scale_factor))
        min_x = int(np.floor(min_x * scale_factor))
        min_y = int(np.floor(min_y * scale_factor))

        max_x -= min_x
        max_y -= min_y

        #print("画布大小", max_x, max_y)

        result_img = np.zeros((max_y, max_x, 4), dtype=np.uint8)

        nodes = sorted(nodes, key=lambda node: self.bfs.index(node))

        for node in nodes:
            pts = (np.asarray(node["data"].exterior.coords) * scale_factor - [min_x, min_y]).astype(np.int32)
            pts[:,1] = max_y - pts[:,1]
            cv2.fillPoly(result_img, [pts], (255,255,255,255))
            
            # 最外圈（parent为None）用红色，内部用黑色
            if node["parent"] is None:
                cv2.polylines(result_img, [pts], True, (0,0,255,255), 2)  # 红色 (BGR: 0,0,255)
            else:
                cv2.polylines(result_img, [pts], True, (0,0,0,255), 2)    # 黑色

        # 画无关紧要的东西（dispensable线段）
        multiPolygon = MultiPolygon([node["data"] for node in nodes])
        multiPolygon = make_valid(multiPolygon.buffer(0))
        
        # 获取最外圈轮廓（parent为None的节点）
        outer_contours = [node["data"].exterior for node in nodes if node["parent"] is None]
            
        for shape in self.dispensable:
            if multiPolygon.contains(shape):
                pts = (np.asarray(shape.coords) * scale_factor - [min_x, min_y]).astype(np.int32)
                pts[:,1] = max_y - pts[:,1]
                
                # 判断是否与最外圈轮廓粘连（距离非常近）
                is_connected_to_outer = False
                line_length = shape.length  # 线段长度（PLT单位）
                
                # 距离阈值（PLT单位），小于此值视为粘连
                CONNECT_THRESHOLD = 1
                # 线段长度限制（PLT单位）
                MAX_LINE_LENGTH = 155
                
                if line_length < MAX_LINE_LENGTH:  # 只检查短线段
                    for outer in outer_contours:
                        distance = shape.distance(outer)
                        if distance < CONNECT_THRESHOLD:
                            is_connected_to_outer = True
                            break
                
                if is_connected_to_outer:
                    # ========== 调试开关 ==========
                    # True = 蓝色标记模式（用于调试查看哪些线被检测到）
                    # False = 镂空模式（正式使用，挖空0.2cm宽度）
                    DEBUG_MODE = False
                    
                    if DEBUG_MODE:
                        # 调试模式：蓝色标记，方便查看检测结果
                        cv2.polylines(result_img, [pts], False, (255,0,0,255), 2)  # 蓝色加粗
                    else:
                        # 正式模式：镂空效果（方形端点）
                        # 用矩形多边形代替polylines，避免圆形端点
                        cutout_half_width = int(15 * scale_factor)  # 半宽度
                        if cutout_half_width < 1:
                            cutout_half_width = 1
                        
                        # 对每个线段创建矩形多边形
                        for i in range(len(pts) - 1):
                            p1 = pts[i]
                            p2 = pts[i + 1]
                            
                            # 计算线段方向的垂直向量
                            dx = p2[0] - p1[0]
                            dy = p2[1] - p1[1]
                            length = np.sqrt(dx*dx + dy*dy)
                            if length < 0.001:
                                continue
                                
                            # 垂直方向单位向量
                            nx = -dy / length * cutout_half_width
                            ny = dx / length * cutout_half_width
                            
                            # 构建矩形的四个角点
                            rect_pts = np.array([
                                [p1[0] + nx, p1[1] + ny],
                                [p1[0] - nx, p1[1] - ny],
                                [p2[0] - nx, p2[1] - ny],
                                [p2[0] + nx, p2[1] + ny]
                            ], dtype=np.int32)
                            
                            # 用透明色填充矩形
                            cv2.fillPoly(result_img, [rect_pts], (0,0,0,0))
                else:
                    cv2.polylines(result_img, [pts], False, (0,0,0,255), 1)    # 黑色

        # 画parent为None的id
        if show_id:
            for node in nodes:
                if node["parent"] is None:
                    centroid = node["data"].centroid
                    text_pos = (int((centroid.x * scale_factor) - min_x), int(max_y - (centroid.y * scale_factor) - min_y))
                    cv2.putText(result_img, str(node["index"]), text_pos, cv2.FONT_HERSHEY_SIMPLEX, 3, (255,0,0,255), 5)

        return result_img

    def full_sheet(self, scale_factor: float = 0.1) -> np.ndarray:
        '''整幅输出'''
        return self._draw_nodes(self.bfs, scale_factor)

    def debug_full_sheet(self, scale_factor: float = 0.1) -> np.ndarray:
        """
        用随机颜色画出所有多边形和被过滤掉的线段
        """
        # 计算画布大小，画布大小应该考虑dispensable
        all_shapes = [node["data"] for node in self.nodes] + self.dispensable
        temp = []
        for shape in all_shapes:
            if isinstance(shape, LineString):
                convex_hull = shape.convex_hull
                if isinstance(convex_hull, Polygon):
                    temp.append(convex_hull)
            else:
                temp.append(shape)
        min_x, min_y, max_x, max_y = MultiPolygon(temp).bounds
        max_x = int(np.ceil(max_x * scale_factor))
        max_y = int(np.ceil(max_y * scale_factor))
        min_x = int(np.floor(min_x * scale_factor))
        min_y = int(np.floor(min_y * scale_factor))
        max_x -= min_x
        max_y -= min_y

        result_img = np.zeros((max_y, max_x, 4), dtype=np.uint8)

        # 画多边形
        for node in self.nodes:
            #overlay = result_img.copy()
            color = tuple(np.random.randint(0,256,size=3).tolist() + [255])
            pts = (np.asarray(node["data"].exterior.coords) * scale_factor - [min_x, min_y]).astype(np.int32)
            pts[:,1] = max_y - pts[:,1]
            cv2.fillPoly(result_img, [pts], color)
            if node["parent"] is None:
                cv2.polylines(result_img, [pts], True, (0,0,255,255), 1)
            else:
                cv2.polylines(result_img, [pts], True, (0,0,0,255), 1)
            #alpha = 0.5
            #cv2.addWeighted(overlay, alpha, result_img, 1 - alpha, 0, result_img)

        # 画被过滤掉的线段
        for shape in self.dispensable:
            #overlay = result_img.copy()
            color = tuple(np.random.randint(0,256,size=3).tolist() + [255])
            pts = (np.asarray(shape.coords) * scale_factor - [min_x, min_y]).astype(np.int32)
            pts[:,1] = max_y - pts[:,1]
            cv2.polylines(result_img, [pts], False, color, 1)
            #cv2.addWeighted(overlay, alpha, result_img, 1 - alpha, 0, result_img)

        # 画parent为None的id
        for node in self.nodes:
            if node["parent"] is None:
                centroid = node["data"].centroid
                text_pos = (int((centroid.x * scale_factor) - min_x), int(max_y - (centroid.y * scale_factor) - min_y))
                cv2.putText(result_img, str(node["index"]), text_pos, cv2.FONT_HERSHEY_SIMPLEX, 3, (0,0,255,255), 5)


        return result_img

    def _distance_between_cluster(self, cluster1, cluster2):
        min_distance = float("inf")
        for node1 in cluster1:
            for node2 in cluster2:
                assert isinstance(node1["data"], Polygon)
                assert isinstance(node2["data"], Polygon)

                d = node1["data"].distance(node2["data"])
                if d < min_distance:
                    min_distance = d
        return min_distance

    def get_single_size_info(self, size_num: int) -> typing.List[typing.List]:
        '''获取单码信息'''
        # 距离聚类
        clusters = [[node] for node in self.nodes]

        while len(clusters) > size_num:
            # 每个循环将合并距离最小的两个类
            min_distance_clusters = (None, None)
            min_distance = float("inf")
            for cluster1, cluster2 in itertools.combinations(clusters, 2):
                d = self._distance_between_cluster(cluster1, cluster2)
                if d < min_distance:
                    min_distance = d
                    min_distance_clusters = (cluster1, cluster2)

            clusters.remove(min_distance_clusters[0])
            clusters.remove(min_distance_clusters[1])
            clusters.append(min_distance_clusters[0] + min_distance_clusters[1])

        # 排序
        sorted_clusters = sorted(clusters, key=lambda x: min([node["index"] for node in x]))

        return list(sorted_clusters)

    def single_size(self, size_num: int, scale_factor: float = 0.1) -> typing.List[np.ndarray]:
        '''单码输出'''
        sorted_clusters = self.get_single_size_info(size_num)

        result = []
        for cluster in sorted_clusters:
            result.append(self._draw_nodes(cluster, scale_factor))
        return result

    def single_piece(self, scale_factor: float = 0.1, show_id: bool = False) -> typing.List[typing.Tuple[int, np.ndarray]]:
        '''单片输出
        返回: [(裁片ID, 图像), ...]
        '''
        result = []
        for node in self.nodes:
            # 如果是顶层则新建一副图像
            if node["parent"] is None:
                nodes = [node] + node['child']
                img = self._draw_nodes(nodes, scale_factor, show_id=show_id)
                result.append((node["index"], img))
        return result


class _VirtualMachine(object):
    '''运行PLT文件的虚拟机'''
    def __init__(self, tolerance) -> None:
        self.tolerance = tolerance

    def run(self, abstract_syntax_tree: AbstractSyntaxTree) -> VirtualMachineOutput:
        context = abstract_syntax_tree.copy()
        current_x = 0
        current_y = 0
        self.pen_state = 'UP'
        lines :typing.List[typing.List[typing.Tuple]] = []
        lines_flag :typing.List[int] = []

        while len(context) != 0:
            code = context[0]

            try:
                if code["cmd"] == PltCommand.IN:
                    print("初始化")

                elif code["cmd"] == PltCommand.PA:
                    if self.pen_state == 'UP':
                        pass
                    elif self.pen_state == 'DOWN':
                        new_x = float(code["args"][0])
                        new_y = float(code["args"][1])

                        # 判断是否与已有的线相连接
                        for i in range(len(lines) - 1, -1, -1):
                            if lines_flag[i] == 1:
                                continue
                            if lines[i][0] == (current_x, current_y):
                                lines[i].insert(0, (new_x, new_y))
                                if lines[i][0] == lines[i][-1]:
                                    lines_flag[i] = 1
                                break
                            elif lines[i][0] == (new_x, new_y):
                                lines[i].insert(0, (current_x, current_y))
                                if lines[i][0] == lines[i][-1]:
                                    lines_flag[i] = 1
                                break
                            elif lines[i][-1] == (current_x, current_y):
                                lines[i].append((new_x, new_y))
                                if lines[i][0] == lines[i][-1]:
                                    lines_flag[i] = 1
                                break
                            elif lines[i][-1] == (new_x, new_y):
                                lines[i].append((current_x, current_y))
                                if lines[i][0] == lines[i][-1]:
                                    lines_flag[i] = 1
                                break
                        else:
                            # 不与已有的线相连接
                            lines.append([(current_x, current_y), (new_x, new_y)])
                            lines_flag.append(0)

                    current_x = float(code["args"][0])
                    current_y = float(code["args"][1])

                elif code["cmd"] == PltCommand.PU:
                    self.pen_state = 'UP'

                elif code["cmd"] == PltCommand.PD:
                    self.pen_state = 'DOWN'

            except Exception as e:
                assert_warning(False, f"[code {code}] 执行指令异常: {e}")
            finally:
                context = context[1:]

        # 创建LineString :typing.List[LineString] = []
        lineStrings :typing.List[LineString] = []
        lineStrings = [LineString(line) for line in lines if len(set(line)) > 1]

        '''
        import matplotlib.pyplot as plt
        for line in lineStrings:
            x, y = line.xy
            plt.plot(x, y)
        plt.show()
        '''

        return self.__build_polygons_and_dispensable(lineStrings)

    def __build_polygons_and_dispensable(self, lineStrings: typing.List[LineString]) -> VirtualMachineOutput:

        # 为非常近的线搭桥
        tree = shapely.STRtree(lineStrings)

        bridges = []
        # 遍历所有线段的端点，寻找需要搭桥的地方
        for line in lineStrings:
            endpoints = [shapely.Point(line.coords[0]), shapely.Point(line.coords[-1])]

            for p in endpoints:
                for idx in tree.query(p.buffer(self.tolerance)):
                    other_line = lineStrings[idx]
                    if other_line == line:
                        continue

                    dist = p.distance(other_line)

                    if 0 < dist <= self.tolerance:
                        _, p_target = shapely.ops.nearest_points(p, other_line)
                        bridge = LineString([p, p_target])
                        bridges.append(bridge)

        lineStrings.extend(bridges)

        # 节点化并提取闭合区域
        all_shapes = shapely.ops.unary_union(lineStrings)
        assert isinstance(all_shapes, MultiLineString)
        polygons :typing.List[Polygon] = list(shapely.ops.polygonize(all_shapes))

        # 合并重叠多边形
        merged_polygons = shapely.unary_union(polygons)
        assert isinstance(merged_polygons, MultiPolygon)
        polygons = list(merged_polygons.geoms)

        # 只保留外环
        polygons = [Polygon(poly.exterior) for poly in polygons]

        # 找到多余的线段
        dispensable :typing.List[LineString] = []
        closed_rings = [poly.exterior for poly in polygons]
        if closed_rings:
            leftovers = all_shapes.difference(shapely.ops.unary_union(closed_rings))
            assert isinstance(leftovers, shapely.MultiLineString)
            for shape in leftovers.geoms:
                dispensable.append(shape)
        else:
            dispensable = list(all_shapes.geoms)

        for i in range(len(polygons)):
            polygons[i] = polygons[i].buffer(0)

        print(f"共生成多边形{len(polygons)}个，线段{len(dispensable)}条")

        output = VirtualMachineOutput(polygons, dispensable)

        return output

class _Preprocessor(object):
    '''PLT文件预处理器'''
    def __init__(self) -> None:
        pass

    def get_abstract_syntax_tree(self, src_content: str) -> AbstractSyntaxTree:
        content = src_content

        # 换行符转分号 某些文件的语句末尾没有分号
        content = content.replace("\n", ";")
        content = content.replace(" ", ",")

        # 删除所有空白字符
        for space in string.whitespace:
            content = content.replace(space, "")

        result = []

        # 遍历每一行命令
        for line_index, line in enumerate(content.split(";")):
            if line == "":
                continue

            # 提取开头的英文
            pattern = r'^[a-zA-Z]+'
            match = re.search(pattern, line)

            if match is None:
                assert_warning(False, f"[line {line_index}] 未找到命令: {line}")
                continue

            # 获取对应的指令
            command_str = match.group().upper()
            try:
                command = PltCommand[command_str]
            except KeyError:
                assert_warning(False, f"[line {line_index}] 不支持的指令: {command_str}")
                continue

            # 获取参数
            if len(command_str) == len(line):
                stack = []
            else:
                stack = line[len(command_str):].split(",")

            if command == PltCommand.IN:
                result.append({'cmd': PltCommand.IN, 'args': []})

            elif command == PltCommand.PA:
                for i in range(len(stack) // 2):
                    result.append({'cmd': PltCommand.PA, 'args': [stack[0], stack[1]]})
                    stack = stack[2:]

            elif command == PltCommand.PU:
                result.append({'cmd': PltCommand.PU, 'args': []})
                for i in range(len(stack) // 2):
                    result.append({'cmd': PltCommand.PA, 'args': [stack[0], stack[1]]})
                    stack = stack[2:]

            elif command == PltCommand.PD:
                result.append({'cmd': PltCommand.PD, 'args': []})
                for i in range(len(stack) // 2):
                    result.append({'cmd': PltCommand.PA, 'args': [stack[0], stack[1]]})
                    stack = stack[2:]

            assert_warning(len(stack) == 0, f"[line {line_index}] 栈非空: {stack}")

        return result


class PltReader(object):
    '''PLT文件读取器'''

    def __init__(self, pltfile: typing.TextIO) -> None:
        self.pltfile = pltfile

        preprocessor = _Preprocessor()
        self.abstract_syntax_tree = preprocessor.get_abstract_syntax_tree(self.pltfile.read())

    def get_output(self, tolerance=5):
        machine = _VirtualMachine(tolerance=tolerance)
        output = machine.run(self.abstract_syntax_tree)
        return output


if __name__ == "__main__":
    filepath = "默认方案-POLO衫长袖-A料-标准.plt"

    with open(filepath, 'r') as f:
        reader = PltReader(f)
        output = reader.get_output()

    # 调试整幅输出（随机颜色）
    debug_full_sheet_path = "debug_full_sheet.png"
    cv2.imwrite(debug_full_sheet_path, output.debug_full_sheet(scale_factor=72/1016))
    print(f"调试整幅输出已写入: {debug_full_sheet_path}")

    # 整幅输出
    full_sheet_path = "full_sheet.png"
    cv2.imwrite(full_sheet_path, output.full_sheet(scale_factor=72/1016))
    print(f"整幅输出已写入: {full_sheet_path}")

    # 单码输出
    single_size_path = "single_size"
    os.makedirs(single_size_path, exist_ok=True)
    for index, img in enumerate(output.single_size(size_num=5, scale_factor=72/1016)):
        cv2.imwrite(os.path.join(single_size_path, f"size_{index}.png"), img)
    print(f"单码输出已写入: {single_size_path}/ (共{len(os.listdir(single_size_path))}个文件)")

    # 单片输出 - 每个裁片独立保存为PNG
    single_piece_path = "single_piece"
    os.makedirs(single_piece_path, exist_ok=True)
    pieces = output.single_piece(scale_factor=72/1016, show_id=False)
    for piece_id, img in pieces:
        cv2.imwrite(os.path.join(single_piece_path, f"piece_{piece_id}.png"), img)
    print(f"单片输出已写入: {single_piece_path}/ (共{len(pieces)}个裁片)")