"""Split labeled YOLO data into train/val/test folders."""

from __future__ import annotations

import argparse

import random

import shutil

from pathlib import Path

from rich.console import Console

ROOT = Path(__file__).resolve().parents[1]

DATA = ROOT / "datasets" / "mlb26_ball_yolo"

IMG_ALL = DATA / "images" / "all"

LBL_ALL = DATA / "labels" / "all"

console = Console()

def reset_split_dirs() -> None:

    for kind in ("images", "labels"):

        for split in ("train", "val", "test"):

            d = DATA / kind / split

            if d.exists():

                shutil.rmtree(d)

            d.mkdir(parents=True, exist_ok=True)

def main() -> int:

    ap = argparse.ArgumentParser(description="Split YOLO ball dataset.")

    ap.add_argument("--val", type=float, default=0.18, help="Validation fraction.")

    ap.add_argument("--test", type=float, default=0.05, help="Test fraction.")

    ap.add_argument("--seed", type=int, default=26)

    ap.add_argument(

        "--max-neg-ratio",

        type=float,

        default=4.0,

        help="Maximum empty/no-ball images per positive image. Use 0 to keep all negatives.",

    )

    args = ap.parse_args()

    if not IMG_ALL.exists() or not LBL_ALL.exists():

        console.print(f"[red]Expected images at {IMG_ALL} and labels at {LBL_ALL}.[/red]")

        return 2

    images = sorted([p for p in IMG_ALL.iterdir() if p.suffix.lower() in {".jpg", ".jpeg", ".png"}])

    positives = []

    negatives = []

    for img in images:

        label = LBL_ALL / f"{img.stem}.txt"

        if label.exists():

            if label.stat().st_size > 0:

                positives.append((img, label))

            else:

                negatives.append((img, label))

    if not positives:

        console.print("[red]No labeled image/label pairs found.[/red]")

        return 3

    rng = random.Random(args.seed)

    rng.shuffle(positives)

    rng.shuffle(negatives)

    if args.max_neg_ratio and negatives:

        max_negs = int(round(len(positives) * args.max_neg_ratio))

        negatives = negatives[:max_negs]

    def split_pairs(pairs):

        n = len(pairs)

        n_test = int(round(n * args.test))

        n_val = int(round(n * args.val))

        return {

            "test": pairs[:n_test],

            "val": pairs[n_test:n_test + n_val],

            "train": pairs[n_test + n_val:],

        }

    pos_buckets = split_pairs(positives)

    neg_buckets = split_pairs(negatives)

    buckets = {

        split: pos_buckets[split] + neg_buckets[split]

        for split in ("test", "val", "train")

    }

    reset_split_dirs()

    for split, split_pairs in buckets.items():

        for img, label in split_pairs:

            shutil.copy2(img, DATA / "images" / split / img.name)

            shutil.copy2(label, DATA / "labels" / split / label.name)

        pos_count = sum(1 for _, label in split_pairs if label.stat().st_size > 0)

        neg_count = len(split_pairs) - pos_count

        console.print(f"[green]{split}: {len(split_pairs)} ({pos_count} ball, {neg_count} empty)[/green]")

    console.print(f"[bold green]Dataset ready at {DATA}[/bold green]")

    return 0

if __name__ == "__main__":

    raise SystemExit(main())