#!/usr/bin/env python
"""把 Label Studio 的 YOLO 导出装配成可训练的数据集，划分 train/val 放到 data/ 下。

    uv run python scripts/build_yolo_dataset.py \\
        --export data/film_yolo_shap.zip \\
        --images runs/boxes_shap_label_studio/images \\
        --out data/film_yolo

Label Studio 的 YOLO 导出**只有标签**：本地存储的图片它不会打包进去（`get_local_path` 对
local-files 直接返回磁盘路径，不复制，`YOLO_WITH_IMAGES` 也一样）。所以图片从我们自己准备
复核材料的那个目录取——文件名是一一对应的。

```
<export>.zip                       runs/<...>_label_studio/images/
├── classes.txt                    └── <类别>_<序号>.png
├── labels/<类别>_<序号>.txt
└── images/  ← 空的
```

装配结果：

```
data/<name>/
├── images/{train,val}/<类别>_<序号>.png   # 硬链接，不占额外磁盘
├── labels/{train,val}/<类别>_<序号>.txt
├── classes.txt                            # 原样沿用导出的顺序
├── data.yaml
└── summary.json
```

**类别顺序以导出的 classes.txt 为准。** Label Studio 的转换器把标签收进 set 再排序编号，
与 labeling config 里写的顺序无关，也与我们 `generate_boxes.py` 时的顺序无关。标签里的
class_id 是按导出的顺序写的，所以 data.yaml 必须照抄，不能另排。

**空标签保留为负样本。** 导出后仍为空的图（本例是 normal，预测已在 Label Studio 里批量删掉）
表示"这张图没有缺陷"，在检测训练里是压低误检的负样本。`--drop-empty` 可以丢掉它们。
注意：没复核、也没把预测转成标注的任务同样会导成空标签，那种"空"是信息缺失而不是负样本，
装配前应先在 Label Studio 里用 Create Annotations From Predictions 补上。
"""

from __future__ import annotations

import argparse
import json
import random
import shutil
import sys
import zipfile
from collections import Counter, defaultdict
from pathlib import Path
from typing import NamedTuple

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

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

logger = None

EMPTY = "(无框)"


class Sample(NamedTuple):
    """一张图、它的标签文本，以及分层划分用的类别。"""

    image: Path
    label_text: str
    stratum: str


class ExportError(RuntimeError):
    """导出内容不自洽，装配无法继续。"""


def read_export(path: Path) -> tuple[list[str], dict[str, str]]:
    """读 Label Studio 的 YOLO 导出（zip 或已解压目录）。

    Returns:
        (类别名列表, {文件名主干: 标签文本})
    """
    if path.is_dir():
        classes_file = path / "classes.txt"
        if not classes_file.exists():
            raise ExportError(f"{path} 下没有 classes.txt，这不像 Label Studio 的 YOLO 导出")
        classes = classes_file.read_text(encoding="utf-8").split()
        labels = {
            p.stem: p.read_text(encoding="utf-8")
            for p in sorted((path / "labels").glob("*.txt"))
        }
    else:
        with zipfile.ZipFile(path) as archive:
            names = archive.namelist()
            if "classes.txt" not in names:
                raise ExportError(f"{path} 里没有 classes.txt，这不像 Label Studio 的 YOLO 导出")
            classes = archive.read("classes.txt").decode("utf-8").split()
            labels = {
                Path(name).stem: archive.read(name).decode("utf-8")
                for name in sorted(names)
                if name.startswith("labels/") and name.endswith(".txt")
            }
    if not classes:
        raise ExportError("classes.txt 是空的")
    if not labels:
        raise ExportError("导出里没有任何标签文件")
    return classes, labels


def parse_boxes(text: str, stem: str, num_classes: int) -> list[int]:
    """校验一个标签文件，返回其中所有框的类别下标。

    坐标越界或类别越界都直接报错：这类问题在训练时不会崩，只会悄悄拉低指标。
    """
    class_ids = []
    for number, line in enumerate(text.splitlines(), start=1):
        line = line.strip()
        if not line:
            continue
        parts = line.split()
        if len(parts) != 5:
            raise ExportError(f"{stem}.txt 第 {number} 行有 {len(parts)} 个字段，YOLO 需要 5 个：{line!r}")
        try:
            class_id = int(parts[0])
            coordinates = [float(v) for v in parts[1:]]
        except ValueError as exc:
            raise ExportError(f"{stem}.txt 第 {number} 行无法解析：{line!r}") from exc
        if not 0 <= class_id < num_classes:
            raise ExportError(f"{stem}.txt 第 {number} 行的类别下标 {class_id} 超出 0~{num_classes - 1}")
        if not all(0.0 <= v <= 1.0 for v in coordinates):
            raise ExportError(f"{stem}.txt 第 {number} 行的坐标不在 [0,1] 内：{coordinates}")
        class_ids.append(class_id)
    return class_ids


