src/iohmm_evac/types.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright (C) 2026 SWGY, Inc.
"""Semantic type aliases and enums shared across the package."""

from __future__ import annotations

from enum import IntEnum, StrEnum

import numpy as np
from numpy.typing import NDArray

__all__ = [
    "BoolArray",
    "EvacPath",
    "FloatArray",
    "IntArray",
    "State",
    "Zone",
]


FloatArray = NDArray[np.float64]
"""1-D or N-D array of float64 values."""

IntArray = NDArray[np.int64]
"""Array of int64 values (general-purpose integer)."""

BoolArray = NDArray[np.bool_]
"""Array of boolean values."""


class State(IntEnum):
    """Latent behavioral state of a household.

    The integer values are stable and used as compact array codes.
    """

    UA = 0
    AW = 1
    PR = 2
    ER = 3
    SH = 4

    @classmethod
    def n_states(cls) -> int:
        """Return the number of distinct states."""
        return len(cls)


class Zone(StrEnum):
    """Coarse evacuation zone derived from coastal distance."""

    A = "A"
    B = "B"
    C = "C"


class EvacPath(StrEnum):
    """Whether a household ended up sheltering away or in place.

    Set when SH is first entered. Households that have not yet entered SH
    carry ``NONE``.
    """

    NONE = "none"
    AWAY = "away"
    HOME = "home"