import affine
import numpy as np
import xarray as xr

import imod
from imod.util.imports import MissingOptionalModule

# since rasterio is a big dependency that is sometimes hard to install
# and not always required, we made this an optional dependency
try:
    import rasterio
    import rasterio.warp
except ImportError:
    rasterio = MissingOptionalModule("rasterio")


def _reproject_dst(source, src_crs, dst_crs, src_transform):
    """
    Prepares destination transform Affine and DataArray for projection.
    """
    src_height, src_width = source.y.size, source.x.size
    bounds = rasterio.transform.array_bounds(src_height, src_width, src_transform)
    dst_transform, dst_width, dst_height = rasterio.warp.calculate_default_transform(
        src_crs, dst_crs, src_width, src_height, *bounds
    )
    # from: http://xarray.pydata.org/en/stable/generated/xarray.open_rasterio.html
    x, y = dst_transform * np.meshgrid(
        np.arange(dst_width) + 0.5, np.arange(dst_height) + 0.5
    )
    dst = xr.DataArray(
        data=np.zeros((dst_height, dst_width), source.dtype),
        coords={"y": y[:, 0], "x": x[0, :]},
        dims=("y", "x"),
    )
    return dst_transform, dst


def reproject(
    source,
    like=None,
    src_crs=None,
    dst_crs=None,
    method="nearest",
    use_src_attrs=False,
    src_nodata=np.nan,
    **reproject_kwargs,
):
    """
    Reprojects and/or resamples a 2D xarray DataArray to a
    different cellsize or coordinate system.

    * To resample to a new cellsize in the same projection: provide only ``like``.
    * To only reproject: provide only ``src_crs`` and ``src_crs``.
    * To reproject and resample to a specific domain: provide ``src_crs``, ``src_crs``, and ``like``.

    Note: when only ``like`` is provided, Cartesian (projected) coordinates are a
    ssumed for resampling. In case of non-Cartesian coordinates, specify
    ``src_crs`` and ``dst_crs`` for correct resampling.

    Parameters
    ----------
    source: xarray DataArray
        The DataArray to be resampled and/or reprojected. Must contain dimensions
        ``y`` and ``x``.
    like: xarray DataArray
        Example DataArray that shows what the resampled result should look like
        in terms of coordinates. Must contain dimensions ``y`` and ``x``.
    src_crs: string, dict, rasterio.crs.CRS
        Coordinate system of ``source``. Options:

        * string: e.g. ``"EPSG:4326"``
        * rasterio.crs.CRS
    dst_crs: string, dict, rasterio.crs.CRS
        Coordinate system of result. Options:

        * string: e.g. ``"EPSG:4326"``
        * rasterio.crs.CRS
    use_src_attrs: boolean
        If True: Use metadata in ``source.attrs``, as generated by ``xarray.open_rasterio()``, to do
        reprojection.
    method: string
        The method to use for resampling/reprojection.
        Defaults to "nearest". GDAL methods are available:

        * nearest
        * bilinear
        * cubic
        * cubic_spline
        * lanczos
        * average
        * mode
        * gauss
        * max
        * min
        * med (50th percentile)
        * q1 (25th percentile)
        * q3 (75th percentile)
    reproject_kwargs: dict, optional
        keyword arguments for ``rasterio.warp.reproject()``.

    Returns
    -------
    xarray.DataArray
        Resampled/reprojected DataArray.

    Examples
    --------
    Resample a DataArray ``a`` to a new cellsize, using an existing DataArray ``b``:

    >>> c = imod.rasterio.reproject(source=a, like=b)

    Resample a DataArray to a new cellsize of 100.0, by creating a ``like`` DataArray first:
    (Note that dy must be negative, as is usual for geospatial grids.)

    >>> dims = ("y", "x")
    >>> coords = {"y": np.arange(200_000.0, 100_000.0, -100.0), "x": np.arange(0.0, 100_000.0, 100.0)}
    >>> b = xr.DataArray(data=np.empty((200, 100)), coords=coords, dims=dims)
    >>> c = imod.rasterio.reproject(source=a, like=b)

    Reproject a DataArray from one coordinate system (WGS84, EPSG:4326) to another (UTM30N, EPSG:32630):

    >>> c = imod.rasterio.reproject(source=a, src_crs="EPSG:4326", dst_crs="EPSG:32630")

    Get the reprojected DataArray in the desired shape and coordinates by providing ``like``:

    >>> c = imod.rasterio.reproject(source=a, like=b, src_crs="EPSG:4326", dst_crs="EPSG:32630")

    Open a single band raster, and reproject to RD new coordinate system (EPSG:28992), without explicitly specifying ``src_crs``.
    ``src_crs`` is taken from ``a.attrs``, so the raster file has to include coordinate system metadata for this to work.

    >>> a = rioxarray.open_rasterio("example.tif").squeeze("band")
    >>> c = imod.rasterio.reproject(source=a, use_src_attrs=True, dst_crs="EPSG:28992")

    In case of a rotated ``source``, provide ``src_transform`` directly or ``use_src_attrs=True`` to rely on generated attributes:

    >>> rotated = rioxarray.open_rasterio("rotated_example.tif").squeeze("band")
    >>> c = imod.rasterio.reproject(source=rotated, dst_crs="EPSG:28992", reproject_kwargs={"src_transform":affine.Affine(...)})
    >>> c = imod.rasterio.reproject(source=rotated, dst_crs="EPSG:28992", use_src_attrs=True)
    """
    # Make sure the rio accessor is avaible.
    import rioxarray  # noqa pylint: F401

    if not source.dims == ("y", "x"):
        raise ValueError(
            "reproject does not support dimensions other than ``x`` and ``y`` for ``source``."
        )
    if like is not None:
        if not like.dims == ("y", "x"):
            raise ValueError(
                "reproject does not support dimensions other than ``x`` and ``y`` for ``like``."
            )
    if use_src_attrs:  # only provided when reproject is necessary
        try:
            src_crs = source.attrs["crs"]
        except KeyError:
            src_crs = source.rio.crs

        if isinstance(src_crs, str):
            if "epsg:" in src_crs.lower():
                # Workaround for rioxarray.open_rasterio generation proj4 strings
                # https://github.com/mapbox/rasterio/issues/1809
                epsg_code = src_crs.lower().split("epsg:")[-1]
                src_crs = rasterio.crs.CRS.from_epsg(epsg_code)
            else:
                src_crs = rasterio.crs.CRS.from_string(src_crs)
        elif isinstance(src_crs, rasterio.crs.CRS):
            pass
        else:
            raise ValueError(
                f"Invalid src_crs: {src_crs}. Must be either str or rasterio.crs.CRS object"
            )

        if "nodatavals" in source.attrs:
            src_nodata = source.attrs["nodatavals"][0]
        else:
            rio_nodata = source.rio.nodata
            if rio_nodata is not None:
                src_nodata = rio_nodata

    resampling_methods = {e.name: e for e in rasterio.enums.Resampling}

    if isinstance(method, str):
        try:
            resampling_method = resampling_methods[method]
        except KeyError as e:
            raise ValueError(
                "Invalid resampling method. Available methods are: {}".format(
                    resampling_methods.keys()
                )
            ) from e
    elif isinstance(method, rasterio.enums.Resampling):
        resampling_method = method
    else:
        raise TypeError("method must be a string or rasterio.enums.Resampling")

    # Givens: source, like, method. No reprojection necessary.
    if src_crs is None and dst_crs is None:
        if like is None:
            raise ValueError(
                "If crs information is not provided, ``like`` must be provided."
            )
        if resampling_method == rasterio.enums.Resampling.nearest:
            # this can be handled with xarray
            # xarray 0.10.9 needs .compute()
            # see https://github.com/pydata/xarray/issues/2454
            return source.compute().reindex_like(like, method="nearest")
        else:
            # if no crs is defined, assume it should remain the same
            # in this case use UTM30, ESPG:32630, as a dummy value for GDAL
            # (Any projected coordinate system should suffice, Cartesian plane == Cartesian plane)
            dst = xr.DataArray(
                data=np.zeros(like.shape, source.dtype),
                coords={"y": like.y, "x": like.x},
                dims=("y", "x"),
            )
            src_crs = dst_crs = rasterio.crs.CRS.from_epsg(32630)
        src_transform = imod.util.spatial.transform(source)
        dst_transform = imod.util.spatial.transform(like)

    elif src_crs and dst_crs:
        if use_src_attrs:
            # TODO: modify if/when xarray uses affine by default for transform
            try:
                src_transform = affine.Affine(*source.attrs["transform"][:6])
            except KeyError:
                src_transform = source.rio.transform()
        elif "src_transform" in reproject_kwargs.keys():
            src_transform = reproject_kwargs.pop("src_transform")
        else:
            src_transform = imod.util.spatial.transform(source)

        # If no like is provided, just reproject to different coordinate system
        if like is None:
            dst_transform, dst = _reproject_dst(source, src_crs, dst_crs, src_transform)
        else:
            dst_transform = imod.util.spatial.transform(like)
            dst = xr.DataArray(
                data=np.zeros(like.shape, source.dtype),
                coords={"y": like.y, "x": like.x},
                dims=("y", "x"),
            )

    else:
        raise ValueError(
            "At least ``like``, or crs information for source and destination must be provided."
        )

    if not src_transform[0] > 0:
        raise ValueError("dx of 'source' must be positive")
    if not src_transform[4] < 0:
        raise ValueError("dy of 'source' must be negative")
    if not dst_transform[0] > 0:
        raise ValueError("dx of 'like' must be positive")
    if not dst_transform[4] < 0:
        raise ValueError("dy of 'like' must be negative")

    rasterio.warp.reproject(
        source.values,
        dst.values,
        src_transform=src_transform,
        dst_transform=dst_transform,
        src_crs=src_crs,
        dst_crs=dst_crs,
        resampling=resampling_method,
        src_nodata=src_nodata,
        **reproject_kwargs,
    )

    dst.attrs = source.attrs
    dst.attrs["transform"] = dst_transform
    dst.attrs["res"] = (abs(dst_transform[0]), abs(dst_transform[4]))
    dst.attrs["crs"] = dst_crs
    # TODO: what should be the type of "crs" field in attrs?
    # Doesn't actually matter for functionality
    # rasterio accepts string, dict, and CRS object
    return dst
