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'

# ----------------------------------------------------
# ※※※ 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)

# 全局统计颜色频率与 class 像素频率
global_color_counter = Counter()
global_class_counter = Counter() # 将用于统计（ori_label 模式）
color_class_counter = defaultdict(int)  # (R,G,B) → count
color_to_old_class = {}  # (R,G,B) → class_id
# *** ori_label 模式: 移除了 remap_class_dict ***

# 自动提取颜色映射（跳过背景）
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 != bg_PALETTE and color[0] == color[1] == color[2]: # 使用元组 bg_PALETTE
            class_id = color[0] - 1
            if class_id >= 0:
                color_map[color] = class_id
                global_class_counter[class_id] += count # 使用 global_class_counter 统计
                color_class_counter[color] += count
                color_to_old_class[color] = class_id
    return color_map

# 处理单张图片
def process_image(img_path, save_path_list):
    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)

    lines = []
    for rgb, old_class_id in color_to_class.items():
        mask = np.all(img_rgb == rgb, axis=-1).astype(np.uint8) * 255
        contours, _ = cv2.findContours(mask, 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()
            
            # *** ori_label 模式: 直接使用 old_class_id ***
            line = f"{old_class_id} " + " ".join([f"{x:.6f}" for x in flat])
            lines.append(line)

    if lines:
        for save_path in save_path_list:
            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:
        print(f"[⚠] {os.path.basename(img_path)} 没有轮廓")

# 第一次遍历图像，仅提取颜色信息用于构建 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)

# *** ori_label 模式: 移除了重映射 (remap) 逻辑块 ***
# (保留 class_pixel_count 用于统计)
class_pixel_count = {color_to_old_class[c]: count for c, count in color_class_counter.items()}

# 第二次遍历图像，正式处理并生成标签
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')
        process_image(img_path, [txt_path_1, txt_path_2])

# ----------------------------------------------------
# ※※※ 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} (此颜色不应被处理)")

# ----------------------------------------------------
# ※※※ 3. 修改: 打印详细的类别统计 (ori_label 模式) ※※※
# ----------------------------------------------------
print("\n✅ 有效类别统计（按 原 class_id 排序 → 总像素数）：")
# 按照 old_id (原类别索引) 排序输出
sorted_classes = sorted(class_pixel_count.items(), key=lambda item: item[0])

for old_id, pixel_count in sorted_classes:
    label_name, label_color = get_label_info(old_id)
    print(f"class {old_id} (标签: '{label_name}', 颜色: {label_color}): {pixel_count} pixels")

# ----------------------------------------------------
# ※※※ 4. 修改: 额外输出 原 class_id 到颜色的映射 (ori_label 模式) ※※※
# ----------------------------------------------------
print(f"\n🎨 找到的 原 class_id 与标签颜色 (Annotate_PALETTE) 映射表：")

# 1. 按照 old_id (0, 1, 2...) 排序并按指定格式打印
#    我们使用 class_pixel_count.keys() 来获取所有实际找到的 old_id
sorted_old_ids = sorted(class_pixel_count.keys())

for old_id in sorted_old_ids:
    if 0 <= old_id < len(Annotate_PALETTE):
        color_tuple = Annotate_PALETTE[old_id]
        color_list = list(color_tuple)
        print(f"  {old_id}: {color_list}")
    else:
        # 预防性代码
        print(f"  {old_id}: [颜色未在调色板中定义]")
# ----------------------------------------------------
# ※※※ 修改结束 ※※※
# ----------------------------------------------------

print(f"\n✅ 全部图像处理完毕。标签输出目录：{output_dir_1}、{output_dir_2}")