Source code for spectral_indices.roi.bbox

"""Implement BoundingBox class to carry CRS and location information."""

from __future__ import annotations

from pathlib import Path
from typing import Annotated, Dict, List, Union

from geopandas import GeoDataFrame
from pydantic import BaseModel, Field, field_validator
from pyproj import CRS, Transformer
from shapely.geometry import Polygon, box
from shapely.ops import unary_union
from typeguard import typechecked
from xarray import DataArray, Dataset

from spectral_indices.roi.utils import read_raster_file, read_shape_file




class BoundingBox:
    """Initialize the bounding box with coordinates and a CRS.

    Args:
        bbox (Union[List[float], Tuple[float, ...]]): Xmin, Ymin, Xmax, Ymax.
        crs (Union[str, CRS]): Coordinate reference system.
    """

    def __init__(self, bbox, crs):
        self.bbox = bbox
        self.crs = crs if isinstance(crs, CRS) else CRS.from_user_input(crs)

    @field_validator("crs")
    def validate_crs(cls, value):
        val = value if isinstance(value, CRS) else CRS.from_user_input(value)
        return val

    @property
    def polygon(self) -> Polygon:
        return box(*self.bbox)

    @property
    def __geo_interface__(self) -> Dict[str, Union[str, list]]:
        """Return the coordinates of the bbox for the GeoJSON interface.

        The coordinates will be returned in the current CRS.

        Returns:
            dict: GeoJSON representation of the bbox.
        """
        return {
            "type": "Polygon",
            "coordinates": [
                [
                    [self.bbox[0], self.bbox[1]],  # Southwest corner
                    [self.bbox[2], self.bbox[1]],  # Southeast corner
                    [self.bbox[2], self.bbox[3]],  # Northeast corner
                    [self.bbox[0], self.bbox[3]],  # Northwest corner
                    [
                        self.bbox[0],
                        self.bbox[1],
                    ],  # Back to southwest corner
                ]
            ],
        }



    def to_geodataframe(self) -> GeoDataFrame:
        """Convert the bbox to a GeoDataFrame with CRS.

        Returns:
            gpd.GeoDataFrame: GeoDataFrame containing the bbox.
        """
        return GeoDataFrame({"geometry": [self.polygon]}, crs=self.crs.to_string())




    def transform(self, target_crs: Union[str, CRS]) -> BoundingBox:
        """Reproject the bbox to another CRS and return a new BoundingBox.

        Args:
            target_crs (Union[str, CRS]): The target coordinate reference system.

        Returns:
            BoundingBox: A new BoundingBox reprojected to the target CRS.
        """
        target_crs = CRS.from_user_input(target_crs)
        transformer = Transformer.from_crs(self.crs, target_crs, always_xy=True)
        xmin, ymin = transformer.transform(self.bbox[0], self.bbox[1])
        xmax, ymax = transformer.transform(self.bbox[2], self.bbox[3])
        bbox = (xmin, ymin, xmax, ymax)
        return BoundingBox(bbox=bbox, crs=target_crs)


    @classmethod
    @typechecked
    def union(cls, boxes: List[BoundingBox]) -> BoundingBox:
        crs = set([b.crs for b in boxes])
        if len(crs) > 1:
            raise ValueError(f"All bounding boxes should have the same CRS, got {crs}")
        crs = list(crs)[0]
        polygons = [bbox.polygon for bbox in boxes]
        total_bounds = unary_union(polygons).bounds
        return BoundingBox(bbox=list(total_bounds), crs=crs)



    @classmethod
    @typechecked
    def from_geodataframe(
        cls, shape_data: Union[GeoDataFrame, Union[str, Path]]
    ) -> BoundingBox:
        """Create BoundingBox from shape file or geoDataFrame.

        Args:
            shape_data (``Union[gpd.GeoDataFrame, Union[str, Path]]``): Path to shape file or geoDataFrame.

        Raises:
            ValueError: CRS is null.

        Returns:
            ``BoundingBox``:
                - Shape corresponding BoundingBox.
        """
        if isinstance(shape_data, (str, Path)):
            shape_data = read_shape_file(shape_data)

        if shape_data.crs is None:
            raise Exception("The GeoDataFrame must have a defined CRS.")

        # Calculate the bounding box
        bounds = shape_data.total_bounds  # (xmin, ymin, xmax, ymax)
        crs = shape_data.crs
        bbox = (bounds[0], bounds[1], bounds[2], bounds[3])
        return cls(bbox=bbox, crs=crs)




    @classmethod
    @typechecked
    def from_raster(
        cls, raster_data: Union[Union[str, Path], Union[Dataset, DataArray]]
    ) -> BoundingBox:
        """Create an instance of BoundingBox from a raster file or rio xarray.

        Args:
            raster_data (``Union[Union[str, Path], Union[Dataset, DataArray]]``): Path to raster file or xarray with rio attr.

        Returns:
            ``BoundingBox``:
                - Raster corresponding BoundingBox
        """
        if isinstance(raster_data, (str, Path)):
            raster_data = read_raster_file(raster_data)

        # Get the raster coordinates
        bounds = raster_data.rio.bounds()
        crs = raster_data.rio.crs.to_string()  # Get the CRS of the raster
        bbox = (bounds[0], bounds[1], bounds[2], bounds[3])
        return cls(bbox, crs)


    def __repr__(self) -> str:
        return f"BoundingBox: {self.bbox}, CRS: {self.crs.to_string()}"