"""
This module is an extension for xarray to provide rasterio capabilities
to xarray datasets/dataarrays.
"""

# pylint: disable=too-many-lines
import json
import math
import warnings
from collections.abc import Hashable, Iterable
from typing import Any, Literal, Optional, Union

import numpy
import pyproj
import rasterio.warp
import rasterio.windows
import xarray
from affine import Affine
from pyproj.aoi import AreaOfInterest
from pyproj.database import query_utm_crs_info
from rasterio.control import GroundControlPoint
from rasterio.crs import CRS

from rioxarray._options import EXPORT_GRID_MAPPING, get_option
from rioxarray.crs import crs_from_user_input
from rioxarray.exceptions import (
    DimensionError,
    DimensionMissingCoordinateError,
    InvalidDimensionOrder,
    MissingCRS,
    MissingSpatialDimensionError,
    NoDataInBounds,
    OneDimensionalRaster,
    RioXarrayError,
    TooManyDimensions,
)

DEFAULT_GRID_MAP = "spatial_ref"


def _affine_has_rotation(affine: Affine) -> bool:
    """
    Determine if the affine has rotation.

    Parameters
    ----------
    affine: :obj:`affine.Affine`
        The affine of the grid.

    Returns
    -------
    bool
    """
    return affine.b == affine.d != 0


def _resolution(affine: Affine) -> tuple[float, float]:
    """
    Determine if the resolution of the affine.
    If it has rotation, the sign of the resolution is lost.

    Based on: https://github.com/mapbox/rasterio/blob/6185a4e4ad72b5669066d2d5004bf46d94a6d298/rasterio/_base.pyx#L943-L951

    Parameters
    ----------
    affine: :obj:`affine.Affine`
        The affine of the grid.


    Returns
    --------
    x_resolution: float
        The X resolution of the affine.
    y_resolution: float
        The Y resolution of the affine.
    """
    if not _affine_has_rotation(affine):
        return affine.a, affine.e
    return (
        math.sqrt(affine.a**2 + affine.d**2),
        math.sqrt(affine.b**2 + affine.e**2),
    )


def affine_to_coords(
    affine: Affine, width: int, height: int, *, x_dim: str = "x", y_dim: str = "y"
) -> dict[str, numpy.ndarray]:
    """Generate 1d pixel centered coordinates from affine.

    Based on code from the xarray rasterio backend.

    Parameters
    ----------
    affine: :obj:`affine.Affine`
        The affine of the grid.
    width: int
        The width of the grid.
    height: int
        The height of the grid.
    x_dim: str, optional
        The name of the X dimension. Default is 'x'.
    y_dim: str, optional
        The name of the Y dimension. Default is 'y'.

    Returns
    -------
    dict: x and y coordinate arrays.

    """
    transform = affine * affine.translation(0.5, 0.5)
    if affine.is_rectilinear and not _affine_has_rotation(affine):
        x_coords, _ = transform * (numpy.arange(width), numpy.zeros(width))
        _, y_coords = transform * (numpy.zeros(height), numpy.arange(height))
    else:
        x_coords, y_coords = transform * numpy.meshgrid(
            numpy.arange(width),
            numpy.arange(height),
        )
    return {y_dim: y_coords, x_dim: x_coords}


def _generate_spatial_coords(
    *, affine: Affine, width: int, height: int
) -> dict[Hashable, Any]:
    """get spatial coords in new transform"""
    new_spatial_coords = affine_to_coords(affine, width, height)
    if new_spatial_coords["x"].ndim == 1:
        return {
            "x": xarray.IndexVariable("x", new_spatial_coords["x"]),
            "y": xarray.IndexVariable("y", new_spatial_coords["y"]),
        }
    return {
        "xc": (("y", "x"), new_spatial_coords["x"]),
        "yc": (("y", "x"), new_spatial_coords["y"]),
    }


def _get_nonspatial_coords(
    src_data_array: Union[xarray.DataArray, xarray.Dataset]
) -> dict[Hashable, Union[xarray.Variable, xarray.IndexVariable]]:
    coords: dict[Hashable, Union[xarray.Variable, xarray.IndexVariable]] = {}
    for coord in set(src_data_array.coords) - {
        src_data_array.rio.x_dim,
        src_data_array.rio.y_dim,
        DEFAULT_GRID_MAP,
    }:
        # skip 2D spatial coords
        if (
            src_data_array.rio.x_dim in src_data_array[coord].dims
            and src_data_array.rio.y_dim in src_data_array[coord].dims
        ):
            continue
        if src_data_array[coord].ndim == 1:
            coords[coord] = xarray.IndexVariable(
                src_data_array[coord].dims,
                src_data_array[coord].values,
                src_data_array[coord].attrs,
            )
        else:
            coords[coord] = xarray.Variable(
                src_data_array[coord].dims,
                src_data_array[coord].values,
                src_data_array[coord].attrs,
            )
    return coords


