#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
脚本1：从源目录挑选图片，处理图片并生成 YOLO 标签

图片处理：原图居中保持原样，四周补白，画布尺寸变为原图的 SCALE 倍
标注格式：class_id center_x center_y width height（精确反映原图在画布中的位置）

配置说明：
- DECIPHERED_ROOT: HUST-OBC/deciphered 路径
- VAL_ROOT: HUST-OBC/GuoXueDaShi_1390 路径（验证集来源）
- OUTPUT_ROOT: 输出目录
- TRAIN_MAX: 训练集每类最多多少张
- TEST_MAX: 测试集每类多少张
- VAL_MAX: 验证集每类最多多少张（None表示不限制）
- CLASS_MAX: 最终取多少个类别
- MIN_TRAIN_IMAGES: 训练图片少于该数量的类别直接跳过
- SCALE: 画布相对于原图的放大倍数
- BG_COLOR: 补白背景色
"""

import os
import shutil
import json
from pathlib import Path
from PIL import Image

# ==================== 配置区域 ====================
DECIPHERED_ROOT = Path('/Users/HonestQiao/Downloads/HUST-OBS/data/HUST-OBC/deciphered')
VAL_ROOT = Path('/Users/HonestQiao/Downloads/HUST-OBS/data/HUST-OBC/GuoXueDaShi_1390')
OUTPUT_ROOT = Path('/Users/HonestQiao/Downloads/HUST-OBS/data/HUST-OBC-Selected')

TRAIN_MAX = 30   # 训练集每类最多几张
TEST_MAX = 1     # 测试集每类最多几张
VAL_MAX = 2      # 验证集每类最多几张（设为 None 则不限制）
CLASS_MAX = 10   # 最终取多少类的图片
MIN_TRAIN_IMAGES = 20  # 训练图片少于该数量的类别直接跳过

SCALE = 2        # 画布相对于原图的放大倍数
BG_COLOR = (255, 255, 255)  # 白色背景

ID_JSON = Path('/Users/HonestQiao/Downloads/HUST-OBS/data/HUST-OBC/deciphered/ID_to_chinese.json')
# =================================================

IMAGE_EXTS = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp'}


def is_image(filename):
    return Path(filename).suffix.lower() in IMAGE_EXTS


def collect_leaf_dirs(root):
    """收集所有叶子目录（最深层目录，不含子目录的目录）"""
    leaf_dirs = []
    for dirpath, dirnames, filenames in os.walk(root):
        if not dirnames:
            leaf_dirs.append(dirpath)
    return leaf_dirs


def get_allowed_classes():
    """根据 CLASS_MAX 配置和训练图片数量过滤，返回允许处理的类别集合"""
    leaf_dirs = collect_leaf_dirs(DECIPHERED_ROOT)

    # 收集验证集中存在的类别
    val_classes = set()
    if VAL_ROOT.exists():
        for d in VAL_ROOT.iterdir():
            if d.is_dir():
                val_classes.add(d.name)

    # 统计每个有效类别的训练图片数量
    class_counts = []
    for leaf in sorted(leaf_dirs):
        class_id = Path(leaf).name
        if not (len(class_id) == 4 and class_id.isdigit()):
            continue

        # 只保留验证集中也存在的类别
        if class_id not in val_classes:
            continue

        images = [f for f in Path(leaf).iterdir() if f.is_file() and is_image(f.name)]
        count = len(images)

        # 跳过训练图片少于 MIN_TRAIN_IMAGES 的类别
        if count < MIN_TRAIN_IMAGES:
            continue

        class_counts.append((class_id, count))

    # 按类别ID排序后取前 CLASS_MAX 个
    class_counts.sort(key=lambda x: x[0])
    allowed = set([c[0] for c in class_counts[:CLASS_MAX]])

    print(f"验证集有 {len(val_classes)} 个类别")
    print(f"训练图片>={MIN_TRAIN_IMAGES} 且验证集存在的类别: {len(class_counts)} 个")
    print(f"本次处理前 {len(allowed)} 个类别")
    return allowed


def process_image(src_path, dst_path):
    """处理单张图片：原图居中保持原样，四周补白，画布变为 SCALE 倍"""
    with Image.open(src_path) as img:
        # 转换为 RGB（处理 PNG 透明通道等）
        if img.mode in ('RGBA', 'P'):
            img = img.convert('RGBA')
            background = Image.new('RGBA', img.size, (*BG_COLOR, 255))
            background.paste(img, mask=img.split()[3] if img.mode == 'RGBA' else None)
            img = background.convert('RGB')
        elif img.mode != 'RGB':
            img = img.convert('RGB')

        orig_w, orig_h = img.size
        canvas_w = orig_w * SCALE
        canvas_h = orig_h * SCALE

        # 创建白色背景画布
        canvas = Image.new('RGB', (canvas_w, canvas_h), BG_COLOR)

        # 计算居中位置
        paste_x = (canvas_w - orig_w) // 2
        paste_y = (canvas_h - orig_h) // 2

        # 粘贴原图（保持原样）
        canvas.paste(img, (paste_x, paste_y))

        # 保存
        canvas.save(dst_path, quality=95)

        return orig_w, orig_h, canvas_w, canvas_h


def write_label(label_path, class_idx, orig_w, orig_h, canvas_w, canvas_h):
    """生成精确的 YOLO 标签：原图在画布中的相对位置"""
    # 中心点（画布坐标系中）
    center_x = 0.5  # 居中
    center_y = 0.5  # 居中
    # 原图宽高占画布的比例
    box_w = orig_w / canvas_w
    box_h = orig_h / canvas_h

    label_content = f"{class_idx} {center_x} {center_y} {box_w} {box_h}\n"
    with open(label_path, 'w', encoding='utf-8') as f:
        f.write(label_content)


def process_deciphered(allowed_classes, id_to_class):
    """处理训练集和测试集"""
    leaf_dirs = collect_leaf_dirs(DECIPHERED_ROOT)
    print(f"发现 {len(leaf_dirs)} 个叶子目录")

    for leaf in sorted(leaf_dirs):
        leaf_path = Path(leaf)
        class_id = leaf_path.name

        if not (len(class_id) == 4 and class_id.isdigit()):
            continue
        if class_id not in allowed_classes:
            continue

        images = sorted([f for f in leaf_path.iterdir() if f.is_file() and is_image(f.name)])
        total = len(images)
        if total == 0:
            continue

        class_idx = id_to_class[class_id]

        # 创建目录
        img_train_dir = OUTPUT_ROOT / 'images' / 'train' / class_id
        img_test_dir = OUTPUT_ROOT / 'images' / 'test' / class_id
        lbl_train_dir = OUTPUT_ROOT / 'labels' / 'train' / class_id
        lbl_test_dir = OUTPUT_ROOT / 'labels' / 'test' / class_id
        for d in [img_train_dir, img_test_dir, lbl_train_dir, lbl_test_dir]:
            d.mkdir(parents=True, exist_ok=True)

        # 训练集
        train_images = images[:TRAIN_MAX]
        for img in train_images:
            new_name = f"{class_id}-{img.stem}.jpg"
            new_stem = f"{class_id}-{img.stem}"
            orig_w, orig_h, canvas_w, canvas_h = process_image(img, img_train_dir / new_name)
            write_label(lbl_train_dir / f"{new_stem}.txt",
                        class_idx, orig_w, orig_h, canvas_w, canvas_h)

        # 测试集
        if total > TRAIN_MAX:
            test_images = images[TRAIN_MAX:TRAIN_MAX + TEST_MAX]
        else:
            test_images = images[-TEST_MAX:] if total >= TEST_MAX else images

        for img in test_images:
            new_name = f"{class_id}-{img.stem}.jpg"
            new_stem = f"{class_id}-{img.stem}"
            orig_w, orig_h, canvas_w, canvas_h = process_image(img, img_test_dir / new_name)
            write_label(lbl_test_dir / f"{new_stem}.txt",
                        class_idx, orig_w, orig_h, canvas_w, canvas_h)

        print(f"[{class_id}] 总计: {total:3d} | train: {len(train_images):2d} | test: {len(test_images):2d}")


def process_val(allowed_classes, id_to_class):
    """处理验证集"""
    if not VAL_ROOT.exists():
        print(f"验证集源目录不存在: {VAL_ROOT}")
        return

    img_val_dir = OUTPUT_ROOT / 'images' / 'val'
    lbl_val_dir = OUTPUT_ROOT / 'labels' / 'val'

    copied_dirs = 0
    copied_files = 0

    for subdir in sorted(VAL_ROOT.iterdir()):
        if not subdir.is_dir():
            continue

        class_id = subdir.name
        if class_id not in allowed_classes:
            continue

        class_idx = id_to_class[class_id]
        dst_img_subdir = img_val_dir / class_id
        dst_lbl_subdir = lbl_val_dir / class_id
        dst_img_subdir.mkdir(parents=True, exist_ok=True)
        dst_lbl_subdir.mkdir(parents=True, exist_ok=True)

        images = sorted([f for f in subdir.iterdir() if f.is_file() and is_image(f.name)])
        if VAL_MAX is not None:
            images = images[:VAL_MAX]

        for img in images:
            new_name = f"{class_id}-{img.stem}.jpg"
            new_stem = f"{class_id}-{img.stem}"
            orig_w, orig_h, canvas_w, canvas_h = process_image(img, dst_img_subdir / new_name)
            write_label(dst_lbl_subdir / f"{new_stem}.txt",
                        class_idx, orig_w, orig_h, canvas_w, canvas_h)

        copied_dirs += 1
        copied_files += len(images)

    print(f"验证集复制完成: {copied_dirs} 个子目录, {copied_files} 个文件")


def generate_data_yaml(id_to_class, id_to_name):
    """生成 data.yaml"""
    # 按 class_idx 排序
    names = {}
    for oid, idx in id_to_class.items():
        names[idx] = id_to_name[oid]

    yaml_lines = [
        f"path: {OUTPUT_ROOT}",
        "train: images/train",
        "val: images/val",
        "test: images/test",
        "names:"
    ]
    for i in range(len(names)):
        yaml_lines.append(f"  {i}: {names[i]}")

    with open(OUTPUT_ROOT / 'data.yaml', 'w', encoding='utf-8') as f:
        f.write('\n'.join(yaml_lines) + '\n')
    print("data.yaml 已生成")


def generate_map_yaml(id_to_class, id_to_name):
    """生成 map.yaml，用于 serve_image.py 图片服务"""
    map_lines = ["names:"]
    for class_id in sorted(id_to_class.keys()):
        name = id_to_name[class_id]
        train_dir = OUTPUT_ROOT / 'images' / 'train' / class_id
        images = sorted([f for f in train_dir.iterdir() if f.is_file() and is_image(f.name)])
        if images:
            rel_path = f"images/train/{class_id}/{images[0].name}"
            map_lines.append(f"  {name}: {rel_path}")

    with open(OUTPUT_ROOT / 'map.yaml', 'w', encoding='utf-8') as f:
        f.write('\n'.join(map_lines) + '\n')
    print("map.yaml 已生成")


def main():
    # 清空输出目录
    if OUTPUT_ROOT.exists():
        shutil.rmtree(OUTPUT_ROOT)
    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)

    # 读取类别映射
    with open(ID_JSON, 'r', encoding='utf-8') as f:
        id_to_name = json.load(f)

    # 获取允许处理的类别
    allowed_classes = get_allowed_classes()

    # 建立 class_id -> class_idx 映射
    sorted_allowed = sorted(allowed_classes)
    id_to_class = {oid: idx for idx, oid in enumerate(sorted_allowed)}

    print("=" * 50)
    print("开始处理训练集和测试集...")
    print("=" * 50)
    process_deciphered(allowed_classes, id_to_class)

    print("\n" + "=" * 50)
    print("开始处理验证集...")
    print("=" * 50)
    process_val(allowed_classes, id_to_class)

    print("\n" + "=" * 50)
    print("生成 data.yaml...")
    print("=" * 50)
    generate_data_yaml(id_to_class, id_to_name)

    print("\n" + "=" * 50)
    print("生成 map.yaml...")
    print("=" * 50)
    generate_map_yaml(id_to_class, id_to_name)

    print("\n" + "=" * 50)
    print("全部完成！")
    print("=" * 50)


if __name__ == '__main__':
    main()
