training_status

@dataclass
class TrainUpdate:
    """Update from the training worker."""

    type: TrainUpdateType
    """Type of the update."""

    # TODO should we split into subclasses to make the value type more specific?
    value: Optional[Union[int, float, str, TrainingState, CAREamist, Exception]] = None
    """Content of the update."""

class TrainUpdateType(str, Enum):
    """Type of training update."""

    MAX_EPOCH = "max_epochs"
    """Number of epochs."""

    EPOCH = "epoch_idx"
    """Index of the current epoch."""

    MAX_BATCH = "max_batches"
    """Number of batches."""

    BATCH = "batch_idx"
    """Index of the current batch."""

    LOSS = "loss"
    """Current loss value."""

    VAL_LOSS = "val_loss"
    """Current validation loss value."""

    STATE = "state"
    """Current state of the training process."""

    CAREAMIST = "careamist"
    """CAREamist instance."""

    DEBUG = "debug message"
    """Debug message."""

    EXCEPTION = "exception"
    """Exception raised during the training process."""

class TrainingState(IntEnum):
    """Training state."""

    IDLE = 0
    """Training is idle."""

    TRAINING = 1
    """Training is ongoing."""

    DONE = 2
    """Training is done."""

    STOPPED = 3
    """Training was stopped."""

    CRASHED = 4
    """Training crashed."""

@evented
@dataclass
class TrainingStatus:
    """Status of the training thread.

    This dataclass is used to update the training UI with the current status and
    progress of the training. Listeners can be registered to the `events` attribute to
    be notified of changes in the value of the attributes (see `psygnal` documentation
    for more details).
    """

    if TYPE_CHECKING:
        events: TrainingStatusSignalGroup
        """Attribute allowing the registration of parameter-specific listeners."""

    max_epochs: int = -1
    """Number of epochs."""

    max_batches: int = -1
    """Number of batches."""

    epoch_idx: int = -1
    """Index of the current epoch."""

    batch_idx: int = -1
    """Index of the current batch."""

    loss: float = -1
    """Current loss value."""

    val_loss: float = -1
    """Current validation loss value."""

    state: TrainingState = TrainingState.IDLE
    """Current state of the training process."""

    def update(self, new_update: TrainUpdate) -> None:
        """Update the status with the new values.

        Exceptions, debugging messages and CAREamist instances are ignored.

        Parameters
        ----------
        new_update : PredictionUpdate
            New update to apply.
        """
        if (
            new_update.type != TrainUpdateType.CAREAMIST
            and new_update.type != TrainUpdateType.EXCEPTION
            and new_update.type != TrainUpdateType.DEBUG
        ):
            setattr(self, new_update.type.value, new_update.value)

class TrainingStatusSignalGroup(SignalGroup):
    """Signal group for the training status dataclass."""

    max_epochs: SignalInstance
    """Number of epochs."""

    max_batches: SignalInstance
    """Number of batches."""

    epoch_idx: SignalInstance
    """Index of the current epoch."""

    batch_idx: SignalInstance
    """Index of the current batch."""

    loss: SignalInstance
    """Current loss value."""

    val_loss: SignalInstance
    """Current validation loss value."""

    state: SignalInstance
    """Current state of the training process."""