"""Standalone EarlyStopping utility for PyTorch training loops.

Extracted from the ``ml_pytorch_vision_classification`` and ``r2plus1d``
BYOM codebases where it was duplicated verbatim.  This is a simpler,
lower-level alternative to :class:`~matrice_models.training.callbacks.EarlyStoppingCallback`
— use it directly inside a custom training loop when you don't want the
full callback system.
"""

from __future__ import annotations

import logging


logger = logging.getLogger(__name__)


class EarlyStopping:
    """Stop training when validation loss stops improving.

    Tracks the lowest observed loss.  If the loss does not decrease by
    at least *min_delta* for *patience* consecutive calls to
    :meth:`update`, the :attr:`stop` flag is set to ``True``.

    Example::

        es = EarlyStopping(patience=10, min_delta=0.001)
        for epoch in range(max_epochs):
            loss = train_and_validate(...)
            es.update(loss)
            if es.stop:
                break

    Args:
        patience: Number of updates without improvement before stopping.
        min_delta: Minimum decrease in loss to count as an improvement.
    """

    def __init__(self, patience: int = 5, min_delta: float = 0.001) -> None:
        """Initialize early-stopping state.

        Args:
            patience: Number of updates without improvement before stopping.
            min_delta: Minimum decrease in loss to count as an improvement.
        """
        self.patience = patience
        self.min_delta = abs(min_delta)
        self.counter: int = 0
        self.lowest_loss: float | None = None
        self.stop: bool = False

    def update(self, val_loss: float) -> None:
        """Record a new validation loss and check whether to stop.

        Args:
            val_loss: Current epoch's validation loss.
        """
        if self.lowest_loss is None:
            self.lowest_loss = val_loss
            return

        if self.lowest_loss - val_loss > self.min_delta:
            self.lowest_loss = val_loss
            self.counter = 0
        else:
            self.counter += 1
            logger.info(
                "EarlyStopping: no improvement for %d/%d epochs",
                self.counter,
                self.patience,
            )
            if self.counter >= self.patience:
                self.stop = True
                logger.info("EarlyStopping triggered — stopping training")

    def reset(self) -> None:
        """Reset internal state so the instance can be reused."""
        self.counter = 0
        self.lowest_loss = None
        self.stop = False

    def __repr__(self) -> str:
        """Return a developer-friendly representation."""
        return (
            f"EarlyStopping(patience={self.patience}, min_delta={self.min_delta}, "
            f"counter={self.counter}, stop={self.stop})"
        )
