import os
import cv2
import numpy as np
from collections import Counter, defaultdict
from PIL import Image
from Tool_Classes_And_Palette import Annotate_CLASSES, Annotate_PALETTE, bg_PALETTE

# ※ 需要修改输入输出路径 ※ #
input_dir = 'ORI_GT_label_fold'
output_dir_1 = 'Data/labels/train'
output_dir_2 = 'Data/labels/val'
output_GT_dir_1 = 'Data/labels_GT/train'
output_GT_dir_2 = 'Data/labels_GT/val'

# 1. --- 您提供的类别和调色板 ---
# (1) 从 Annotate_CLASSES (元组) 中移除 '背景'
if Annotate_CLASSES and Annotate_CLASSES[-1] == '背景':
    Annotate_CLASSES = Annotate_CLASSES[:-1]
# (2) 从 Annotate_PALETTE (列表) 中移除背景颜色 [0,0,0]
if Annotate_PALETTE and (Annotate_PALETTE[-1] == [0, 0, 0] or Annotate_PALETTE[-1] == (0, 0, 0)):
    Annotate_PALETTE.pop()
# (3) 确保调色板中的所有颜色都是元组 (Tuple)，以便后续查找
Annotate_PALETTE = [tuple(color) for color in Annotate_PALETTE]
# (4) 确保 bg_PALETTE 是元组，以匹配 Counter 中的键类型
bg_PALETTE = tuple(bg_PALETTE) 
# (5) 检查类别和调色板长度是否一致
if len(Annotate_CLASSES) != len(Annotate_PALETTE):
    print(f"[警告] 移除背景后，类别 ({len(Annotate_CLASSES)}) 和调色板 ({len(Annotate_PALETTE)}) 长度不一致!")
else:
    print(f"--- 成功移除背景项，剩余 {len(Annotate_CLASSES)} 个有效类别。 ---")

os.makedirs(output_dir_1, exist_ok=True)
os.makedirs(output_dir_2, exist_ok=True)
os.makedirs(output_GT_dir_1, exist_ok=True)
os.makedirs(output_GT_dir_2, exist_ok=True)

# 全局统计颜色频率与 class 像素频率
global_color_counter = Counter()
global_class_counter = Counter()
color_class_counter = defaultdict(int)  # (R,G,B) → count
color_to_old_class = {}  # (R,G,B) → class_id
remap_class_dict = {}  # old_class_id → new_class_id

# 自动提取颜色映射（跳过背景）
def extract_color_mapping(img_path):
    img = Image.open(img_path).convert('RGB')
    pixels = list(img.getdata())
    counter = Counter(pixels)
    color_map = {}
    for color, count in counter.items():
        global_color_counter[color] += count
        if color != (0, 0, 0) and color[0] == color[1] == color[2]:
            class_id = color[0] - 1
            if class_id >= 0:
                color_map[color] = class_id
                global_class_counter[class_id] += count
                color_class_counter[color] += count
                color_to_old_class[color] = class_id
    return color_map

