"""Train and export a one-class MLB ball detector with Ultralytics YOLO."""

from __future__ import annotations

import argparse

import shutil

from pathlib import Path

from rich.console import Console

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

DATA_YAML = ROOT / "configs" / "ball_yolo.yaml"

MODEL_DIR = ROOT / "configs" / "models"

console = Console()

def main() -> int:

    ap = argparse.ArgumentParser(description="Train/export MLB ball YOLO model.")

    ap.add_argument("--model", default="yolo11n.pt", help="Base model, e.g. yolo11n.pt or yolov8n.pt.")

    ap.add_argument("--epochs", type=int, default=80)

    ap.add_argument("--imgsz", type=int, default=640)

    ap.add_argument("--batch", type=float, default=None, help="Batch size. Omit for Ultralytics default.")

    ap.add_argument("--device", default=None, help="cuda device id, cpu, or omit for auto.")

    ap.add_argument("--name", default="mlb26_ball")

    ap.add_argument("--workers", type=int, default=0, help="Dataloader workers. Use 0 on Windows/network paths.")

    args = ap.parse_args()

    from ultralytics import YOLO

    model = YOLO(args.model)

    train_kwargs = {

        "data": str(DATA_YAML),

        "epochs": args.epochs,

        "imgsz": args.imgsz,

        "name": args.name,

        "project": str(ROOT / "runs" / "detect"),

        "single_cls": True,

        "patience": 25,

        "amp": False,

        "workers": args.workers,

        "mosaic": 0.0,

        "fliplr": 0.0,

        "erasing": 0.0,

        "hsv_h": 0.0,

        "hsv_s": 0.15,

        "hsv_v": 0.15,

    }

    if args.batch is not None:

        train_kwargs["batch"] = int(args.batch) if args.batch >= 1 else args.batch

    if args.device:

        train_kwargs["device"] = args.device

    results = model.train(**train_kwargs)

    best = Path(results.save_dir) / "weights" / "best.pt"

    if not best.exists():

        console.print(f"[red]Training finished but best.pt was not found at {best}[/red]")

        return 2

    trained = YOLO(str(best))

    onnx_path = Path(trained.export(format="onnx", imgsz=args.imgsz, simplify=True, opset=12))

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

    final_onnx = MODEL_DIR / "mlb26_ball.onnx"

    shutil.copy2(onnx_path, final_onnx)

    shutil.copy2(best, MODEL_DIR / "mlb26_ball.pt")

    console.print(f"[bold green]Exported ONNX -> {final_onnx}[/bold green]")

    console.print(f"[green]Saved PyTorch weights -> {MODEL_DIR / 'mlb26_ball.pt'}[/green]")

    return 0

if __name__ == "__main__":

    raise SystemExit(main())