#!/usr/bin/env python
"""把 generate_boxes.py 的产出整理成 YOLO 训练所需的数据集，放到 data/ 下。

    uv run python scripts/prepare_yolo_dataset.py --source runs/boxes_gradcam

generate_boxes.py 的产出是**按类别**分目录的（便于人工核对）：

```
runs/boxes_gradcam/
├── images/<类别>/<图名>.png
└── labels/<类别>/<图名>.txt
```

而 YOLO 检测训练要的是 **train/val 划分**——类别信息在标签文件的第一列里，
不靠目录名。这一步做三件事：

1. 把按类别分的目录拍平，按 train/val 重新组织；
2. **按类别分层**切分，小类不至于在验证集里挂零；
3. 图片**统一成 PNG**：本来就是 PNG 的硬链接过去，不占额外磁盘；BMP 等其他格式
   （pearlescent / blue / uv 原始数据都是 BMP）无损转成 PNG。不缩放像素，
   归一化的框坐标原样有效。

4. 文件**按类别重新编号**：`<类别>_<序号>.png`，序号在类内按源路径排序，重跑结果一致。
   原文件名里的正反面、位置、时间戳等信息记在 `name_map.csv` 里，不会丢。
   想保留原名加 `--keep-names`。

```
data/<name>/
├── images/{train,val}/<类别>_<序号>.png
├── labels/{train,val}/<类别>_<序号>.txt
├── classes.txt      # 每行一个类别名，行号即 class_id，给标注工具复核用
├── name_map.csv     # 新名 ↔ 源文件对照
├── data.yaml
└── summary.json
```

没有框的图（标签为空）默认保留：YOLO 训练里它们是负样本，有助于压低误检。但**空标签的
语义是"这张图什么都没有"**，不是"有缺陷但没定位出来"。来自缺陷类别的空标签如果照搬，
等于在教模型抑制真实缺陷，所以可以用 `--empty-policy` 按来源类别分别处置：

    --empty-policy drop normal=keep black=full
    # 缺陷类没定位出来的直接丢弃；normal 留作负样本；整图属性的类退化成整图框
"""

from __future__ import annotations

import argparse
import csv
import errno
import json
import os
import random
import shutil
import sys
from collections import Counter, defaultdict
from collections.abc import Sequence
from pathlib import Path
from typing import NamedTuple

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import yaml  # noqa: E402

from datasets.yolo import IMAGE_SUFFIXES, write_classes_txt, write_data_yaml  # noqa: E402
from utils import setup_logging  # noqa: E402

logger = None

EMPTY_POLICIES = ("keep", "drop", "full")
EMPTY_POLICY_DESC = {
    "keep": "保留为负样本",
    "drop": "丢弃",
    "full": "退化成整图框",
}


PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
# PNG 能原样存下的像素模式；其余（CMYK、YCbCr 等）先转 RGB
PNG_MODES = frozenset({"1", "L", "LA", "I", "I;16", "P", "RGB", "RGBA"})


def is_png(path: Path) -> bool:
    """看文件头而不是后缀：后缀写着 .png、内容却是 BMP/JPEG 的也要转。"""
    with path.open("rb") as handle:
        return handle.read(len(PNG_SIGNATURE)) == PNG_SIGNATURE


def materialize_png(source: Path, destination: Path) -> str:
    """把图片落成 PNG：本来就是 PNG 的硬链接，其余格式无损转换。

    返回 "hard" / "copy" / "convert"。转换不缩放像素，归一化的框坐标原样有效。
    """
    if is_png(source):
        return link_or_copy(source, destination)

    from PIL import Image

    if destination.exists() or destination.is_symlink():
        destination.unlink()
    destination.parent.mkdir(parents=True, exist_ok=True)
    with Image.open(source) as image:
        if image.mode not in PNG_MODES:
            image = image.convert("RGB")
        image.save(destination, format="PNG")
    return "convert"


def link_or_copy(source: Path, destination: Path) -> str:
    """硬链接；跨文件系统时退回复制。"""
    if destination.exists() or destination.is_symlink():
        destination.unlink()
    destination.parent.mkdir(parents=True, exist_ok=True)
    try:
        os.link(source, destination)
        return "hard"
    except OSError as exc:
        if exc.errno != errno.EXDEV:
            raise
        shutil.copy2(source, destination)
        return "copy"