# --- 修改：函数签名增加了 gt_save_paths ---
def process_image(img_path, txt_save_paths, gt_save_paths):
    color_to_class = extract_color_mapping(img_path)
    if not color_to_class:
        print(f"[跳过] {os.path.basename(img_path)} 无有效目标")
        return

    img = cv2.imread(img_path)
    h, w = img.shape[:2]
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # --- 新增：创建新的 GT 掩码 (灰度图)，0 为背景 ---
    new_gt_mask = np.zeros((h, w), dtype=np.uint8)
    lines = []

    for rgb, old_class_id in color_to_class.items():
        # --- 修改：获取 new_class_id 并填充 new_gt_mask ---
        new_class_id = remap_class_dict.get(old_class_id, old_class_id)
        
        # 创建布尔掩码
        mask_bool = np.all(img_rgb == rgb, axis=-1)
        
        # --- 新增：在 new_gt_mask 上填充新的 class_id
        # (使用 new_class_id + 1，因为 0 是背景)
        new_gt_mask[mask_bool] = new_class_id + 1

        # --- 修改：基于布尔掩码创建 255 掩码用于 findContours ---
        mask_255 = mask_bool.astype(np.uint8) * 255
        
        contours, _ = cv2.findContours(mask_255, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        for contour in contours:
            if contour.shape[0] < 3:
                continue
            norm_pts = contour.squeeze(1).astype(np.float32)
            norm_pts[:, 0] /= w
            norm_pts[:, 1] /= h
            flat = norm_pts.flatten()
            
            # --- 修改：此处 new_class_id 已在循环外层定义 ---
            line = f"{new_class_id} " + " ".join([f"{x:.6f}" for x in flat])
            lines.append(line)

    # --- 新增：保存处理后的 GT 掩码图像 ---
    # 只要存在有效类别 (color_to_class 非空)，就保存 GT 掩码
    for save_path in gt_save_paths:
        try:
            # 使用 PIL 保存灰度 PNG 图像
            Image.fromarray(new_gt_mask).save(save_path)
            print(f"[✔] GT 掩码保存成功: {os.path.basename(save_path)}")
        except Exception as e:
            print(f"[✘] GT 掩码保存失败: {os.path.basename(save_path)} - {e}")

    # 保存 .txt 标签文件 (仅当找到轮廓时)
    if lines:
        for save_path in txt_save_paths:
            with open(save_path, 'w') as f:
                for line in lines:
                    f.write(line + "\n")
            print(f"[✔] 转换成功: {os.path.basename(save_path)}，共 {len(lines)} 个实例")
    else:
        # 即使没有轮廓，GT 掩码也已保存
        print(f"[⚠] {os.path.basename(img_path)} 没有轮廓 (但 GT 掩码已保存)")

# 第一次遍历图像，仅提取颜色信息用于构建 class 映射
for fname in os.listdir(input_dir):
    if fname.lower().endswith('.png'):
        img_path = os.path.join(input_dir, fname)
        extract_color_mapping(img_path)

# 构建 class_id 重映射表（按像素数从大到小排序）
class_pixel_count = {color_to_old_class[c]: count for c, count in color_class_counter.items()}
sorted_classes = sorted(class_pixel_count.items(), key=lambda x: x[1], reverse=True)
remap_class_dict = {old_cls: new_idx for new_idx, (old_cls, _) in enumerate(sorted_classes)}

# 第二次遍历图像，正式处理并生成标签
for fname in os.listdir(input_dir):
    if fname.lower().endswith('.png'):
        img_path = os.path.join(input_dir, fname)
        base_name = os.path.splitext(fname)[0]

        txt_path_1 = os.path.join(output_dir_1, base_name + '.txt')
        txt_path_2 = os.path.join(output_dir_2, base_name + '.txt')

        # --- 新增：定义 GT 掩码输出路径 (保存为 .png) ---
        gt_path_1 = os.path.join(output_GT_dir_1, base_name + '.png')
        gt_path_2 = os.path.join(output_GT_dir_2, base_name + '.png')

        process_image(img_path, [txt_path_1, txt_path_2], [gt_path_1, gt_path_2])

# 打印颜色统计和类别映射
print("\n📊 所有图像颜色统计：")

def get_label_info(class_id):
    """辅助函数：根据 class_id 获取标签名和颜色"""
    if 0 <= class_id < len(Annotate_CLASSES) and 0 <= class_id < len(Annotate_PALETTE):
        return Annotate_CLASSES[class_id], Annotate_PALETTE[class_id]
    else:
        return "未知标签", (255, 255, 255) # 返回一个默认值

for color, count in global_color_counter.most_common():
    if color == bg_PALETTE: # 使用 bg_PALETTE 变量
        print(f"背景颜色 {color} 出现次数: {count}")
    elif color[0] == color[1] == color[2]:
        class_id = color[0] - 1
        label_name, label_color = get_label_info(class_id)
        print(f"颜色 {color} → class {class_id} (标签: '{label_name}', 颜色: {label_color})，出现次数: {count}")
    else:
        # 此情况理论上不应出现，因为 extract_color_mapping 已过滤
        print(f"[⚠] 非灰阶颜色 {color}，出现次数: {count} (此颜色不应被处理)")

# 打印详细的类别统计 
print("\n✅ 有效类别统计（原 class_id → 新 class_id → 总像素数）：")
# 按照 new_id (新类别索引) 排序输出
sorted_remap = sorted(remap_class_dict.items(), key=lambda item: item[1])
for old_id, new_id in sorted_remap:
    label_name, label_color = get_label_info(old_id)
    print(f"class {old_id} (标签: '{label_name}', 颜色: {label_color}) → {new_id}: {class_pixel_count[old_id]} pixels")

# 4. 额外输出新 class_id 到颜色的映射 ※※※
print(f"\n🎨 新 class_id 与标签颜色 (Annotate_PALETTE) 映射表：")
# 1. 创建一个 new_id -> color 的映射
new_id_to_color = {}
for old_id, new_id in remap_class_dict.items():
    if 0 <= old_id < len(Annotate_PALETTE):
        # 从 Annotate_PALETTE 获取原始颜色（它是一个元组）
        color_tuple = Annotate_PALETTE[old_id]
        # 转换为列表 [R, G, B] 以匹配您要的格式
        new_id_to_color[new_id] = list(color_tuple)
    else:
        # 预防性代码，以防 old_id 超出范围
        new_id_to_color[new_id] = [-1, -1, -1] # 表示错误/未找到

# 2. 按照 new_id (0, 1, 2...) 排序并按指定格式打印
sorted_new_ids = sorted(new_id_to_color.keys())
for new_id in sorted_new_ids:
    color_list = new_id_to_color[new_id]
    # 格式化输出：  {id}: {color_list} 
    # (注意：颜色列表的格式会自然包含逗号和空格)
    print(f"  {new_id}: {color_list}")

print(f"\n✅ 全部图像处理完毕。")
print(f"   标签 (.txt) 输出目录: {output_dir_1}, {output_dir_2}")
print(f"   GT 掩码 (.png) 输出目录: {output_GT_dir_1}, {output_GT_dir_2}")
