#!/usr/bin/env python3
import random
import shutil
from pathlib import Path


HERE = Path(__file__).resolve().parent
SOURCE = HERE / "dataset/label-studio"
IMAGES = SOURCE / "images"
LABELS = SOURCE / "labels"
OUTPUT = HERE / "training_new"
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}


def main():
    if not IMAGES.is_dir():
        raise SystemExit(f"Нет каталога {IMAGES}")
    images = sorted(path for path in IMAGES.iterdir() if path.suffix.lower() in IMAGE_EXTENSIONS)
    if not images:
        raise SystemExit(f"Нет кадров в {IMAGES}")
    if not LABELS.is_dir():
        raise SystemExit(f"Сначала экспортируйте YOLO-разметку в {LABELS}")
    classes_file = SOURCE / "classes.txt"
    if not classes_file.is_file():
        raise SystemExit(f"Нет {classes_file}")

    random.Random(42).shuffle(images)
    test_count = max(1, round(len(images) * 0.1)) if len(images) >= 3 else 0
    valid_count = max(1, round(len(images) * 0.2)) if len(images) >= 3 else 0
    parts = {
        "test": images[:test_count],
        "valid": images[test_count : test_count + valid_count],
        "train": images[test_count + valid_count :],
    }

    if OUTPUT.exists():
        shutil.rmtree(OUTPUT)

    for part, part_images in parts.items():
        image_dir = OUTPUT / part / "images"
        label_dir = OUTPUT / part / "labels"
        image_dir.mkdir(parents=True)
        label_dir.mkdir(parents=True)
        for image in part_images:
            shutil.copy2(image, image_dir / image.name)
            label = LABELS / f"{image.stem}.txt"
            destination = label_dir / f"{image.stem}.txt"
            if label.exists():
                shutil.copy2(label, destination)
            else:
                destination.touch()

    classes = classes_file.read_text(encoding="utf-8").splitlines()
    names = "\n".join(f"  {number}: {name}" for number, name in enumerate(classes))
    (OUTPUT / "data.yaml").write_text(
        f"path: {OUTPUT}\ntrain: train/images\nval: valid/images\ntest: test/images\n\nnames:\n{names}\n",
        encoding="utf-8",
    )

    print(f"Готово: train={len(parts['train'])}, valid={len(parts['valid'])}, test={len(parts['test'])}")
    print(OUTPUT / "data.yaml")


if __name__ == "__main__":
    main()