import os.path

import numpy as np
import cv2

# 根据输出结果输出json问价还有边界
def segment_instance_and_contour(seg_map, class_names):
    """
    参数：
        seg_map: HxW 的语义分割图像，像素值表示类别编号
        class_names: 类别编号到类别名的映射，如 {0:'Background', 1:'Farmland', 2:'Forest'}
    返回：
        实例列表，每个包含id、类别名、轮廓点
    """
    instance_id = 1
    results = []

    for cls_id in np.unique(seg_map):
        if cls_id == 0:  # 忽略背景
            continue
        mask = (seg_map == cls_id).astype(np.uint8)

        # 找连通域（不同地块）
        num_labels, labels = cv2.connectedComponents(mask)

        for i in range(1, num_labels):  # 第0是背景
            instance_mask = (labels == i).astype(np.uint8)

            # 找轮廓
            contours, _ = cv2.findContours(instance_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            for contour in contours:
                contour_list = contour.squeeze().tolist()
                if len(contour_list) < 3:
                    continue  # 过滤掉太小的碎片
                results.append({
                    'id': instance_id,
                    'class': class_names.get(cls_id, f"Class{cls_id}"),
                    'contour': contour_list
                })
                instance_id += 1

    return results

def draw_instances_with_ids(seg_map, results, save_path="instance_visual.png"):
    """
    在原始图或灰度图上绘制每个地块的轮廓及其编号。

    参数：
        seg_map: 原始图像或分割图（HxW 或 HxWx3）
        results: 上一步输出的地块信息列表
        save_path: 保存路径
    """
    if len(seg_map.shape) == 2:  # 若为灰度图，转为BGR
        vis_img = cv2.cvtColor(seg_map, cv2.COLOR_GRAY2BGR)
    else:
        vis_img = seg_map.copy()
    # h, w = vis_img.shape[:2]
    for item in results:
        contour = np.array(item['contour']).reshape(-1, 1, 2).astype(np.int32)

        # 随机颜色（根据类别也可指定颜色）
        color = tuple(np.random.randint(0, 255, size=3).tolist())
        cv2.drawContours(vis_img, [contour], -1, color, 2)#线的宽度

        # 计算重心位置以便放编号
        M = cv2.moments(contour)
        if M['m00'] != 0:
            cx = int(M['m10'] / M['m00'])
            cy = int(M['m01'] / M['m00'])
            cv2.putText(vis_img, str(item['id']), (cx, cy),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA)#序号的大小与线宽度

        # # 替代 moments 重心方式，使用外接矩形中心
        # x, y, w, h = cv2.boundingRect(contour)
        # cx = x + w // 2
        # cy = y + h // 2
        # cv2.putText(vis_img, str(item['id']), (cx, cy),
        #             cv2.FONT_HERSHEY_SIMPLEX, 3, color, 2, cv2.LINE_AA)

        # # 创建该块的 instance mask，用于质心计算
        # instance_mask = np.zeros((h, w), dtype=np.uint8)
        # cv2.drawContours(instance_mask, [contour], -1, 1, -1)
        #
        # # 计算质心（内部像素点的几何中心）
        # ys, xs = np.where(instance_mask > 0)
        # if len(xs) > 0 and len(ys) > 0:
        #     cx = int(np.mean(xs))
        #     cy = int(np.mean(ys))
        #     cv2.putText(vis_img, str(item['id']), (cx, cy),
        #                 cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2, cv2.LINE_AA)

    cv2.imwrite(save_path, vis_img)
    print(f"✅ 可视化图保存至: {save_path}")

def color_map_to_label(seg_img, palette):
    """
    将彩色语义图映射为类别ID图

    参数：
        seg_img: HxWx3 彩色图像（例如读取的png）
        palette: 类别调色板，如 [[0,0,0], [139,69,19], ...]

    返回：
        label_map: HxW 图像，每个像素为类别ID（int）
    """
    h, w, _ = seg_img.shape
    label_map = np.zeros((h, w), dtype=np.uint8)

    # 将调色板转换为字典
    color2label = {tuple(color): idx for idx, color in enumerate(palette)}

    # 对每个像素映射
    for color, label in color2label.items():
        mask = np.all(seg_img == color, axis=-1)
        label_map[mask] = label

    return label_map

# 定义类别调色板和类别名
PALETTE = [
    [0, 0, 0],         # background
    [139, 69, 19],     # barren
    [0, 255, 0],       # forest
    [255, 255, 0],     # farmland
    [0, 0, 255],       # water
    [128, 128, 128],   # road
    [0, 255, 255]      # building
]
CLASS_NAMES = {
    0: 'background',
    1: 'barren',
    2: 'forest',
    3: 'farmland',
    4: 'water',
    5: 'road',
    6: 'building'
}


def draw_json_boundary(pic_url):


    # 读取彩色分割图
    seg_color = cv2.imread(pic_url)  # 注意是 BGR 格式
    # seg_color = cv2.imread(r"D:\project\UAV_model\prediction_results\unetformer_UAV_6000X4000\patch_0028.png")  # 注意是 BGR 格式
    # seg_color = cv2.imread(r"J:\uhr\pytools\numer_coordinate\patch_0047_rgb.png")  # 注意是 BGR 格式
    seg_color = cv2.cvtColor(seg_color, cv2.COLOR_BGR2RGB)

    # 转为类别索引图
    label_map = color_map_to_label(seg_color, PALETTE)

    # 继续后续处理
    results = segment_instance_and_contour(label_map, CLASS_NAMES)
    dir_name=os.path.dirname(pic_url)
    final_vis_png_name=os.path.basename(pic_url)+"final_vis.png"
    final_vis_png_path=os.path.join(dir_name,final_vis_png_name) #输出大本地的只有边缘线条

    instance_results_json_name=os.path.basename(pic_url)+"instance_results.json"
    instance_results_json_path=os.path.join(dir_name,instance_results_json_name)  # 输出到本地的，分割的边缘的json文件
    draw_instances_with_ids(label_map, results, save_path=final_vis_png_path)



    # 保存为 JSON（可选）
    import json
    with open(instance_results_json_path, "w") as f:
        json.dump(results, f, indent=2)

    return final_vis_png_path,instance_results_json_path