#!/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. 全部用**硬链接**，不占额外磁盘。

```
data/<name>/
├── images/{train,val}/<图名>.png
├── labels/{train,val}/<图名>.txt
└── data.yaml
```

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

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

from __future__ import annotations

import argparse
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_data_yaml  # noqa: E402
from utils import setup_logging  # noqa: E402

logger = None

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


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 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("--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])

    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)

    modes: Counter[str] = Counter()
    seen: dict[str, Path] = {}
    for split, items in splits.items():
        for sample in items:
            image = sample.image
            # 拍平之后不同类别目录下的同名文件会撞车，显式挡住
            if image.name in seen:
                logger.error(
                    "拍平后文件名冲突：%s 与 %s 同名。源数据集的文件名需要全局唯一，"
                    "用 scripts/export_clusters.py --prefix-parent 重新导出即可",
                    seen[image.name], image,
                )
                return 1
            seen[image.name] = image
            modes[link_or_copy(image, out_dir / "images" / split / image.name)] += 1
            destination = out_dir / "labels" / split / f"{image.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")
    (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},
                "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", dict(modes))
    logger.info("完成 -> %s", out_dir)
    return 0


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