from typing import Annotated, Literal, TypeVar, Union

import numpy as np
import numpy.typing as npt

DType = TypeVar("DType", bound=np.generic)

# Masks are type bool and have 2 dimensions
MaskArray = Annotated[npt.NDArray[np.bool_], Literal[2]]

# Grayscale images are type uint8 and have 2 dimensions
GrayArray = Annotated[npt.NDArray[np.uint8], Literal[2]]

# RGB images are type uint8 and have 3 dimensions
RGBArray = Annotated[npt.NDArray[np.uint8], Literal[3]]

# Generic 1D, 2D and 3D arrays
Generic1DArray = Annotated[npt.NDArray[DType], Literal[1]]
Generic2DArray = Annotated[npt.NDArray[DType], Literal[2]]
Generic3DArray = Annotated[npt.NDArray[DType], Literal[3]]

# Raw value arrays are type uint16 (raw bytes) or int64 (PNG) and have 2 dimensions
RawArray = Union[Generic2DArray[np.int64], Generic2DArray[np.uint16]]

Temperature = Literal["kelvin", "celsius", "fahrenheit"]
