"""Train a YOLO11-seg mosaic detector on a generated synthetic dataset.

Prereqs: pip install -e ".[yolo]" plus PyTorch (CUDA build recommended — see README).

Usage:
    python scripts/training/train_mosaic.py --data dataset_mosaic/data.yaml --epochs 100

The resulting weights (runs/segment/<name>/weights/best.pt) drop straight into the
app: Параметры → Детектор «YOLO» → выбрать этот .pt. The class is named "mosaic",
which YoloDetector maps to CensorType.MOSAIC.
"""

from __future__ import annotations

import argparse


def main() -> None:
    ap = argparse.ArgumentParser(description="Train YOLO11-seg mosaic detector")
    ap.add_argument("--data", required=True, help="path to data.yaml from the generator")
    ap.add_argument("--model", default="yolo11n-seg.pt", help="base model (n/s/m...-seg)")
    ap.add_argument("--epochs", type=int, default=100)
    ap.add_argument("--imgsz", type=int, default=640)
    ap.add_argument("--batch", default="-1", help="batch size (-1 = auto)")
    ap.add_argument("--device", default=None, help="cuda / 0 / cpu (default: auto)")
    ap.add_argument("--name", default="mosaic", help="run name under the project dir")
    ap.add_argument("--project", default=None, help="output dir for runs (default: runs/segment)")
    args = ap.parse_args()

    try:
        from ultralytics import YOLO
    except ImportError as exc:  # pragma: no cover
        raise SystemExit('Не установлен ultralytics: pip install -e ".[yolo]"') from exc

    batch = int(args.batch) if str(args.batch).lstrip("-").isdigit() else args.batch
    model = YOLO(args.model)
    results = model.train(
        data=args.data,
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=batch,
        device=args.device,
        name=args.name,
        project=args.project,
    )
    save_dir = getattr(results, "save_dir", "runs/segment/" + args.name)
    print(f"\nГотово. Веса: {save_dir}/weights/best.pt")
    print("Подключите их в приложении: Параметры → Детектор YOLO → выбрать best.pt")


if __name__ == "__main__":
    main()