# -*- coding: utf-8 -*-
"""Scripts to derive (sub)basin geometries.

Based on pre-cooked basin index files, basin maps or flow direction maps.
"""

import logging
import warnings
from os.path import isdir, isfile
from pathlib import Path

import geopandas as gpd
import numpy as np
import xarray as xr
from shapely.geometry import box

from .. import _compat
from ..data_adapter import GeoDataFrameAdapter
from ..data_catalog import DataCatalog
from ..flw import basin_map, flwdir_from_da, outlet_map, stream_map

# local
from ..models import MODELS

if _compat.HAS_XUGRID:
    import xugrid as xu

logger = logging.getLogger(__name__)

__all__ = ["get_basin_geometry", "parse_region"]


def parse_region(region, logger=logger, data_catalog=None):
    """Check and return parsed region arguments.

    Parameters
    ----------
    region : dict
        Dictionary describing region of interest.

        For an exact clip of the region:

        * {'bbox': [xmin, ymin, xmax, ymax]}

        * {'geom': /path/to/polygon_geometry}

        For a region based of another models grid:

        * {'<model_name>': root}

        For a region based of the grid of a raster file:

        * {'grid': /path/to/raster}

        For a region based on a mesh grid of a mesh file:

        * {'mesh': /path/to/mesh}

        Entire basin can be defined based on an ID, one or multiple point location
        (x, y), or a region of interest (bounding box or geometry) for which the
        basin IDs are looked up. The basins withint the area of interest can be further
        filtered to only include basins with their outlet within the area of interest
        ('outlets': true) of stream threshold arguments (e.g.: 'uparea': 1000).

        Common use-cases include:

        * {'basin': ID}

        * {'basin': [ID1, ID2, ..]}

        * {'basin': [x, y]}

        * {'basin': [[x1, x2, ..], [y1, y2, ..]]}

        * {'basin': /path/to/point_geometry}

        * {'basin': [xmin, ymin, xmax, ymax]}

        * {'basin': [xmin, ymin, xmax, ymax], 'outlets': true}

        * {'basin': [xmin, ymin, xmax, ymax], '<variable>': threshold}

        Subbasins are defined by its outlet locations and include all area upstream
        from these points. The outlet locations can be passed as xy coordinate pairs,
        but also derived from the most downstream cell(s) within a area of interest
        defined by a bounding box or geometry, optionally refined by stream threshold
        arguments.

        The method can be speed up by providing an additional ``bounds`` argument which
        should contain all upstream cell. If cells upstream of the subbasin are not
        within the provide bounds a warning will be raised. Common use-cases include:

        * {'subbasin': [x, y], '<variable>': threshold}

        * {
            'subbasin': [[x1, x2, ..], [y1, y2, ..]],
            '<variable>': threshold, 'bounds': [xmin, ymin, xmax, ymax]
            }

        * {'subbasin': /path/to/point_geometry, '<variable>': threshold}

        * {'subbasin': [xmin, ymin, xmax, ymax], '<variable>': threshold}

        * {'subbasin': /path/to/polygon_geometry, '<variable>': threshold}

        Interbasins are similar to subbasins but are bounded by a bounding box or
        geometry and do not include all upstream area. Common use-cases include:

        * {'interbasin': [xmin, ymin, xmax, ymax], '<variable>': threshold}

        * {'interbasin': [xmin, ymin, xmax, ymax], 'xy': [x, y]}

        * {'interbasin': /path/to/polygon_geometry, 'outlets': true}
    logger:
        The logger to use.

    Returns
    -------
    kind : {'basin', 'subbasin', 'interbasin', 'geom', 'bbox', 'grid'}
        region kind
    kwargs : dict
        parsed region json
    """
    if data_catalog is None:
        data_catalog = DataCatalog()
    kwargs = region.copy()
    # NOTE: the order is important to prioritize the arguments
    options = {
        "basin": ["basid", "geom", "bbox", "xy"],
        "subbasin": ["geom", "bbox", "xy"],
        "interbasin": ["geom", "bbox", "xy"],  # FIXME remove interbasin & xy combi?
        "outlet": ["geom", "bbox"],  # deprecated!
        "geom": ["geom"],
        "bbox": ["bbox"],
        "grid": ["RasterDataArray"],
        "mesh": ["UgridDataArray"],
    }
    kind = next(iter(kwargs))  # first key of region
    value0 = kwargs.pop(kind)
    if kind in MODELS:
        model_class = MODELS.load(kind)
        kwargs = dict(mod=model_class(root=value0, mode="r", logger=logger))
        kind = "model"
    elif kind == "grid":
        kwargs = {"grid": data_catalog.get_rasterdataset(value0, driver_kwargs=kwargs)}
    elif kind == "mesh":
        if _compat.HAS_XUGRID:
            if isinstance(value0, (str, Path)) and isfile(value0):
                kwarg = dict(mesh=xu.open_dataset(value0))
            elif isinstance(value0, (xu.UgridDataset, xu.UgridDataArray)):
                kwarg = dict(mesh=value0)
            elif isinstance(value0, (xu.Ugrid1d, xu.Ugrid2d)):
                kwarg = dict(
                    mesh=xu.UgridDataset(value0.to_dataset(optional_attributes=True))
                )
            else:
                raise ValueError(
                    f"Unrecognised type {type(value0)}."
                    "Should be a path, data catalog key or xugrid object."
                )
            kwargs.update(kwarg)
        else:
            raise ImportError("xugrid is required to read mesh files.")
    elif kind not in options:
        k_lst = '", "'.join(list(options.keys()) + list(MODELS))
        raise ValueError(f'Region key "{kind}" not understood, select from "{k_lst}"')
    else:
        kwarg = _parse_region_value(value0, data_catalog=data_catalog)
        if len(kwarg) == 0 or next(iter(kwarg)) not in options[kind]:
            v_lst = '", "'.join(list(options[kind]))
            raise ValueError(
                f'Region value "{value0}" for kind={kind} not understood, '
                f'provide one of "{v_lst}"'
            )
        kwargs.update(kwarg)
    kwargs_str = dict()
    for k, v in kwargs.items():
        if isinstance(v, gpd.GeoDataFrame):
            v = f"GeoDataFrame {v.total_bounds} (crs = {v.crs})"
        elif isinstance(v, xr.DataArray):
            v = f"DataArray {v.raster.bounds} (crs = {v.raster.crs})"
        kwargs_str.update({k: v})
    logger.debug(f"Parsed region (kind={kind}): {str(kwargs_str)}")
    return kind, kwargs


