"""
Training script for YOLO enemy/boss detection model.

Usage:
    python train.py --data path/to/data.yaml --epochs 200
    python train.py --data path/to/data.yaml --model yolo11m --imgsz 1280 --epochs 300

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

import argparse
import glob
import os


def run_training(args):
    """Run YOLO training. Called from main() or manage.py."""
    from ultralytics import YOLO

    model = YOLO(f"{args.model}.pt")

    model.train(
        data=args.data,
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=args.batch,
        device=args.device,
        name=args.name,
        patience=30,

        # Learning rate (fine-tuning pretrained, not from scratch)
        lr0=0.001,
        lrf=0.01,
        cos_lr=True,
        warmup_epochs=5,
        weight_decay=0.001,

        # Augmentation tuned for boss glow/morph effects
        hsv_h=0.03,
        hsv_s=0.8,
        hsv_v=0.6,
        scale=0.7,
        translate=0.2,
        degrees=5.0,
        mixup=0.15,
        close_mosaic=15,
        erasing=0.3,

        workers=0,          # avoid multiprocessing paging file issues on Windows
        save=True,
        save_period=10,
        plots=True,
        verbose=True,
    )

    # Find best.pt — try the trainer's save_dir first, then scan runs/detect/
    best_path = None
    save_dir = getattr(model.trainer, "save_dir", None)
    if save_dir:
        candidate = os.path.join(str(save_dir), "weights", "best.pt")
        if os.path.exists(candidate):
            best_path = candidate
    if not best_path:
        run_base = os.path.join("runs", "detect")
        candidates = sorted(glob.glob(os.path.join(run_base, f"{args.name}*", "weights", "best.pt")))
        best_path = candidates[-1] if candidates else os.path.join(run_base, 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)

    # If boss is set (from manage.py), deploy as boss-{boss}.pt; otherwise use run name
    boss = getattr(args, "boss", None)
    model_filename = f"boss-{boss}.pt" if boss else f"{args.name}.pt"
    output_path = os.path.join(output_dir, model_filename)
    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")


def main():
    parser = argparse.ArgumentParser(description="Train YOLO enemy/boss detector")
    parser.add_argument("--data", required=True, help="Path to data.yaml")
    parser.add_argument("--model", default="yolo11s", help="YOLO model variant (yolo11n, yolo11s, yolo11m)")
    parser.add_argument("--epochs", type=int, default=200, help="Training epochs")
    parser.add_argument("--imgsz", type=int, default=1280, help="Image size")
    parser.add_argument("--batch", type=int, default=8, 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()
    run_training(args)


if __name__ == "__main__":
    main()