#!/usr/bin/env python3
"""Fine-tune YOLO11n (COCO-pretrained) as a single-class face detector on converted WIDER FACE.

Adds video-like degradation (motion blur + JPEG compression) on top of ultralytics' default augmentation,
implemented with plain cv2 so no albumentations / second OpenCV build is needed.

Usage:
  python train.py                       # full run (40 epochs)
  python train.py --fraction 0.02 --epochs 1 --name smoke   # quick pipeline check
"""
import argparse
import random
from pathlib import Path

import cv2
import numpy as np
from ultralytics import YOLO
from ultralytics.data import augment


def _motion_blur(img, max_k=15):
    k = random.choice(range(3, max_k + 1, 2))
    kernel = np.zeros((k, k), np.float32)
    kernel[k // 2, :] = 1.0
    rot = cv2.getRotationMatrix2D((k / 2 - 0.5, k / 2 - 0.5), random.uniform(0, 180), 1.0)
    kernel = cv2.warpAffine(kernel, rot, (k, k))
    s = kernel.sum()
    return img if s == 0 else cv2.filter2D(img, -1, kernel / s)


def _jpeg(img, qmin=20, qmax=90):
    ok, enc = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, random.randint(qmin, qmax)])
    return cv2.imdecode(enc, cv2.IMREAD_COLOR) if ok else img


def install_video_degradation(p_blur=0.3, p_jpeg=0.3):
    """Wrap ultralytics' Albumentations step (runs after mosaic/affine, train only) with cv2 degradations."""
    orig = augment.Albumentations.__call__

    def call(self, labels):
        img = labels["img"]
        if random.random() < p_blur:
            img = _motion_blur(img)
        if random.random() < p_jpeg:
            img = _jpeg(img)
        labels["img"] = img
        return orig(self, labels)

    augment.Albumentations.__call__ = call


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default="dataset/data.yaml")
    ap.add_argument("--model", default="yolo11n.pt")
    ap.add_argument("--epochs", type=int, default=40)
    ap.add_argument("--batch", type=int, default=64)
    ap.add_argument("--imgsz", type=int, default=640)
    ap.add_argument("--workers", type=int, default=16)
    ap.add_argument("--fraction", type=float, default=1.0, help="fraction of train set (smoke tests)")
    ap.add_argument("--cache", default="disk", help="ram | disk | False")
    ap.add_argument("--project", default=str(Path(__file__).resolve().parent / "runs"))
    ap.add_argument("--name", default="face_yolo11n")
    ap.add_argument("--no-degrade", action="store_true")
    ap.add_argument("--resume", default=None,
                    help="path to a stopped run's last.pt; continues it and finishes at --epochs "
                         "(the LR schedule is recomputed for the new total)")
    a = ap.parse_args()

    if not a.no_degrade:
        install_video_degradation()

    if a.resume:
        import torch
        ck = torch.load(a.resume, map_location="cpu", weights_only=False)
        ck["train_args"]["epochs"] = a.epochs
        torch.save(ck, a.resume)
        model = YOLO(a.resume)
        model.train(resume=True)
    else:
        model = YOLO(a.model)
        model.train(
            data=a.data, epochs=a.epochs, imgsz=a.imgsz, batch=a.batch, workers=a.workers,
            fraction=a.fraction, cache=False if a.cache == "False" else a.cache, amp=True,
            close_mosaic=5, patience=0, project=a.project, name=a.name, exist_ok=True,
            plots=False,
        )
    m = model.val(data=a.data, imgsz=a.imgsz, batch=a.batch)
    print(f"RESULT mAP50={m.box.map50:.4f} mAP50-95={m.box.map:.4f} P={m.box.mp:.4f} R={m.box.mr:.4f}")


if __name__ == "__main__":
    main()