"""
SemanticSegmentationTrainer tuned for Apple Silicon throughput.

Two fixes, both measured on an M5 Max (see ml/isolate_bottleneck.py):

1. Dataloader workers. ultralytics/engine/trainer.py:162 hard-forces
   `args.workers = 0` whenever the device is cpu OR mps ("faster CPU training as
   time dominated by inference") -- a CPU-era assumption. Measured batch-load
   time at imgsz=512/batch=32: 136 ms with workers=0 vs 28 ms with workers=8.
   Not the dominant cost, but free to reclaim.

2. Validation frequency. The semantic validator accumulates full-resolution
   confusion statistics and cost ~18 min for the 2993-image val split, which
   over a long run exceeds the training time itself. `val_period` does not
   exist in this ultralytics build, so we skip intermediate validations and
   keep the last known metrics/fitness, which leaves best-checkpoint selection
   and EarlyStopping functional (they simply do not update on skipped epochs).
   The final epoch always validates.

Usage:
    model.train(trainer=MPSWorkersTrainer, workers=8, ...)          # fix 1 only
    MPSWorkersTrainer.val_period = 4                                # + fix 2
"""
from ultralytics.models.yolo.semantic.train import SemanticSegmentationTrainer
from ultralytics.utils import DEFAULT_CFG, LOGGER


class MPSWorkersTrainer(SemanticSegmentationTrainer):
    #: Validate every Nth epoch (1 = every epoch, ultralytics' behaviour).
    val_period: int = 1

    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
        overrides = overrides or {}
        requested = int(overrides.pop("workers_mps", 0) or overrides.get("workers") or 0)
        vp = overrides.pop("val_period", type(self).val_period)
        self.val_period = 1 if vp is None else int(vp)   # NB: 0 is valid (= never validate)
        super().__init__(cfg, overrides, _callbacks)
        if requested > 0 and self.device.type == "mps" and self.args.workers != requested:
            LOGGER.info(f"MPSWorkersTrainer: restoring dataloader workers 0 -> {requested} on MPS")
            self.args.workers = requested
        if self.val_period > 1:
            LOGGER.info(f"MPSWorkersTrainer: validating every {self.val_period} epochs (plus the final epoch)")

    def validate(self):
        """Validate on every val_period-th epoch, on the final epoch, and whenever
        the trainer is stopping; otherwise reuse the last known metrics.

        val_period == 0 disables validation entirely, including the final epoch
        that ultralytics validates unconditionally (trainer.py:575). Intended for
        throughput benchmarking, where the 2993-image val split would dominate.
        """
        if self.val_period == 0:
            return {}, 0.0
        if self.val_period > 1:
            epoch = getattr(self, "epoch", 0)          # 0-based
            final = (epoch + 1) >= self.epochs
            due = ((epoch + 1) % self.val_period == 0)
            if not (final or due or getattr(self, "stop", False)):
                metrics = getattr(self, "metrics", None) or {}
                fitness = getattr(self, "fitness", None)
                if fitness is None:                     # nothing measured yet
                    fitness = float(-self.loss.detach().cpu().numpy()) if self.loss is not None else 0.0
                return metrics, fitness
        return super().validate()