"""
Training script for YOLOv11n enemy detection model.

Usage:
    python train.py --data path/to/data.yaml --epochs 100

Expects YOLO-format dataset with data.yaml pointing to train/val image directories.
Export from Roboflow in "YOLOv11" format.
"""

import argparse
import os


def main():
    parser = argparse.ArgumentParser(description="Train YOLOv11n enemy detector")
    parser.add_argument("--data", required=True, help="Path to data.yaml")
    parser.add_argument("--epochs", type=int, default=100, help="Training epochs")
    parser.add_argument("--imgsz", type=int, default=640, help="Image size")
    parser.add_argument("--batch", type=int, default=16, help="Batch size")
    parser.add_argument("--device", default="0", help="CUDA device (0, cpu)")
    parser.add_argument("--name", default="enemy-v1", help="Run name")
    args = parser.parse_args()

    from ultralytics import YOLO

    model = YOLO("yolo11n.pt")  # start from pretrained nano

    model.train(
        data=args.data,
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=args.batch,
        device=args.device,
        name=args.name,
        patience=20,       # early stopping
        save=True,
        save_period=10,
        plots=True,
        verbose=True,
    )

    # Copy best weights to models directory
    best_path = os.path.join("runs", "detect", args.name, "weights", "best.pt")
    output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
    os.makedirs(output_dir, exist_ok=True)

    output_path = os.path.join(output_dir, f"{args.name}.pt")
    if os.path.exists(best_path):
        import shutil
        shutil.copy2(best_path, output_path)
        print(f"\nBest model copied to: {output_path}")
    else:
        print(f"\nWarning: {best_path} not found — check training output")


if __name__ == "__main__":
    main()