def _make_coords(
    *,
    src_data_array: Union[xarray.DataArray, xarray.Dataset],
    dst_affine: Affine,
    dst_width: int,
    dst_height: int,
    force_generate: bool = False,
) -> dict[Hashable, Any]:
    """Generate the coordinates of the new projected `xarray.DataArray`"""
    coords = _get_nonspatial_coords(src_data_array)
    if (
        force_generate
        or (
            src_data_array.rio.x_dim in src_data_array.coords
            and src_data_array.rio.y_dim in src_data_array.coords
        )
        or ("xc" in src_data_array.coords and "yc" in src_data_array.coords)
    ):
        new_coords = _generate_spatial_coords(
            affine=dst_affine, width=dst_width, height=dst_height
        )
        new_coords.update(coords)
        return new_coords
    return coords


def _get_data_var_message(obj: Union[xarray.DataArray, xarray.Dataset]) -> str:
    """
    Get message for named data variables.
    """
    try:
        return f" Data variable: {obj.name}" if obj.name else ""
    except AttributeError:
        return ""


def _get_spatial_dims(
    obj: Union[xarray.Dataset, xarray.DataArray], *, var: Union[Any, Hashable]
) -> tuple[str, str]:
    """
    Retrieve the spatial dimensions of the dataset
    """
    try:
        return obj[var].rio.x_dim, obj[var].rio.y_dim
    except MissingSpatialDimensionError as err:
        try:
            obj[var].rio.set_spatial_dims(
                x_dim=obj.rio.x_dim, y_dim=obj.rio.y_dim, inplace=True
            )
            return obj.rio.x_dim, obj.rio.y_dim
        except MissingSpatialDimensionError:
            raise err from None


def _has_spatial_dims(
    obj: Union[xarray.Dataset, xarray.DataArray], *, var: Union[Any, Hashable]
) -> bool:
    """
    Check to see if the variable in the Dataset has spatial dimensions
    """
    try:
        # pylint: disable=pointless-statement
        _get_spatial_dims(obj, var=var)
    except MissingSpatialDimensionError:
        return False
    return True


def _order_bounds(
    *,
    minx: float,
    miny: float,
    maxx: float,
    maxy: float,
    resolution_x: float,
    resolution_y: float,
) -> tuple[float, float, float, float]:
    """
    Make sure that the bounds are in the correct order
    """
    if resolution_y < 0:
        top = maxy
        bottom = miny
    else:
        top = miny
        bottom = maxy
    if resolution_x < 0:
        left = maxx
        right = minx
    else:
        left = minx
        right = maxx

    return left, bottom, right, top


