Raster

class Raster(Spatial):
    """
    Represents a binary mask.

    Parameters
    ----------
    value : Dict[str, typing.Union[np.ndarray, str, None]], optional
        An raster value.

    Attributes
    ----------
    area
    array
    geometry
    height
    width

    Raises
    ------
    TypeError
        If `encoding` is not a string.

    Examples
    --------
    Generate a random mask.
    >>> import numpy.random
    >>> height = 640
    >>> width = 480
    >>> array = numpy.random.rand(height, width)

    Convert to binary mask.
    >>> mask = (array > 0.5)

    Create Raster.
    >>> Raster.from_numpy(mask)
    """

    def __init__(
        self,
        value: typing.Dict[
            str, typing.Union[np.ndarray, Box, Polygon, MultiPolygon, None]
        ],
    ):
        """
        Initialize and instance of a raster.

        Parameters
        ----------
        value : Dict[str, Union[np.ndarray, Box, Polygon, MultiPolygon, None]]
            The raster in dictionary format {"mask": <np.ndarray>, "geometry": <geometry | None>}.
        """
        super().__init__(value)

    @classmethod
    def __validate__(cls, value: typing.Any):
        """
        Validates typing.

        Parameters
        ----------
        value : Any
            The value to validate.

        Raises
        ------
        TypeError
            If the value type is not supported.
        """
        if not isinstance(value, dict):
            raise TypeError(
                "Raster should contain a dictionary describing a mask and optionally a geometry."
            )
        elif set(value.keys()) != {"mask", "geometry"}:
            raise ValueError(
                "Raster should be described by a dictionary with keys 'mask' and 'geometry'"
            )
        elif not isinstance(value["mask"], np.ndarray):
            raise TypeError(
                f"Expected mask to have type '{np.ndarray}' receieved type '{value['mask']}'"
            )
        elif len(value["mask"].shape) != 2:
            raise ValueError("raster only supports 2d arrays")
        elif value["mask"].dtype != bool:
            raise ValueError(
                f"Expecting a binary mask (i.e. of dtype bool) but got dtype {value['mask'].dtype}"
            )
        elif (
            value["geometry"] is not None
            and not Polygon.supports(value["geometry"])
            and not MultiPolygon.supports(value["geometry"])
        ):
            raise TypeError(
                "Expected geometry to conform to either Polygon or MultiPolygon or be 'None'"
            )

    def encode_value(self) -> typing.Any:
        """Encode object to JSON compatible dictionary."""
        value = self.get_value()
        if value is None:
            return None
        f = io.BytesIO()
        PIL.Image.fromarray(value["mask"]).save(f, format="PNG")
        f.seek(0)
        mask_bytes = f.read()
        f.close()
        return {
            "mask": b64encode(mask_bytes).decode(),
            "geometry": value["geometry"],
        }

    @classmethod
    def decode_value(cls, value: typing.Any):
        """Decode object from JSON compatible dictionary."""
        if value is None:
            return None
        if not (
            isinstance(value, dict)
            and set(value.keys()) == {"mask", "geometry"}
        ):
            raise ValueError(
                f"Improperly formatted raster encoding. Received '{value}'"
            )
        mask_bytes = b64decode(value["mask"])
        with io.BytesIO(mask_bytes) as f:
            img = PIL.Image.open(f)
            value = {
                "mask": np.array(img),
                "geometry": value["geometry"],
            }
        return cls(value=value)

    @classmethod
    def from_numpy(cls, mask: np.ndarray):
        """
        Create a Raster object from a NumPy array.

        Parameters
        ----------
        mask : np.ndarray
            The 2D binary array representing the mask.

        Returns
        -------
        Raster

        Raises
        ------
        ValueError
            If the input array is not 2D or not of dtype bool.
        """
        return cls(value={"mask": mask, "geometry": None})

    @classmethod
    def from_geometry(
        cls,
        geometry: typing.Union[Box, Polygon, MultiPolygon],
        height: int,
        width: int,
    ):
        """
        Create a Raster object from a geometric mask.

        Parameters
        ----------
        geometry : Union[Box, Polygon, MultiPolygon]
            Defines the bitmask as a geometry. Overrides any existing mask.
        height : int
            The intended height of the binary mask.
        width : int
            The intended width of the binary mask.

        Returns
        -------
        Raster
        """
        bitmask = np.full((int(height), int(width)), False)
        return cls(value={"mask": bitmask, "geometry": geometry.get_value()})

    @property
    def area(self) -> Float:
        """
        Symbolic representation of area.
        """
        if not isinstance(self._value, Symbol):
            raise ValueError
        return Float.symbolic(
            name=self._value._name,
            key=self._value._key,
            attribute="area",
        )

    @property
    def array(self) -> np.ndarray:
        """
        The bitmask as a numpy array.

        Returns
        -------
        Optional[np.ndarray]
            A 2D binary array representing the mask if it exists.
        """
        value = self.get_value()
        if value["geometry"] is not None:
            warnings.warn(
                "Raster array does not contain bitmask as this is a geometry-defined raster.",
                RuntimeWarning,
            )
        return value["mask"]

    @property
    def geometry(self) -> typing.Union[Box, Polygon, MultiPolygon]:
        """
        The geometric mask if it exists.

        Returns
        -------
        Box | Polygon | MultiPolygon | None
            The geometry if it exists.
        """
        return self.get_value()["geometry"]

    @property
    def height(self) -> int:
        """Returns the height of the raster if it exists."""
        return self.array.shape[0]

    @property
    def width(self) -> int:
        """Returns the width of the raster if it exists."""
        return self.array.shape[1]