def _parse_region_value(value, data_catalog):
    kwarg = {}
    if isinstance(value, np.ndarray):
        value = value.tolist()  # array to list

    if isinstance(value, list):
        if np.all([isinstance(p0, int) and abs(p0) > 180 for p0 in value]):  # all int
            kwarg = dict(basid=value)
        elif len(value) == 4:  # 4 floats
            kwarg = dict(bbox=value)
        elif len(value) == 2:  # 2 floats
            kwarg = dict(xy=value)
    elif isinstance(value, tuple) and len(value) == 2:  # tuple of x and y coords
        kwarg = dict(xy=value)
    elif isinstance(value, int):  # single int
        kwarg = dict(basid=value)
    elif isinstance(value, (str, Path)) and isdir(value):
        kwarg = dict(root=value)
    elif isinstance(value, (str, Path)):
        geom = data_catalog.get_geodataframe(value)
        kwarg = dict(geom=geom)
    elif isinstance(value, gpd.GeoDataFrame):  # geometry
        kwarg = dict(geom=value)
    else:
        raise ValueError(f"Region value {value} not understood.")

    if "geom" in kwarg and np.all(kwarg["geom"].geometry.type == "Point"):
        xy = (
            kwarg["geom"].geometry.x.values,
            kwarg["geom"].geometry.y.values,
        )
        kwarg = dict(xy=xy)
    return kwarg


