#!/usr/bin/env python
"""按比例切分成 YOLO 的 train / val / test 目录结构。

    python scripts/split_dataset.py --images data/film/images --labels data/film/labels \
        --val-ratio 0.2 --data-yaml data/film/data.yaml

默认建软链（不复制数据），产出：

```
images/train/  images/val/
labels/train/  labels/val/
```

原始文件保留在 `images/` `labels/` 根目录下不动；重复执行会先清掉旧的划分目录。
"""

from __future__ import annotations

import argparse
import random
import shutil
import sys
from pathlib import Path

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

from datasets.yolo import IMAGE_SUFFIXES, load_classes, write_data_yaml  # noqa: E402


def main() -> int:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument("--images", required=True, help="图片目录（划分前的平铺目录）")
    parser.add_argument("--labels", default=None, help="标注目录，缺省为 images 同级的 labels")
    parser.add_argument("--val-ratio", type=float, default=0.2)
    parser.add_argument("--test-ratio", type=float, default=0.0)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--copy", action="store_true", help="复制文件而不是建软链")
    parser.add_argument("--data-yaml", default=None, help="顺带更新/生成 data.yaml")
    parser.add_argument("--classes", nargs="*", default=None, help="写 data.yaml 时的类别顺序")
    args = parser.parse_args()

    images_dir = Path(args.images)
    labels_dir = Path(args.labels) if args.labels else images_dir.parent / "labels"
    splits = ("train", "val", "test")

    files = sorted(
        p
        for p in images_dir.iterdir()
        if p.is_file() and p.suffix.lower() in IMAGE_SUFFIXES and p.parent.name not in splits
    )
    if not files:
        print(f"{images_dir} 下没有待划分的图片（子目录里的已划分文件会被跳过）")
        return 1

    random.Random(args.seed).shuffle(files)
    n_val = int(len(files) * args.val_ratio)
    n_test = int(len(files) * args.test_ratio)
    partition = {
        "val": files[:n_val],
        "test": files[n_val : n_val + n_test],
        "train": files[n_val + n_test :],
    }

    for name, items in partition.items():
        img_out, label_out = images_dir / name, labels_dir / name
        for directory in (img_out, label_out):
            if directory.exists():
                shutil.rmtree(directory)
        if not items:
            continue
        img_out.mkdir(parents=True, exist_ok=True)
        label_out.mkdir(parents=True, exist_ok=True)

        empty = 0
        for image in items:
            label = labels_dir / f"{image.stem}.txt"
            place(image, img_out / image.name, args.copy)
            if label.exists():
                place(label, label_out / label.name, args.copy)
                empty += int(not label.read_text(encoding="utf-8").strip())
            else:
                empty += 1  # 缺 .txt 等价于无缺陷
        print(f"{name}: {len(items)} 张（其中无缺陷 {empty} 张） -> {img_out}")

    if args.data_yaml:
        classes = load_classes(args.classes or args.data_yaml)
        path = write_data_yaml(args.data_yaml, classes, root=images_dir.parent)
        print(f"data.yaml 已更新 {path}")
    return 0


def place(src: Path, dst: Path, copy: bool) -> None:
    if dst.exists() or dst.is_symlink():
        dst.unlink()
    if copy:
        shutil.copy2(src, dst)
    else:
        dst.symlink_to(src.resolve())


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