class Sample(NamedTuple):
    """一张图及其标签。

    Attributes:
        image: 图片路径。
        label: 标签路径。
        class_id: 标签里第一个框的类别下标；无框时为 None。
        group: 来源的类别目录名——也就是人工归类给这张图的类别。
            标签为空时，这是唯一还能知道"这张图本该是什么"的线索。
        full_box: 整图框的类别下标。非 None 时 _emit 会**重写**标签而不是链接原文件。
    """

    image: Path
    label: Path
    class_id: int | None
    group: str
    full_box: int | None = None


def first_class_id(label_path: Path) -> int | None:
    """取标签里第一个框的类别下标；空标签（无框）返回 None。"""
    for line in label_path.read_text(encoding="utf-8").splitlines():
        line = line.strip()
        if line:
            return int(float(line.split()[0]))
    return None


def collect(source: Path) -> list[Sample]:
    """配对 images/ 与 labels/。"""
    images_root, labels_root = source / "images", source / "labels"
    if not images_root.is_dir() or not labels_root.is_dir():
        raise FileNotFoundError(f"{source} 下缺少 images/ 或 labels/，先跑 scripts/generate_boxes.py")

    samples = []
    missing = []
    for image in sorted(p for p in images_root.rglob("*") if p.suffix.lower() in IMAGE_SUFFIXES):
        relative = image.relative_to(images_root)
        label = labels_root / relative.with_suffix(".txt")
        if not label.exists():
            missing.append(image)
            continue
        group = relative.parent.name  # 顶层就是图片时为 ""
        samples.append(Sample(image, label, first_class_id(label), group))
    if missing:
        raise FileNotFoundError(
            f"有 {len(missing)} 张图没有对应的标签，例如 {missing[0]}；"
            "产出目录可能不完整，重新跑一次 generate_boxes.py"
        )
    if not samples:
        raise ValueError(f"{source} 下没有找到任何图片")
    return samples


def class_of(sample: Sample, names: Sequence[str]) -> str:
    """样本在输出里归属的类别名：优先看标签内容，空标签才退回来源目录名。

    以标签为准，重命名后的文件名就永远不会与标签里的 class_id 矛盾。
    """
    if sample.class_id is not None:
        return names[sample.class_id]
    if sample.full_box is not None:
        return names[sample.full_box]
    return sample.group or "background"


def assign_names(samples: Sequence[Sample], names: Sequence[str]) -> dict[Path, str]:
    """按类别重新编号，返回 {源图片: 新文件名主干}，如 `blackdot_0042`。

    序号在类内按源路径排序、与 train/val 划分无关：换 --val-ratio 或 --seed 重跑，
    同一张图拿到的名字不变，已经做过的人工复核记录还能对得上。
    """
    by_class: dict[str, list[Sample]] = defaultdict(list)
    for sample in samples:
        by_class[class_of(sample, names)].append(sample)
    width = max(4, len(str(max((len(v) for v in by_class.values()), default=0))))
    mapping = {}
    for name, items in by_class.items():
        for number, sample in enumerate(sorted(items, key=lambda s: str(s.image)), start=1):
            mapping[sample.image] = f"{name}_{number:0{width}d}"
    return mapping


def parse_empty_policy(tokens: Sequence[str]) -> tuple[str, dict[str, str]]:
    """解析 --empty-policy：一个可选的全局默认，加若干 `类别=策略`。

        --empty-policy drop normal=keep black=full
        -> ("drop", {"normal": "keep", "black": "full"})
    """
    default, per_class = "keep", {}
    for token in tokens:
        if "=" in token:
            name, _, policy = token.partition("=")
            if policy not in EMPTY_POLICIES:
                raise ValueError(f"未知策略 {policy!r}（类别 {name}）；可选 {EMPTY_POLICIES}")
            per_class[name] = policy
        elif token in EMPTY_POLICIES:
            default = token
        else:
            raise ValueError(
                f"看不懂 --empty-policy 的 {token!r}；"
                f"要么是全局默认（{'/'.join(EMPTY_POLICIES)}），要么是 `类别=策略`"
            )
    return default, per_class