def collect(
    classes: list[str], labels: dict[str, str], images_dir: Path
) -> tuple[list[Sample], list[str]]:
    """把标签与图片配对，返回 (样本列表, 多余的图片名)。"""
    samples = []
    missing = []
    for stem, text in labels.items():
        candidates = [
            images_dir / f"{stem}{suffix}" for suffix in (".png", *IMAGE_SUFFIXES)
        ]
        image = next((p for p in candidates if p.exists()), None)
        if image is None:
            missing.append(stem)
            continue
        class_ids = parse_boxes(text, stem, len(classes))
        # 分层用"框最多的那个类别"：一张图可能有多个类别的框，人工复核后尤其常见
        stratum = classes[Counter(class_ids).most_common(1)[0][0]] if class_ids else EMPTY
        samples.append(Sample(image, text, stratum))

    if missing:
        raise ExportError(
            f"有 {len(missing)} 个标签找不到对应的图片，例如 {missing[0]}；"
            f"--images 是否指向了准备复核材料时的 images/ 目录（当前 {images_dir}）"
        )
    known = {s.image.stem for s in samples}
    extra = sorted(
        p.stem for p in images_dir.iterdir()
        if p.suffix.lower() in IMAGE_SUFFIXES and p.stem not in known
    )
    return samples, extra


def split(samples: list[Sample], val_ratio: float, seed: int) -> dict[str, list[Sample]]:
    """按类别分层切分。小类只有几十张，随机整体切很容易让它在验证集里挂零。"""
    grouped: dict[str, list[Sample]] = defaultdict(list)
    for sample in samples:
        grouped[sample.stratum].append(sample)

    rng = random.Random(seed)
    splits: dict[str, list[Sample]] = {"train": [], "val": []}
    for _stratum, group in sorted(grouped.items()):
        order = sorted(group, key=lambda s: s.image.name)     # 先定序，结果才可复现
        rng.shuffle(order)
        n_val = max(1, round(len(order) * val_ratio)) if len(order) > 1 else 0
        splits["val"].extend(order[:n_val])
        splits["train"].extend(order[n_val:])
    return splits


def main() -> int:
    global logger
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--export", required=True, help="Label Studio 的 YOLO 导出（zip 或目录）")
    parser.add_argument(
        "--images", required=True, help="图片目录，即准备复核材料时的 <...>_label_studio/images",
    )
    parser.add_argument("--out", required=True, help="输出目录，建议放在 data/ 下")
    parser.add_argument("--val-ratio", type=float, default=0.2)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--drop-empty", 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()

    out_dir = Path(args.out)
    logger = setup_logging()

    try:
        classes, labels = read_export(Path(args.export))
        samples, extra = collect(classes, labels, Path(args.images))
    except (ExportError, FileNotFoundError) as exc:
        logger.error("%s", exc)
        return 1

    logger.info("导出: %s（%d 个标签，%d 类）", args.export, len(labels), len(classes))
    if extra:
        logger.warning(
            "%s 里有 %d 张图片没有对应的标签，不会进入数据集（例如 %s）",
            args.images, len(extra), extra[0],
        )

    empty = [s for s in samples if s.stratum == EMPTY]
    if args.drop_empty:
        samples = [s for s in samples if s.stratum != EMPTY]
    logger.info(
        "空标签 %d 张（%.1f%%）：%s",
        len(empty), 100 * len(empty) / max(len(labels), 1),
        "已丢弃" if args.drop_empty else "保留为负样本",
    )

    splits = split(samples, args.val_ratio, args.seed)
    counts = {name: Counter(s.stratum for s in items) for name, items in splits.items()}
    boxes = Counter()
    for sample in samples:
        for line in sample.label_text.splitlines():
            if line.strip():
                boxes[classes[int(line.split()[0])]] += 1

    logger.info("划分: train %d / val %d", len(splits["train"]), len(splits["val"]))
    logger.info("%-12s %8s %8s %8s", "类别", "train", "val", "框数")
    for name in [*classes, EMPTY]:
        if counts["train"][name] or counts["val"][name]:
            logger.info(
                "%-12s %8d %8d %8d",
                name, counts["train"][name], counts["val"][name], boxes[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()
    for name, items in splits.items():
        for sample in items:
            modes[link_or_copy(sample.image, out_dir / "images" / name / sample.image.name)] += 1
            label_path = out_dir / "labels" / name / f"{sample.image.stem}.txt"
            label_path.parent.mkdir(parents=True, exist_ok=True)
            label_path.write_text(sample.label_text, encoding="utf-8")

    write_data_yaml(out_dir / "data.yaml", classes, root=out_dir.resolve(),
                    train="images/train", val="images/val")
    write_classes_txt(out_dir / "classes.txt", classes)
    (out_dir / "summary.json").write_text(
        json.dumps(
            {
                "export": str(Path(args.export).resolve()),
                "images": str(Path(args.images).resolve()),
                "num_images": len(samples),
                "num_boxes": sum(boxes.values()),
                "empty_labels": len(empty),
                "dropped_empty": args.drop_empty,
                "images_without_label": len(extra),
                "val_ratio": args.val_ratio,
                "seed": args.seed,
                "split_sizes": {k: len(v) for k, v in splits.items()},
                "per_class_images": {k: dict(v) for k, v in counts.items()},
                "per_class_boxes": dict(boxes),
                "link_modes": dict(modes),
            },
            ensure_ascii=False,
            indent=2,
        ),
        encoding="utf-8",
    )
    logger.info("图片落盘方式: %s", dict(modes))
    logger.info("完成 -> %s（%d 个框）", out_dir, sum(boxes.values()))
    return 0


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