prediction_status

class PredictionState(IntEnum):
    """Prediction state."""

    IDLE = 0
    """Prediction is idle."""

    PREDICTING = 1
    """Prediction is ongoing."""

    DONE = 2
    """Prediction is done."""

    STOPPED = 3
    """Prediction was stopped."""

    CRASHED = 4
    """Prediction crashed."""

@evented
@dataclass
class PredictionStatus:
    """Status of the prediction thread.

    This dataclass is used to update the prediction UI with the current status and
    progress of the prediction. 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: PredictionStatusSignalGroup
        """Attribute allowing the registration of parameter-specific listeners."""

    max_samples: int = -1
    """Number of samples."""

    sample_idx: int = -1
    """Index of the current sample being predicted."""

    state: PredictionState = PredictionState.IDLE
    """Current state of the prediction process."""

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

        Exceptions, debugging messages and samples are ignored.

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

class PredictionStatusSignalGroup(SignalGroup):
    """Signal group for the prediction status dataclass."""

    max_samples: SignalInstance
    """Number of samples."""

    sample_idx: SignalInstance
    """Index of the current sample being predicted."""

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

@dataclass
class PredictionUpdate:
    """Update from the prediction worker."""

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

    value: Optional[Union[int, float, str, NDArray, PredictionState, Exception]] = None
    """Content of the update."""

class PredictionUpdateType(str, Enum):
    """Type of prediction update."""

    MAX_SAMPLES = "max_samples"
    """Number of samples."""

    SAMPLE_IDX = "sample_idx"
    """Index of the current sample being predicted."""

    SAMPLE = "sample"
    """Prediction result."""

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

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

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