def apply_empty_policy(
    samples: Sequence[Sample], default: str, per_class: dict[str, str], names: Sequence[str]
) -> tuple[list[Sample], dict[str, tuple[str, int]]]:
    """按来源类别处置空标签，返回 (保留下来的样本, {类别: (策略, 数量)})。

    空标签在 YOLO 里的语义是"这张图什么都没有"——一个明确的负样本断言，而不是
    "有缺陷但没定位出来"。所以来自缺陷类别的空标签如果照搬，等于在教模型抑制真实缺陷。
    """
    index = {name: i for i, name in enumerate(names)}
    unknown = {
        name for name in per_class if name not in index
    }
    if unknown:
        raise ValueError(
            f"--empty-policy 里的类别 {sorted(unknown)} 不在 data.yaml 中；可选 {list(names)}"
        )

    kept: list[Sample] = []
    stats: dict[str, tuple[str, int]] = {}
    for sample in samples:
        if sample.class_id is not None:
            kept.append(sample)
            continue
        policy = per_class.get(sample.group, default)
        counted = stats.get(sample.group, (policy, 0))
        stats[sample.group] = (policy, counted[1] + 1)
        if policy == "drop":
            continue
        if policy == "full":
            if sample.group not in index:
                raise ValueError(
                    f"{sample.group!r} 要用整图框，但它不在 data.yaml 的类别里；可选 {list(names)}"
                )
            kept.append(sample._replace(full_box=index[sample.group]))
        else:
            kept.append(sample)
    return kept, stats


