"""
将标注的 55 组样本转换为标准 YOLO 目标检测数据集
生成:
  tcaptcha/data/yolo/images/train/
  tcaptcha/data/yolo/images/val/
  tcaptcha/data/yolo/labels/train/
  tcaptcha/data/yolo/labels/val/
  tcaptcha/data/yolo/dataset.yaml
"""

import json
import logging
import random
import shutil
from pathlib import Path
import cv2
import numpy as np

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("prepare_yolo")

RAW_DIR = Path("tcaptcha/data/raw")
LABEL_FILE = Path("tcaptcha/data/labels.jsonl")
YOLO_DIR = Path("tcaptcha/data/yolo")

IMG_W = 672
IMG_H = 480
BOX_SIZE = 54  # 图标大致大小 (约 54x54 像素)


def find_all_icons_in_bg(bg_path: Path, labeled_points: list) -> list:
    """
    结合人工标注点与轮廓补充检测，找到大图中所有图标（包含 3 个目标 + 2~3 个干扰图标）
    返回: [(cx, cy, w, h), ...] 绝对像素坐标
    """
    img = cv2.imread(str(bg_path))
    if img is None:
        return []
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 提取暗色轮廓找干扰图标
    edges = cv2.Canny(gray, 40, 120)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (9, 9))
    closed = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
    contours, _ = cv2.findContours(closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    all_boxes = []
    # 1. 先加入人工标注的 3 个目标（绝对精准）
    for p in labeled_points:
        cx, cy = p["x"], p["y"]
        all_boxes.append((cx, cy, BOX_SIZE, BOX_SIZE))

    # 2. 补充检测未标注的干扰图标
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if 26 <= w <= 85 and 26 <= h <= 85 and (w * h) >= 650:
            cx, cy = x + w // 2, y + h // 2
            # 必须远离已有图标
            if all(np.hypot(cx - lx, cy - ly) > 42 for lx, ly, _, _ in all_boxes):
                # 限制尺寸
                bw = min(68, max(46, w + 8))
                bh = min(68, max(46, h + 8))
                all_boxes.append((cx, cy, bw, bh))

    return all_boxes


def prepare_dataset(val_ratio: float = 0.2, seed: int = 42):
    random.seed(seed)
    if not LABEL_FILE.exists():
        logger.error("标签文件不存在: %s", LABEL_FILE)
        return

    # 清理并创建目录
    if YOLO_DIR.exists():
        shutil.rmtree(YOLO_DIR)

    for split in ["train", "val"]:
        (YOLO_DIR / "images" / split).mkdir(parents=True, exist_ok=True)
        (YOLO_DIR / "labels" / split).mkdir(parents=True, exist_ok=True)

    with open(LABEL_FILE, "r", encoding="utf-8") as f:
        records = [json.loads(line) for line in f if line.strip()]

    random.shuffle(records)
    val_count = max(5, int(len(records) * val_ratio))
    train_records = records[val_count:]
    val_records = records[:val_count]

    logger.info("总标注样本: %d (训练集: %d, 验证集: %d)", len(records), len(train_records), len(val_records))

    def process_split(record_list, split_name):
        total_boxes = 0
        for r in record_list:
            sid = r["sample_id"]
            src_img = RAW_DIR / f"{sid}_bg.png"
            if not src_img.exists():
                continue

            dst_img = YOLO_DIR / "images" / split_name / f"{sid}.png"
            shutil.copy(src_img, dst_img)

            # 提取全部图标坐标
            boxes = find_all_icons_in_bg(src_img, r["points"])
            total_boxes += len(boxes)

            # 写入 YOLO 格式标注: class_id x_center y_center width height (归一化到 0~1)
            dst_label = YOLO_DIR / "labels" / split_name / f"{sid}.txt"
            lines = []
            for cx, cy, bw, bh in boxes:
                # 边界截断保护
                x1 = max(0, cx - bw / 2)
                y1 = max(0, cy - bh / 2)
                x2 = min(IMG_W, cx + bw / 2)
                y2 = min(IMG_H, cy + bh / 2)

                norm_cx = (x1 + x2) / 2.0 / IMG_W
                norm_cy = (y1 + y2) / 2.0 / IMG_H
                norm_w = (x2 - x1) / IMG_W
                norm_h = (y2 - y1) / IMG_H

                lines.append(f"0 {norm_cx:.6f} {norm_cy:.6f} {norm_w:.6f} {norm_h:.6f}")

            dst_label.write_text("\n".join(lines), encoding="utf-8")
        logger.info("  [%s] 处理完成: %d 张图片, 共 %d 个图标标注", split_name, len(record_list), total_boxes)

    process_split(train_records, "train")
    process_split(val_records, "val")

    # 生成 dataset.yaml
    yaml_content = f"""path: {YOLO_DIR.resolve()}
train: images/train
val: images/val

names:
  0: icon
"""
    yaml_path = YOLO_DIR / "dataset.yaml"
    yaml_path.write_text(yaml_content, encoding="utf-8")
    logger.info("✓ YOLO 数据集生成完毕: %s", yaml_path)


if __name__ == "__main__":
    prepare_dataset()