class XRasterBase:
    """This is the base class for the GIS extensions for xarray"""

    # pylint: disable=too-many-instance-attributes
    def __init__(self, xarray_obj: Union[xarray.DataArray, xarray.Dataset]):
        self._obj: Union[xarray.DataArray, xarray.Dataset] = xarray_obj

        self._x_dim: Optional[Hashable] = None
        self._y_dim: Optional[Hashable] = None
        # Determine the spatial dimensions of the `xarray.DataArray`
        if "x" in self._obj.dims and "y" in self._obj.dims:
            self._x_dim = "x"
            self._y_dim = "y"
        elif "longitude" in self._obj.dims and "latitude" in self._obj.dims:
            self._x_dim = "longitude"
            self._y_dim = "latitude"
        else:
            # look for coordinates with CF attributes
            for coord in self._obj.coords:
                # make sure to only look in 1D coordinates
                # that has the same dimension name as the coordinate
                if self._obj.coords[coord].dims != (coord,):
                    continue
                if (self._obj.coords[coord].attrs.get("axis", "").upper() == "X") or (
                    self._obj.coords[coord].attrs.get("standard_name", "").lower()
                    in ("longitude", "projection_x_coordinate")
                ):
                    self._x_dim = coord
                elif (self._obj.coords[coord].attrs.get("axis", "").upper() == "Y") or (
                    self._obj.coords[coord].attrs.get("standard_name", "").lower()
                    in ("latitude", "projection_y_coordinate")
                ):
                    self._y_dim = coord

        # properties
        self._count: Optional[int] = None
        self._height: Optional[int] = None
        self._width: Optional[int] = None
        self._crs: Union[rasterio.crs.CRS, None, Literal[False]] = None
        self._gcps: Optional[list[GroundControlPoint]] = None

    @property
    def crs(self) -> Optional[rasterio.crs.CRS]:
        """:obj:`rasterio.crs.CRS`:
        Retrieve projection from :obj:`xarray.Dataset` | :obj:`xarray.DataArray`
        """
        if self._crs is not None:
            return None if self._crs is False else self._crs

        # look in wkt attributes to avoid using
        # pyproj CRS if possible for performance
        for crs_attr in ("spatial_ref", "crs_wkt"):
            try:
                self._set_crs(
                    self._obj.coords[self.grid_mapping].attrs[crs_attr],
                    inplace=True,
                )
                return self._crs
            except KeyError:
                pass

        # look in grid_mapping
        try:
            self._set_crs(
                pyproj.CRS.from_cf(self._obj.coords[self.grid_mapping].attrs),
                inplace=True,
            )
        except (KeyError, pyproj.exceptions.CRSError):
            try:
                # look in attrs for 'crs'
                self._set_crs(self._obj.attrs["crs"], inplace=True)
            except KeyError:
                self._crs = False
                return None
        return self._crs

    def _get_obj(self, inplace: bool) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Get the object to modify.

        Parameters
        ----------
        inplace: bool
            If True, returns self.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`
        """
        if inplace:
            return self._obj
        obj_copy = self._obj.copy(deep=True)
        # preserve attribute information
        obj_copy.rio._x_dim = self._x_dim
        obj_copy.rio._y_dim = self._y_dim
        obj_copy.rio._width = self._width
        obj_copy.rio._height = self._height
        obj_copy.rio._crs = self._crs
        obj_copy.rio._gcps = self._gcps
        return obj_copy

    def set_crs(
        self, input_crs: Any, inplace: bool = True
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Set the CRS value for the Dataset/DataArray without modifying
        the dataset/data array.

        .. deprecated:: 0.15.8
            It is recommended to use `rio.write_crs()` instead. This
        method will likely be removed in a future release.

        Parameters
        ----------
        input_crs: object
            Anything accepted by `rasterio.crs.CRS.from_user_input`.
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            Dataset with crs attribute.
        """
        warnings.warn(
            "It is recommended to use 'rio.write_crs()' instead. 'rio.set_crs()' will likely"
            "be removed in a future release.",
            FutureWarning,
            stacklevel=2,
        )

        return self._set_crs(input_crs, inplace=inplace)

    def _set_crs(
        self, input_crs: Any, inplace: bool = True
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Set the CRS value for the Dataset/DataArray without modifying
        the dataset/data array.

        Parameters
        ----------
        input_crs: object
            Anything accepted by `rasterio.crs.CRS.from_user_input`.
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        xarray.Dataset | xarray.DataArray
            Dataset with crs attribute.
        """
        crs = crs_from_user_input(input_crs)
        obj = self._get_obj(inplace=inplace)
        obj.rio._crs = crs
        return obj

    @property
    def grid_mapping(self) -> str:
        """
        str: The CF grid_mapping attribute. 'spatial_ref' is the default.
        """
        grid_mapping = self._obj.encoding.get(
            "grid_mapping", self._obj.attrs.get("grid_mapping")
        )
        if grid_mapping is not None:
            return grid_mapping
        grid_mapping = DEFAULT_GRID_MAP
        # search the dataset for the grid mapping name
        if hasattr(self._obj, "data_vars"):
            grid_mappings = set()
            for var in self._obj.data_vars:
                if not _has_spatial_dims(self._obj, var=var):
                    continue
                var_grid_mapping = self._obj[var].encoding.get(
                    "grid_mapping", self._obj[var].attrs.get("grid_mapping")
                )
                if var_grid_mapping is not None:
                    grid_mapping = var_grid_mapping
                    grid_mappings.add(grid_mapping)
            if len(grid_mappings) > 1:
                raise RioXarrayError("Multiple grid mappings exist.")
        return grid_mapping

    def write_grid_mapping(
        self, grid_mapping_name: str = DEFAULT_GRID_MAP, inplace: bool = False
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Write the CF grid_mapping attribute to the encoding.

        Parameters
        ----------
        grid_mapping_name: str, optional
            Name of the grid_mapping coordinate.
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            Modified dataset with CF compliant CRS information.
        """
        data_obj = self._get_obj(inplace=inplace)
        if hasattr(data_obj, "data_vars"):
            for var in data_obj.data_vars:
                try:
                    x_dim, y_dim = _get_spatial_dims(data_obj, var=var)
                except MissingSpatialDimensionError:
                    continue
                # remove grid_mapping from attributes if it exists
                # and update the grid_mapping in encoding
                new_attrs = dict(data_obj[var].attrs)
                new_attrs.pop("grid_mapping", None)
                data_obj[var].rio.update_encoding(
                    {"grid_mapping": grid_mapping_name}, inplace=True
                ).rio.set_attrs(new_attrs, inplace=True).rio.set_spatial_dims(
                    x_dim=x_dim, y_dim=y_dim, inplace=True
                )
        # remove grid_mapping from attributes if it exists
        # and update the grid_mapping in encoding
        new_attrs = dict(data_obj.attrs)
        new_attrs.pop("grid_mapping", None)
        return data_obj.rio.update_encoding(
            {"grid_mapping": grid_mapping_name}, inplace=True
        ).rio.set_attrs(new_attrs, inplace=True)

    def write_crs(
        self,
        input_crs: Optional[Any] = None,
        grid_mapping_name: Optional[str] = None,
        inplace: bool = False,
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Write the CRS to the dataset in a CF compliant manner.

        .. warning:: The grid_mapping attribute is written to the encoding.

        Parameters
        ----------
        input_crs: Any
            Anything accepted by `rasterio.crs.CRS.from_user_input`.
        grid_mapping_name: str, optional
            Name of the grid_mapping coordinate to store the CRS information in.
            Default is the grid_mapping name of the dataset.
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            Modified dataset with CF compliant CRS information.

        Examples
        --------
        Write the CRS of the current `xarray` object:

        >>> raster.rio.write_crs("epsg:4326", inplace=True)

        Write the CRS on a copy:

        >>> raster = raster.rio.write_crs("epsg:4326")
        """
        if input_crs is not None:
            data_obj = self._set_crs(input_crs, inplace=inplace)
        else:
            data_obj = self._get_obj(inplace=inplace)

        # get original transform
        transform = self._cached_transform()
        # remove old grid maping coordinate if exists
        grid_mapping_name = (
            self.grid_mapping if grid_mapping_name is None else grid_mapping_name
        )
        try:
            del data_obj.coords[grid_mapping_name]
        except KeyError:
            pass

        if data_obj.rio.crs is None:
            raise MissingCRS(
                "CRS not found. Please set the CRS with 'rio.write_crs()'."
            )
        # add grid mapping coordinate
        data_obj.coords[grid_mapping_name] = xarray.Variable((), 0)
        grid_map_attrs = {}
        if get_option(EXPORT_GRID_MAPPING):
            try:
                grid_map_attrs = pyproj.CRS.from_user_input(data_obj.rio.crs).to_cf()
            except KeyError:
                pass
        # spatial_ref is for compatibility with GDAL
        crs_wkt = data_obj.rio.crs.to_wkt()
        grid_map_attrs["spatial_ref"] = crs_wkt
        grid_map_attrs["crs_wkt"] = crs_wkt
        if transform is not None:
            grid_map_attrs["GeoTransform"] = " ".join(
                [str(item) for item in transform.to_gdal()]
            )
        data_obj.coords[grid_mapping_name].rio.set_attrs(grid_map_attrs, inplace=True)

        # remove old crs if exists
        data_obj.attrs.pop("crs", None)

        return data_obj.rio.write_grid_mapping(
            grid_mapping_name=grid_mapping_name, inplace=True
        )

    def estimate_utm_crs(self, datum_name: str = "WGS 84") -> rasterio.crs.CRS:
        """Returns the estimated UTM CRS based on the bounds of the dataset.

        .. versionadded:: 0.2

        .. note:: Requires pyproj 3+

        Parameters
        ----------
        datum_name : str, optional
            The name of the datum to use in the query. Default is WGS 84.

        Returns
        -------
        rasterio.crs.CRS
        """
        if self.crs is None:
            raise RuntimeError("crs must be set to estimate UTM CRS.")

        # ensure using geographic coordinates
        if self.crs.is_geographic:  # pylint: disable=no-member
            minx, miny, maxx, maxy = self.bounds(recalc=True)
        else:
            minx, miny, maxx, maxy = self.transform_bounds("EPSG:4326", recalc=True)

        x_center = numpy.mean([minx, maxx]).item()
        y_center = numpy.mean([miny, maxy]).item()

        utm_crs_list = query_utm_crs_info(
            datum_name=datum_name,
            area_of_interest=AreaOfInterest(
                west_lon_degree=x_center,
                south_lat_degree=y_center,
                east_lon_degree=x_center,
                north_lat_degree=y_center,
            ),
        )
        try:
            return CRS.from_epsg(utm_crs_list[0].code)
        except IndexError:
            raise RuntimeError("Unable to determine UTM CRS") from None

    def _cached_transform(self) -> Optional[Affine]:
        """
        Get the transform from:
        1. The GeoTransform metatada property in the grid mapping
        2. The transform attribute.
        """
        try:
            # look in grid_mapping
            transform = numpy.fromstring(
                self._obj.coords[self.grid_mapping].attrs["GeoTransform"], sep=" "
            )
            # Calling .tolist() to assure the arguments are Python float and JSON serializable
            return Affine.from_gdal(*transform.tolist())

        except KeyError:
            try:
                return Affine(*self._obj.attrs["transform"][:6])
            except KeyError:
                pass
        return None

    def write_transform(
        self,
        transform: Optional[Affine] = None,
        grid_mapping_name: Optional[str] = None,
        inplace: bool = False,
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        .. versionadded:: 0.0.30

        Write the GeoTransform to the dataset where GDAL can read it in.

        https://gdal.org/drivers/raster/netcdf.html#georeference

        Parameters
        ----------
        transform: affine.Affine, optional
            The transform of the dataset. If not provided, it will be calculated.
        grid_mapping_name: str, optional
            Name of the grid_mapping coordinate to store the transform information in.
            Default is the grid_mapping name of the dataset.
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            Modified dataset with Geo Transform written.
        """
        transform = transform or self.transform(recalc=True)
        data_obj = self._get_obj(inplace=inplace)
        # delete the old attribute to prevent confusion
        data_obj.attrs.pop("transform", None)
        grid_mapping_name = (
            self.grid_mapping if grid_mapping_name is None else grid_mapping_name
        )
        try:
            grid_map_attrs = data_obj.coords[grid_mapping_name].attrs.copy()
        except KeyError:
            data_obj.coords[grid_mapping_name] = xarray.Variable((), 0)
            grid_map_attrs = data_obj.coords[grid_mapping_name].attrs.copy()
        grid_map_attrs["GeoTransform"] = " ".join(
            [str(item) for item in transform.to_gdal()]
        )
        data_obj.coords[grid_mapping_name].rio.set_attrs(grid_map_attrs, inplace=True)
        return data_obj.rio.write_grid_mapping(
            grid_mapping_name=grid_mapping_name, inplace=True
        )

    def transform(self, recalc: bool = False) -> Affine:
        """
        Parameters
        ----------
        recalc: bool, optional
            If True, it will re-calculate the transform instead of using
            the cached transform.

        Returns
        -------
        :obj:`affine.Affine`:
            The affine of the :obj:`xarray.Dataset` | :obj:`xarray.DataArray`
        """
        transform = self._cached_transform()
        if transform and (
            not transform.is_rectilinear or _affine_has_rotation(transform)
        ):
            if recalc:
                warnings.warn(
                    "Transform that is non-rectilinear or with rotation found. "
                    "Unable to recalculate."
                )
            return transform

        try:
            src_left, _, _, src_top = self._unordered_bounds(recalc=recalc)
            src_resolution_x, src_resolution_y = self.resolution(recalc=recalc)
        except (DimensionMissingCoordinateError, DimensionError):
            return Affine.identity() if transform is None else transform
        return Affine.translation(src_left, src_top) * Affine.scale(
            src_resolution_x, src_resolution_y
        )

    def write_coordinate_system(
        self, inplace: bool = False
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Write the coordinate system CF metadata.

        .. versionadded:: 0.0.30

        Parameters
        ----------
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            The dataset with the CF coordinate system attributes added.
        """
        data_obj = self._get_obj(inplace=inplace)
        # add metadata to x,y coordinates
        is_projected = data_obj.rio.crs and data_obj.rio.crs.is_projected
        is_geographic = data_obj.rio.crs and data_obj.rio.crs.is_geographic
        x_coord_attrs = dict(data_obj.coords[self.x_dim].attrs)
        x_coord_attrs["axis"] = "X"
        y_coord_attrs = dict(data_obj.coords[self.y_dim].attrs)
        y_coord_attrs["axis"] = "Y"
        if is_projected:
            units = None
            if hasattr(data_obj.rio.crs, "linear_units_factor"):
                unit_factor = data_obj.rio.crs.linear_units_factor[-1]
                if unit_factor != 1:
                    units = f"{unit_factor} metre"
                else:
                    units = "metre"
            # X metadata
            x_coord_attrs["long_name"] = "x coordinate of projection"
            x_coord_attrs["standard_name"] = "projection_x_coordinate"
            if units:
                x_coord_attrs["units"] = units
            # Y metadata
            y_coord_attrs["long_name"] = "y coordinate of projection"
            y_coord_attrs["standard_name"] = "projection_y_coordinate"
            if units:
                y_coord_attrs["units"] = units
        elif is_geographic:
            # X metadata
            x_coord_attrs["long_name"] = "longitude"
            x_coord_attrs["standard_name"] = "longitude"
            x_coord_attrs["units"] = "degrees_east"
            # Y metadata
            y_coord_attrs["long_name"] = "latitude"
            y_coord_attrs["standard_name"] = "latitude"
            y_coord_attrs["units"] = "degrees_north"
        data_obj.coords[self.y_dim].attrs = y_coord_attrs
        data_obj.coords[self.x_dim].attrs = x_coord_attrs
        return data_obj

    def set_attrs(
        self, new_attrs: dict, inplace: bool = False
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Set the attributes of the dataset/dataarray and reset
        rioxarray properties to re-search for them.

        Parameters
        ----------
        new_attrs: dict
            A dictionary of new attributes.
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            Modified dataset with new attributes.
        """
        data_obj = self._get_obj(inplace=inplace)
        # set the attributes
        data_obj.attrs = new_attrs
        # reset rioxarray properties depending
        # on attributes to be generated
        data_obj.rio._nodata = None
        data_obj.rio._crs = None
        return data_obj

    def update_attrs(
        self, new_attrs: dict, inplace: bool = False
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Update the attributes of the dataset/dataarray and reset
        rioxarray properties to re-search for them.

        Parameters
        ----------
        new_attrs: dict
            A dictionary of new attributes to update with.
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            Modified dataset with updated attributes.
        """
        data_attrs = dict(self._obj.attrs)
        data_attrs.update(**new_attrs)
        return self.set_attrs(data_attrs, inplace=inplace)

    def set_encoding(
        self, new_encoding: dict, inplace: bool = False
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Set the encoding of the dataset/dataarray and reset
        rioxarray properties to re-search for them.

        .. versionadded:: 0.4

        Parameters
        ----------
        new_encoding: dict
            A dictionary for encoding.
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            Modified dataset with new attributes.
        """
        data_obj = self._get_obj(inplace=inplace)
        # set the attributes
        data_obj.encoding = new_encoding
        # reset rioxarray properties depending
        # on attributes to be generated
        data_obj.rio._nodata = None
        data_obj.rio._crs = None
        return data_obj

    def update_encoding(
        self, new_encoding: dict, inplace: bool = False
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Update the encoding of the dataset/dataarray and reset
        rioxarray properties to re-search for them.

        .. versionadded:: 0.4

        Parameters
        ----------
        new_encoding: dict
            A dictionary with encoding values to update with.
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            Modified dataset with updated attributes.
        """
        data_encoding = dict(self._obj.encoding)
        data_encoding.update(**new_encoding)
        return self.set_encoding(data_encoding, inplace=inplace)

    def set_spatial_dims(
        self, x_dim: str, y_dim: str, inplace: bool = True
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        This sets the spatial dimensions of the dataset.

        Parameters
        ----------
        x_dim: str
            The name of the x dimension.
        y_dim: str
            The name of the y dimension.
        inplace: bool, optional
            If True, it will modify the dataframe in place.
            Otherwise it will return a modified copy.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            Dataset with spatial dimensions set.
        """

        data_obj = self._get_obj(inplace=inplace)
        if x_dim in data_obj.dims:
            data_obj.rio._x_dim = x_dim
        else:
            raise MissingSpatialDimensionError(
                f"x dimension ({x_dim}) not found.{_get_data_var_message(data_obj)}"
            )
        if y_dim in data_obj.dims:
            data_obj.rio._y_dim = y_dim
        else:
            raise MissingSpatialDimensionError(
                f"y dimension ({y_dim}) not found.{_get_data_var_message(data_obj)}"
            )
        return data_obj

    @property
    def x_dim(self) -> Hashable:
        """Hashable: The dimension for the X-axis."""
        if self._x_dim is not None:
            return self._x_dim
        raise MissingSpatialDimensionError(
            "x dimension not found. 'rio.set_spatial_dims()' or "
            "using 'rename()' to change the dimension name to 'x' can address this."
            f"{_get_data_var_message(self._obj)}"
        )

    @property
    def y_dim(self) -> Hashable:
        """Hashable: The dimension for the Y-axis."""
        if self._y_dim is not None:
            return self._y_dim
        raise MissingSpatialDimensionError(
            "y dimension not found. 'rio.set_spatial_dims()' or "
            "using 'rename()' to change the dimension name to 'y' can address this."
            f"{_get_data_var_message(self._obj)}"
        )

    @property
    def width(self) -> int:
        """int: Returns the width of the dataset (x dimension size)"""
        if self._width is not None:
            return self._width
        self._width = self._obj[self.x_dim].size
        return self._width

    @property
    def height(self) -> int:
        """int: Returns the height of the dataset (y dimension size)"""
        if self._height is not None:
            return self._height
        self._height = self._obj[self.y_dim].size
        return self._height

    @property
    def shape(self) -> tuple[int, int]:
        """tuple(int, int): Returns the shape (height, width)"""
        return (self.height, self.width)

    def _check_dimensions(self) -> Optional[str]:
        """
        This function validates that the dimensions 2D/3D and
        they are are in the proper order.

        Returns
        -------
        str or None: Name extra dimension.
        """
        extra_dims = tuple(set(list(self._obj.dims)) - {self.x_dim, self.y_dim})
        if len(extra_dims) > 1:
            raise TooManyDimensions(
                "Only 2D and 3D data arrays supported."
                f"{_get_data_var_message(self._obj)}"
            )
        if extra_dims and self._obj.dims != (extra_dims[0], self.y_dim, self.x_dim):
            dim_info: tuple = (extra_dims[0], self.y_dim, self.x_dim)
            raise InvalidDimensionOrder(
                f"Invalid dimension order. Expected order: {dim_info}. "
                f"You can use `DataArray.transpose{dim_info}`"
                " to reorder your dimensions."
                f"{_get_data_var_message(self._obj)}"
            )
        if not extra_dims and self._obj.dims != (self.y_dim, self.x_dim):
            dim_info = (self.y_dim, self.x_dim)
            raise InvalidDimensionOrder(
                f"Invalid dimension order. Expected order: {dim_info}. "
                f"You can use `DataArray.transpose{dim_info}`"
                " to reorder your dimensions."
                f"{_get_data_var_message(self._obj)}"
            )
        return str(extra_dims[0]) if extra_dims else None

    @property
    def count(self) -> int:
        """int: Returns the band count (z dimension size)"""
        if self._count is not None:
            return self._count
        extra_dim = self._check_dimensions()
        self._count = 1
        if extra_dim is not None:
            self._count = self._obj[extra_dim].size
        return self._count

    def _internal_bounds(self) -> tuple[float, float, float, float]:
        """Determine the internal bounds of the `xarray.DataArray`"""
        if self.x_dim not in self._obj.coords:
            raise DimensionMissingCoordinateError(f"{self.x_dim} missing coordinates.")
        if self.y_dim not in self._obj.coords:
            raise DimensionMissingCoordinateError(f"{self.y_dim} missing coordinates.")
        try:
            left = float(self._obj[self.x_dim][0])
            right = float(self._obj[self.x_dim][-1])
            top = float(self._obj[self.y_dim][0])
            bottom = float(self._obj[self.y_dim][-1])
        except IndexError:
            raise NoDataInBounds(
                "Unable to determine bounds from coordinates."
                f"{_get_data_var_message(self._obj)}"
            ) from None
        return left, bottom, right, top

    def resolution(self, recalc: bool = False) -> tuple[float, float]:
        """
        Determine if the resolution of the grid.
        If the transformation has rotation, the sign of the resolution is lost.

        Parameters
        ----------
        recalc: bool, optional
            Will force the resolution to be recalculated instead of using the
            transform attribute.

        Returns
        -------
        x_resolution, y_resolution: float
            The resolution of the `xarray.DataArray` | `xarray.Dataset`
        """
        transform = self._cached_transform()

        if (
            not recalc or self.width == 1 or self.height == 1
        ) and transform is not None:
            return _resolution(transform)

        # if the coordinates of the spatial dimensions are missing
        # use the cached transform resolution
        try:
            left, bottom, right, top = self._internal_bounds()
        except DimensionMissingCoordinateError:
            if transform is None:
                raise
            return _resolution(transform)

        if self.width == 1 or self.height == 1:
            raise OneDimensionalRaster(
                "Only 1 dimenional array found. Cannot calculate the resolution."
                f"{_get_data_var_message(self._obj)}"
            )

        resolution_x = (right - left) / (self.width - 1)
        resolution_y = (bottom - top) / (self.height - 1)
        return resolution_x, resolution_y

    def _unordered_bounds(
        self, recalc: bool = False
    ) -> tuple[float, float, float, float]:
        """
        Unordered bounds.

        Parameters
        ----------
        recalc: bool, optional
            Will force the bounds to be recalculated instead of using the
            transform attribute.

        Returns
        -------
        left, bottom, right, top: float
            Outermost coordinates of the `xarray.DataArray` | `xarray.Dataset`.
        """
        resolution_x, resolution_y = self.resolution(recalc=recalc)

        try:
            # attempt to get bounds from xarray coordinate values
            left, bottom, right, top = self._internal_bounds()
            left -= resolution_x / 2.0
            right += resolution_x / 2.0
            top -= resolution_y / 2.0
            bottom += resolution_y / 2.0
        except DimensionMissingCoordinateError as error:
            transform = self._cached_transform()
            if not transform:
                raise RioXarrayError("Transform not able to be determined.") from error
            left = transform.c
            top = transform.f
            right = left + resolution_x * self.width
            bottom = top + resolution_y * self.height

        return left, bottom, right, top

    def bounds(self, *, recalc: bool = False) -> tuple[float, float, float, float]:
        """
        Parameters
        ----------
        recalc: bool, optional
            Will force the bounds to be recalculated instead of using the
            transform attribute.

        Returns
        -------
        left, bottom, right, top: float
            Outermost coordinates of the `xarray.DataArray` | `xarray.Dataset`.
        """
        minx, miny, maxx, maxy = self._unordered_bounds(recalc=recalc)
        resolution_x, resolution_y = self.resolution(recalc=recalc)
        return _order_bounds(
            minx=minx,
            miny=miny,
            maxx=maxx,
            maxy=maxy,
            resolution_x=resolution_x,
            resolution_y=resolution_y,
        )

    def isel_window(
        self, window: rasterio.windows.Window, *, pad: bool = False
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Use a rasterio.windows.Window to select a subset of the data.

        .. versionadded:: 0.6.0 pad

        .. warning:: Float indices are converted to integers.

        Parameters
        ----------
        window: :class:`rasterio.windows.Window`
            The window of the dataset to read.
        pad: bool, default=False
            Set to True to expand returned DataArray to dimensions of the window

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            The data in the window.
        """
        (row_start, row_stop), (col_start, col_stop) = window.toranges()
        row_start = 0 if row_start < 0 else math.floor(row_start)
        row_stop = 0 if row_stop < 0 else math.ceil(row_stop)
        col_start = 0 if col_start < 0 else math.floor(col_start)
        col_stop = 0 if col_stop < 0 else math.ceil(col_stop)
        row_slice = slice(int(row_start), int(row_stop))
        col_slice = slice(int(col_start), int(col_stop))
        array_subset = (
            self._obj.isel({self.y_dim: row_slice, self.x_dim: col_slice})
            .copy()  # this is to prevent sharing coordinates with the original dataset
            .rio.set_spatial_dims(x_dim=self.x_dim, y_dim=self.y_dim, inplace=True)
            .rio.write_transform(
                transform=rasterio.windows.transform(
                    rasterio.windows.Window.from_slices(
                        rows=row_slice,
                        cols=col_slice,
                        width=self.width,
                        height=self.height,
                    ),
                    self.transform(recalc=True),
                ),
                inplace=True,
            )
        )
        if pad:
            return array_subset.rio.pad_box(
                *rasterio.windows.bounds(window, self.transform(recalc=True))
            )
        return array_subset

    def slice_xy(
        self,
        minx: float,
        miny: float,
        maxx: float,
        maxy: float,
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """Slice the array by x,y bounds.

        Parameters
        ----------
        minx: float
            Minimum bound for x coordinate.
        miny: float
            Minimum bound for y coordinate.
        maxx: float
            Maximum bound for x coordinate.
        maxy: float
            Maximum bound for y coordinate.


        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            The data in the slice.
        """
        left, bottom, right, top = self._internal_bounds()
        if top > bottom:
            y_slice = slice(maxy, miny)
        else:
            y_slice = slice(miny, maxy)

        if left > right:
            x_slice = slice(maxx, minx)
        else:
            x_slice = slice(minx, maxx)

        subset = (
            self._obj.sel({self.x_dim: x_slice, self.y_dim: y_slice})
            .copy()  # this is to prevent sharing coordinates with the original dataset
            .rio.set_spatial_dims(x_dim=self.x_dim, y_dim=self.y_dim, inplace=True)
            .rio.write_transform(inplace=True)
        )
        return subset

    def transform_bounds(
        self, dst_crs: Any, *, densify_pts: int = 21, recalc: bool = False
    ) -> tuple[float, float, float, float]:
        """Transform bounds from src_crs to dst_crs.

        Optionally densifying the edges (to account for nonlinear transformations
        along these edges) and extracting the outermost bounds.

        Note: this does not account for the antimeridian.

        Parameters
        ----------
        dst_crs: str, :obj:`rasterio.crs.CRS`, or dict
            Target coordinate reference system.
        densify_pts: uint, optional
            Number of points to add to each edge to account for nonlinear
            edges produced by the transform process.  Large numbers will produce
            worse performance.  Default: 21 (gdal default).
        recalc: bool, optional
            Will force the bounds to be recalculated instead of using the transform
            attribute.

        Returns
        -------
        left, bottom, right, top: float
            Outermost coordinates in target coordinate reference system.
        """
        return rasterio.warp.transform_bounds(
            self.crs, dst_crs, *self.bounds(recalc=recalc), densify_pts=densify_pts
        )

    def write_gcps(
        self,
        gcps: Iterable[GroundControlPoint],
        gcp_crs: Any,
        *,
        grid_mapping_name: Optional[str] = None,
        inplace: bool = False,
    ) -> Union[xarray.Dataset, xarray.DataArray]:
        """
        Write the GroundControlPoints to the dataset.

        https://rasterio.readthedocs.io/en/latest/topics/georeferencing.html#ground-control-points

        Parameters
        ----------
        gcp: list of :obj:`rasterio.control.GroundControlPoint`
            The Ground Control Points to integrate to the dataset.
        gcp_crs: str, :obj:`rasterio.crs.CRS`, or dict
            Coordinate reference system for the GCPs.
        grid_mapping_name: str, optional
            Name of the grid_mapping coordinate to store the GCPs information in.
            Default is the grid_mapping name of the dataset.
        inplace: bool, optional
            If True, it will write to the existing dataset. Default is False.

        Returns
        -------
        :obj:`xarray.Dataset` | :obj:`xarray.DataArray`:
            Modified dataset with Ground Control Points written.
        """
        grid_mapping_name = (
            self.grid_mapping if grid_mapping_name is None else grid_mapping_name
        )
        data_obj = self._get_obj(inplace=True)

        if gcp_crs:
            data_obj = data_obj.rio.write_crs(
                gcp_crs, grid_mapping_name=grid_mapping_name, inplace=inplace
            )
        try:
            grid_map_attrs = data_obj.coords[grid_mapping_name].attrs.copy()
        except KeyError:
            data_obj.coords[grid_mapping_name] = xarray.Variable((), 0)
            grid_map_attrs = data_obj.coords[grid_mapping_name].attrs.copy()
        geojson_gcps = _convert_gcps_to_geojson(gcps)
        grid_map_attrs["gcps"] = json.dumps(geojson_gcps)
        data_obj.coords[grid_mapping_name].rio.set_attrs(grid_map_attrs, inplace=True)
        self._gcps = list(gcps)
        return data_obj

    def get_gcps(self) -> Optional[list[GroundControlPoint]]:
        """
        Get the GroundControlPoints from the dataset.

        https://rasterio.readthedocs.io/en/latest/topics/georeferencing.html#ground-control-points

        Returns
        -------
        list of :obj:`rasterio.control.GroundControlPoint` or None
            The Ground Control Points from the dataset or None if not applicable
        """
        if self._gcps is not None:
            return self._gcps
        try:
            geojson_gcps = json.loads(self._obj.coords[self.grid_mapping].attrs["gcps"])
        except (KeyError, AttributeError):
            return None

        def _parse_gcp(gcp) -> GroundControlPoint:
            x, y, *z = gcp["geometry"]["coordinates"]
            z = z[0] if z else None
            return GroundControlPoint(
                x=x,
                y=y,
                z=z,
                row=gcp["properties"]["row"],
                col=gcp["properties"]["col"],
                id=gcp["properties"]["id"],
                info=gcp["properties"]["info"],
            )

        self._gcps = [_parse_gcp(gcp) for gcp in geojson_gcps["features"]]
        return self._gcps


def _convert_gcps_to_geojson(
    gcps: Iterable[GroundControlPoint],
) -> dict:
    """
    Convert GCPs to geojson.

    Parameters
    ----------
    gcps: The list of GroundControlPoint instances.

    Returns
    -------
    A FeatureCollection dict.
    """

    def _gcp_coordinates(gcp):
        if gcp.z is None:
            return [gcp.x, gcp.y]
        return [gcp.x, gcp.y, gcp.z]

    features = [
        {
            "type": "Feature",
            "properties": {
                "id": gcp.id,
                "info": gcp.info,
                "row": gcp.row,
                "col": gcp.col,
            },
            "geometry": {"type": "Point", "coordinates": _gcp_coordinates(gcp)},
        }
        for gcp in gcps
    ]
    return {"type": "FeatureCollection", "features": features}