def _check_size(ds, logger=logger, threshold=12e3**2):
    # warning for large domain
    if (
        np.multiply(*ds.raster.shape) > threshold
    ):  # 12e3 ** 2 > 10x10 degree at 3 arcsec
        logger.warning(
            "Loading very large spatial domain to derive a subbasin. "
            "Provide initial 'bounds' if this takes too long."
        )


def get_basin_geometry(
    ds,
    basin_index=None,
    kind="basin",
    bounds=None,
    bbox=None,
    geom=None,
    xy=None,
    basid=None,
    outlets=False,
    basins_name="basins",
    flwdir_name="flwdir",
    ftype="infer",
    logger=logger,
    buffer=10,
    **stream_kwargs,
):
    """Return a geometry of the (sub)(inter)basin(s).

    This method derives a geometry of sub-, inter- or full basin based on an input
    dataset with flow-direction and optional basins ID raster data in combination
    with a matching basin geometry file containing the bounding boxes of each basin.

    Either ``bbox``, ``geom``, ``xy`` (or ``basid`` in case of ``kind='basin'``) must
    be provided.

    Parameters
    ----------
    ds : xarray.Dataset
        dataset containing basin and flow direction variables
    basin_index: geopandas.GeoDataFrame or GeoDataFrameAdapter
        Dataframe with basin geomtries or bounding boxes with "basid" column
        corresponding to the ``ds[<basins_name>]`` map.
    kind : {"basin", "subbasin", "interbasin"}
        kind of basin description
    bounds: array_like of float, optional
        [xmin, ymin, xmax, ymax] coordinates of total bounding box, i.e. the data is
        clipped to this domain before futher processing.
    bbox : array_like of float, optional
        [xmin, ymin, xmax, ymax] coordinates to infer (sub)(inter)basin(s)
    geom : geopandas.GeoDataFrame, optional
        polygon geometry describing area of interest
    xy : tuple of array_like of float, optional
        x, y coordinates of (sub)basin outlet locations
    basid : int or array_like of int, optional
        basin IDs, must match values in basin maps
    outlets: bool, optional
        If True, include (sub)basins of outlets within domain only.
    flwdir_name : str, optional
        Name of flow direction variable in source, by default "flwdir"
    basins_name : str, optional
        Name of flow direction variable in source, by default "basins"
    ftype : {'d8', 'ldd', 'nextxy'}, optional
        name of flow direction type, by default None; use input ftype.
    stream_kwargs : key-word arguments
        name of variable in ds and threshold value
    buffer:
        The buffer to apply.
    logger:
        The logger to use.

    Returns
    -------
    basin_geom : geopandas.geoDataFrame
        geometry the (sub)basin(s)
    outlet_geom : geopandas.geoDataFrame
        geometry the outlet point location
    """
    kind_lst = ["basin", "subbasin", "interbasin"]
    if kind == "outlet":
        outlets = True
        kind = "basin"
        warnings.warn(
            'kind="outlets" has been deprecated, use outlets=True in combination with '
            'kind="basin" or kind="interbasin" instead.',
            DeprecationWarning,
            stacklevel=2,
        )
    elif kind not in kind_lst:
        msg = f"Unknown kind: {kind}, select from {kind_lst}."
        raise ValueError(msg)
    if bool(stream_kwargs.pop("within", False)):
        warnings.warn(
            '"within" stream argument has been deprecated.',
            DeprecationWarning,
            stacklevel=2,
        )

    # check variables
    dvars = [flwdir_name] + [v for v in stream_kwargs]
    for name in dvars:
        if name not in ds.data_vars:
            raise ValueError(f"Dataset variable {name} not in ds.")

    # for interbasins we can limit the domain based on either bbox / geom or bounds
    if kind == "interbasin" and bounds is None:
        if bbox is None and geom is None:
            raise ValueError('"kind=interbasin" requires either "bbox" or "geom"')
        bounds = bbox if bbox is not None else geom.total_bounds
    # initial clip based on bounds
    if bounds is not None:
        ds = ds.raster.clip_bbox(bounds, buffer=buffer)
    # convert bbox to geom
    if geom is None and bbox is not None:
        geom = gpd.GeoDataFrame(geometry=[box(*bbox)], crs=ds.raster.crs)

    # check basin index
    # TODO understand pfafstetter codes
    gdf_bas = None
    if basin_index is not None:
        if isinstance(basin_index, GeoDataFrameAdapter):
            kwargs = dict(variables=["basid"])
            if geom is not None:
                kwargs.update(geom=geom)
            elif xy is not None:
                xy0 = np.atleast_1d(xy[0])
                xy1 = np.atleast_1d(xy[1])
                kwargs.update(
                    bbox=[
                        min(xy0) - 0.1,
                        min(xy1) - 0.1,
                        max(xy0) + 0.1,
                        max(xy1) + 0.1,
                    ]
                )
            gdf_bas = basin_index.get_data(**kwargs)
        elif isinstance(basin_index, gpd.GeoDataFrame):
            gdf_bas = basin_index
            if "basid" not in gdf_bas.columns:
                raise ValueError("Basin geometries does not have 'basid' column.")
        if gdf_bas.crs != ds.raster.crs:
            logger.warning("Basin geometries CRS does not match the input raster CRS.")
            gdf_bas = gdf_bas.to_crs(ds.raster.crs)

    ## BASINS
    if kind == "basin" or bounds is None:
        dvars = dvars + [basins_name]
        if basins_name not in ds:
            if gdf_bas is not None:
                gdf_bas = None
                logger.warning(
                    "Basin geometries ignored as no corresponding"
                    + " basin map is provided."
                )
            _check_size(ds, logger)
            logger.info(f'basin map "{basins_name}" missing, calculating on the fly.')
            flwdir = flwdir_from_da(ds[flwdir_name], ftype=ftype)
            ds[basins_name] = xr.Variable(ds.raster.dims, flwdir.basins())
        elif (
            ds[basins_name].raster.nodata != 0
            and ds[basins_name].raster.nodata is not None
        ):
            ds[basins_name].where(ds[basins_name] != ds[basins_name].raster.nodata, 0)
        ds[basins_name].raster.set_nodata(0)
        # clip
        ds_clip = ds[dvars]
        if geom is not None:
            ds_clip = ds[dvars].raster.clip_geom(geom, buffer=buffer, mask=True)
        # get basin IDs
        if xy is not None:
            logger.debug("Getting basin IDs at point locations.")
            sel = {
                ds.raster.x_dim: xr.IndexVariable("xy", np.atleast_1d(xy[0])),
                ds.raster.y_dim: xr.IndexVariable("xy", np.atleast_1d(xy[1])),
            }
            basid = np.unique(ds_clip[basins_name].sel(**sel, method="nearest").values)
        elif basid is None:
            if stream_kwargs or outlets:
                if stream_kwargs:
                    stream = stream_map(ds_clip, **stream_kwargs)
                if outlets:
                    outmap = outlet_map(ds_clip[flwdir_name], ftype=ftype)
                    if stream_kwargs:
                        stream = stream.where(outmap, 0)
                    else:
                        stream = outmap
                ds_clip[basins_name] = ds_clip[basins_name].where(stream, 0)
            logger.debug("Getting IDs of intersecting basins.")
            basid = np.unique(ds_clip[basins_name].values)
        basid = np.atleast_1d(basid)
        basid = basid[basid > 0]
        if basid.size == 0:
            raise ValueError("No basins found with given criteria.")
        # clip ds to total basin
        if gdf_bas is not None:
            gdf_match = np.isin(gdf_bas["basid"], basid)
            gdf_bas = gdf_bas.loc[gdf_match]
            if gdf_bas.index.size > 0:
                if geom is not None:
                    xminbas, yminbas, xmaxbas, ymaxbas = gdf_bas.total_bounds
                    # Check that total_bounds is at least bigger
                    # than original geom bounds
                    xmingeom, ymingeom, xmaxgeom, ymaxgeom = geom.total_bounds
                    total_bounds = [
                        min(xminbas, xmingeom),
                        min(yminbas, ymingeom),
                        max(xmaxbas, xmaxgeom),
                        max(ymaxbas, ymaxgeom),
                    ]
                else:
                    total_bounds = gdf_bas.total_bounds
                ds = ds[dvars].raster.clip_bbox(total_bounds, buffer=0)
            elif np.any(gdf_match):
                logger.warning("No matching basin IDs found in basin geometries.")
        # get full basin mask and use this mask ds_basin incl flow direction raster
        _mask = np.isin(ds[basins_name], basid)
        if not np.any(_mask > 0):
            raise ValueError(f"No basins found with IDs: {basid}")
        ds = ds[dvars].raster.mask(_mask)
        ds = ds.raster.clip_mask(ds[basins_name])
        bas_mask = ds[basins_name]

    # INTER- & SUBBASINS
    xy_out = None
    if kind in ["subbasin", "interbasin"]:
        # get flow directions
        _check_size(ds, logger)  # warning for large domain
        mask = False
        # if interbasin, set flwdir mask within geometry / bounding box
        if kind == "interbasin":  # we have checked before that geom is not None
            mask = ds.raster.geometry_mask(geom)
        elif basins_name in ds:  # set flwdir mask based on basin map
            mask = ds[basins_name] > 0
        flwdir = flwdir_from_da(ds[flwdir_name], ftype=ftype, mask=mask)
        # get area of interest (aoi) mask
        if geom is not None:
            aoi = ds.raster.geometry_mask(geom)
        else:
            aoi = xr.DataArray(
                coords=ds.raster.coords,
                dims=ds.raster.dims,
                data=np.full(ds.raster.shape, True, dtype=bool),
            )  # all True
        # Convert xy to tuple
        if xy is not None:
            xy = (np.atleast_1d(xy[0]), np.atleast_1d(xy[1]))
        # get stream mask. Always over entire domain
        # to include cells downstream of aoi!
        kwargs = dict()
        if stream_kwargs:
            stream = stream_map(ds, **stream_kwargs)
            if not np.any(stream):
                raise ValueError(f"No streams found with: {stream_kwargs}.")
            if not outlets and xy is None:  # get aoi outflow cells if none provided
                xy = flwdir.xy(flwdir.outflow_idxs(np.logical_and(stream, aoi).values))
            elif not outlets:
                kwargs.update(stream=stream.values)
        if outlets:
            outmap = aoi.where(outlet_map(ds[flwdir_name], ftype=flwdir.ftype), False)
            if stream_kwargs:
                outmap = outmap.where(stream_kwargs, False)
            idxs_out = np.where(outmap.values.ravel())[0]
            if not np.any(outmap):
                raise ValueError("No outlets found with with given criteria.")
            xy = outmap.raster.idx_to_xy(idxs_out)
        # get subbasin map
        bas_mask, xy_out = basin_map(ds, flwdir, xy, **kwargs)
        # is subbasin with bounds check if all upstream cells are included
        if kind == "subbasin" and bounds is not None:
            geom = gpd.GeoDataFrame(geometry=[box(*bounds)], crs=ds.raster.crs)
            mask = ds.raster.geometry_mask(geom)
            if np.any(np.logical_and(mask == 0, bas_mask != 0)):
                logger.warning("The subbasin does not include all upstream cells.")

    if not np.any(bas_mask > 0):
        raise ValueError(f"No {kind} found with given criteria.")
    bas_mask = bas_mask.astype(np.int32)
    bas_mask.raster.set_crs(ds.raster.crs)
    bas_mask.raster.set_nodata(0)
    w, s, e, n = bas_mask.raster.clip_mask(bas_mask).raster.bounds
    logger.info(f"{kind} bbox: [{w:.4f}, {s:.4f}, {e:.4f}, {n:.4f}]")

    # vectorize basins and outlets
    basin_geom = bas_mask.raster.vectorize()
    outlet_geom = None
    if xy_out is not None:
        points = gpd.points_from_xy(*xy_out)
        outlet_geom = gpd.GeoDataFrame(geometry=points, crs=ds.raster.crs)

    return basin_geom, outlet_geom