def main() -> int:
    global logger
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--source", required=True, help="generate_boxes.py 的产出目录")
    parser.add_argument("--out", default=None, help="输出目录；缺省 data/<source 目录名>")
    parser.add_argument("--data-root", default="data", help="自动推导时的根目录，默认 data")
    parser.add_argument("--val-ratio", type=float, default=0.2)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--empty-policy", nargs="*", default=[], metavar="POLICY|类别=POLICY",
        help=(
            "怎么处置没有框的图（标签为空）。可写一个全局默认，再按类别覆盖，如 "
            "`--empty-policy drop normal=keep black=full`。"
            "keep=当负样本保留（默认）、drop=丢弃、full=退化成整图框。"
            "注意空标签在 YOLO 里的意思是『此图无缺陷』，缺陷类的空标签照搬会教模型抑制真实缺陷"
        ),
    )
    parser.add_argument(
        "--keep-names", action="store_true",
        help="保留源文件名，不按 <类别>_<序号> 重新编号（源文件名需去掉后缀后全局唯一）",
    )
    parser.add_argument("--overwrite", action="store_true", help="输出目录已存在时先清空")
    parser.add_argument("--dry-run", action="store_true", help="只打印计划")
    args = parser.parse_args()

    source = Path(args.source)
    out_dir = Path(args.out) if args.out else Path(args.data_root) / source.name
    logger = setup_logging()

    try:
        samples = collect(source)
    except (FileNotFoundError, ValueError) as exc:
        logger.error("%s", exc)
        return 1

    meta_path = source / "data.yaml"
    if not meta_path.exists():
        logger.error("找不到 %s", meta_path)
        return 1
    meta = yaml.safe_load(meta_path.read_text(encoding="utf-8"))
    names = [meta["names"][i] for i in sorted(meta["names"])]

    logger.info("来源: %s（%d 张，%d 类）", source, len(samples), len(names))
    try:
        default_policy, per_class_policy = parse_empty_policy(args.empty_policy)
        samples, empty_stats = apply_empty_policy(samples, default_policy, per_class_policy, names)
    except ValueError as exc:
        logger.error("%s", exc)
        return 1

    total_empty = sum(count for _, count in empty_stats.values())
    if total_empty:
        logger.info("无框图片 %d 张，按来源类别处置：", total_empty)
        for group, (policy, count) in sorted(empty_stats.items(), key=lambda kv: -kv[1][1]):
            logger.info("  %-14s %5d 张 -> %s", group or "(顶层)", count, EMPTY_POLICY_DESC[policy])
        # 只对"默认落到 keep"的类别告警——显式写了 `类别=keep` 的是用户的判断，不再唠叨
        as_negative = sorted(
            g for g, (policy, _) in empty_stats.items()
            if policy == "keep" and g not in per_class_policy
        )
        if as_negative:
            logger.warning(
                "以上 %s 的空标签会以『此图无缺陷』的身份进入训练。"
                "如果其中某些类其实是有缺陷、只是没能定位出来，它们会教模型抑制真实缺陷——"
                "用 --empty-policy 逐类改成 drop 或 full。",
                "/".join(as_negative),
            )

    # 按类别分层切分；无框的图单独作为一组分配，避免它们全落进同一侧
    grouped: dict[int | None, list[Sample]] = defaultdict(list)
    for sample in samples:
        # 退化成整图框的那些已经有了确定类别，跟着本类一起分层
        grouped[sample.full_box if sample.class_id is None else sample.class_id].append(sample)

    rng = random.Random(args.seed)
    splits: dict[str, list] = {"train": [], "val": []}
    for _class_id, group in sorted(grouped.items(), key=lambda kv: (kv[0] is None, kv[0])):
        order = group[:]
        rng.shuffle(order)
        n_val = max(1, round(len(order) * args.val_ratio)) if len(order) > 1 else 0
        splits["val"].extend(order[:n_val])
        splits["train"].extend(order[n_val:])

    counts = {
        split: Counter(
            names[s.class_id] if s.class_id is not None
            else names[s.full_box] if s.full_box is not None
            else "(无框)"
            for s in items
        )
        for split, items in splits.items()
    }
    logger.info("划分: train %d / val %d", len(splits["train"]), len(splits["val"]))
    logger.info("%-14s %8s %8s", "类别", "train", "val")
    for name in names + ["(无框)"]:
        if counts["train"][name] or counts["val"][name]:
            logger.info("%-14s %8d %8d", name, counts["train"][name], counts["val"][name])

    to_convert = sum(not is_png(sample.image) for sample in samples)
    if to_convert:
        logger.info(
            "其中 %d 张不是 PNG，会无损转成 PNG（无法硬链接，占额外磁盘）；其余 %d 张硬链接",
            to_convert, len(samples) - to_convert,
        )

    if args.dry_run:
        logger.info("--dry-run：未写入任何文件")
        return 0

    if out_dir.exists() and any(out_dir.iterdir()):
        if not args.overwrite:
            logger.error("%s 已存在且非空；确认要覆盖请加 --overwrite", out_dir)
            return 1
        shutil.rmtree(out_dir)

    renamed = None if args.keep_names else assign_names(samples, names)
    modes: Counter[str] = Counter()
    seen: dict[str, Path] = {}
    name_rows = []
    for split, items in splits.items():
        for sample in items:
            image = sample.image
            stem = image.stem if renamed is None else renamed[image]
            output_name = f"{stem}.png"
            # 保留原名时，拍平并统一后缀后不同目录下的 a.png 与 a.bmp 会撞成同一个 a.png；
            # 重新编号时只有类别名仅大小写不同才会撞。按小写比较：macOS 默认文件系统
            # 不区分大小写，A.png 与 a.png 也会互相覆盖。
            key = output_name.lower()
            if key in seen:
                hint = (
                    "源数据集的文件名（去掉后缀后）需要全局唯一，"
                    "去掉 --keep-names 按类别重新编号，"
                    "或用 scripts/export_clusters.py --prefix-parent 重新导出"
                    if renamed is None
                    else "data.yaml 里有只差大小写的类别名，请先改成互不相同"
                )
                logger.error("输出文件名冲突：%s 与 %s 都会落成 %s。%s", seen[key], image, output_name, hint)
                return 1
            seen[key] = image
            modes[materialize_png(image, out_dir / "images" / split / output_name)] += 1
            name_rows.append((split, output_name, class_of(sample, names), str(image.resolve())))
            destination = out_dir / "labels" / split / f"{stem}.txt"
            if sample.full_box is None:
                link_or_copy(sample.label, destination)
            else:
                # 整图框：内容与源标签不同，只能新写一份，不能链接
                destination.parent.mkdir(parents=True, exist_ok=True)
                destination.write_text(f"{sample.full_box} 0.5 0.5 1.0 1.0\n", encoding="utf-8")

    write_data_yaml(out_dir / "data.yaml", names, root=out_dir.resolve(),
                    train="images/train", val="images/val")
    write_classes_txt(out_dir / "classes.txt", names)
    # utf-8-sig：带 BOM，Excel / Numbers 直接打开中文源路径不乱码
    with (out_dir / "name_map.csv").open("w", encoding="utf-8-sig", newline="") as handle:
        writer = csv.writer(handle)
        writer.writerow(["split", "name", "class", "source"])
        writer.writerows(sorted(name_rows, key=lambda row: (row[2], row[1])))
    (out_dir / "summary.json").write_text(
        json.dumps(
            {
                "source": str(source.resolve()),
                "num_images": len(samples),
                "empty_labels": {
                    group: {"policy": policy, "count": count}
                    for group, (policy, count) in sorted(empty_stats.items())
                },
                "empty_policy": {"default": default_policy, **per_class_policy},
                "renamed": renamed is not None,
                "val_ratio": args.val_ratio,
                "seed": args.seed,
                "split_sizes": {k: len(v) for k, v in splits.items()},
                "per_class": {k: dict(v) for k, v in counts.items()},
                "link_modes": dict(modes),
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )
    logger.info("图片落盘方式: %s（hard=硬链接, copy=跨盘复制, convert=转成 PNG）", dict(modes))
    logger.info("完成 -> %s", out_dir)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())