"""
Train YOLO26 semantic segmentation on CelebAMask-HQ (19-class face parsing).

Variants: nano (realtime, ships in the iOS app) and large (accuracy reference).
imgsz=512 == deploy resolution (masks are natively 512). ENet inverse-log class
weighting is automatic in the semantic trainer (rare classes: glasses/hat/
earring/necklace).

Augmentation notes:
  * fliplr=0.0  -- CRITICAL. RandomFlip flips geometry WITHOUT swapping the
    left/right class labels, so any hflip corrupts l_eye/r_eye, l_brow/r_brow,
    l_ear/r_ear. Must stay disabled.
  * mosaic + degrees/translate/scale -- scale/position/tilt robustness to close
    the aligned-crop (CelebA-HQ) -> live-selfie domain gap.

Throughput: uses MPSWorkersTrainer to undo ultralytics' forced workers=0 on MPS.

Timing probe:  ml/train.py --epochs 1 --fraction 0.04 --no-val --name bench_X
Full run:      ml/train.py --model yolo26l-sem.pt --epochs 60 --name celeba_large
"""
import argparse
from ultralytics import YOLO

from fast_trainer import MPSWorkersTrainer

DATA = "/Users/ari/FaceSegmentation/dataset_celebamaskhq_semantic/data.yaml"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", default=DATA)
    ap.add_argument("--model", default="yolo26n-sem.pt")
    ap.add_argument("--epochs", type=int, default=100)
    ap.add_argument("--imgsz", type=int, default=512)
    ap.add_argument("--batch", type=int, default=32)
    ap.add_argument("--fraction", type=float, default=1.0)
    ap.add_argument("--device", default="mps")
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--name", default="celeba_face")
    ap.add_argument("--project", default="/Users/ari/FaceSegmentation/runs_semantic")
    ap.add_argument("--patience", type=int, default=30)
    ap.add_argument("--cache", default="False")
    ap.add_argument("--mosaic", type=float, default=0.5)
    ap.add_argument("--val", action=argparse.BooleanOptionalAction, default=True)
    ap.add_argument("--plots", action=argparse.BooleanOptionalAction, default=True)
    # deterministic=True (ultralytics default) forces slow MPS kernels in the
    # semantic loss -- measured 5037 -> 2293 ms/iter when disabled.
    ap.add_argument("--deterministic", action=argparse.BooleanOptionalAction, default=False)
    ap.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True)
    ap.add_argument("--val-period", type=int, default=1,
                    help="validate every Nth epoch (plus the final epoch)")
    # ENet weights are (1/ln(1.02+p))**cls_pw. Ultralytics defaults cls_pw=0.0,
    # which makes every weight exactly 1.0 -- i.e. class weighting silently OFF.
    # At 1.0 the rare classes (necklace/eyes/glasses/earring) get ~45-50x the
    # weight of hair/background/skin (~3.4-4.2), a 14.6x dynamic range.
    ap.add_argument("--cls-pw", type=float, default=1.0,
                    help="ENet class-weight exponent; 0 disables class weighting")
    # The stock loss upsamples stride-8 logits to full mask resolution (159M
    # elements at 512/b32). Evaluating it lower is much cheaper but gives up
    # some sub-grid boundary supervision. 0 = stock behaviour.
    ap.add_argument("--loss-size", type=int, default=0,
                    help="compute the loss at this resolution instead of full mask res")
    args = ap.parse_args()

    if args.loss_size:
        import loss_patch
        loss_patch.apply(args.loss_size)
        print(f"loss_patch: computing semantic loss at {args.loss_size}px")

    MPSWorkersTrainer.val_period = args.val_period
    model = YOLO(args.model)
    model.train(
        trainer=MPSWorkersTrainer,
        data=args.data,
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=args.batch,
        device=args.device,
        workers=args.workers,
        project=args.project,
        name=args.name,
        fraction=args.fraction,
        patience=args.patience,
        cache=args.cache,
        deterministic=args.deterministic,
        amp=args.amp,
        cls_pw=args.cls_pw,
        # --- augmentation tuned for face parsing ---
        fliplr=0.0,        # CRITICAL: preserves left/right class semantics
        flipud=0.0,
        degrees=10.0,      # head tilt
        translate=0.1,
        scale=0.5,
        mosaic=args.mosaic,
        close_mosaic=10,
        mixup=0.0,
        copy_paste=0.0,
        # default hsv_* kept for lighting robustness
        val=args.val,
        plots=args.plots,
        seed=0,
        exist_ok=True,
    )


if __name__ == "__main__":
    main()