"""Basemap workflows for Wflow plugin."""

import logging
import math
from typing import Optional

import geopandas as gpd
import numpy as np
import pyflwdir
import xarray as xr
from hydromt import DataCatalog
from hydromt.gis import flw
from hydromt.model.processes.basin_mask import get_basin_geometry
from hydromt.model.processes.region import (
    _parse_region_value,
    parse_region_bbox,
    parse_region_geom,
)
from pyflwdir import core_conversion, core_d8, core_ldd

logger = logging.getLogger(f"hydromt.{__name__}")

__all__ = ["hydrography", "topography", "parse_region"]


def hydrography(
    ds: xr.Dataset,
    res: float,
    xy: Optional[gpd.GeoDataFrame] = None,
    upscale_method: str = "ihu",
    flwdir_name: str = "flwdir",
    uparea_name: str = "uparea",
    basins_name: str = "basins",
    strord_name: str = "strord",
    ftype: str = "infer",
):
    """Return hydrography maps (see list below) and FlwdirRaster object.

    Based on gridded flow direction and elevation data input.

    The output maps are:

    - flwdir : flow direction [-]
    - basins : basin map [-]
    - uparea : upstream area [km2]
    - strord : stream order [-]

    If the resolution is lower than the source resolution, the flow direction data is
    upscaled and river length and slope are based on subgrid flow paths and
    the following maps are added:

    - subare : contributing area to each subgrid outlet pixel \
(unit catchment area) [km2]
    - subelv : elevation at subgrid outlet pixel [m+REF]

    Parameters
    ----------
    ds : xarray.DataArray
        Dataset containing gridded flow direction and elevation data.
    res : float
        output resolution
    xy : geopandas.GeoDataFrame, optional
        Subbasin pits. Only required when upscaling a subbasin.
    upscale_method : {'ihu', 'eam', 'dmm'}
        Upscaling method for flow direction data, by default 'ihu', see [1]_
    ftype : {'d8', 'ldd', 'nextxy', 'nextidx', 'infer'}, optional
        name of flow direction type, infer from data if 'infer', by default is 'infer'
    flwdir_name, elevtn_name, uparea_name, basins_name, strord_name : str, optional
        Name of flow direction [-], elevation [m], upstream area [km2], basin index [-]
        and stream order [-] variables in ds

    Returns
    -------
    ds_out : xarray.DataArray
        Dataset containing gridded hydrography data
    flwdir_out : pyflwdir.FlwdirRaster
        Flow direction raster object.

    References
    ----------
    .. [1] Eilander et al. (2021). A hydrography upscaling method for scale-invariant \
parametrization of distributed hydrological models.
           Hydrology and Earth System Sciences, 25(9), 5287–5313. https://doi.org/10.5194/hess-25-5287-2021

    See Also
    --------
    pyflwdir.FlwdirRaster.upscale_flwdir

    """
    # TODO add check if flwdir in ds, calculate if not
    flwdir = None
    basins = None
    outidx = None
    if "mask" not in ds.coords and xy is None:
        ds.coords["mask"] = xr.Variable(
            dims=ds.raster.dims, data=np.ones(ds.raster.shape, dtype=bool)
        )
    elif "mask" not in ds.coords:
        # NOTE if no subbasin mask is provided calculate it here
        logger.debug(f"Delineate {xy[0].size} subbasin(s).")
        flwdir = flw.flwdir_from_da(ds[flwdir_name], ftype=ftype)
        basins = flwdir.basins(xy=xy).astype(np.int32)
        ds.coords["mask"].data = basins != 0
        if not np.any(ds.coords["mask"]):
            raise ValueError("Delineating subbasins not successful.")
    elif xy is not None:
        # NOTE: this mask is passed on from get_basin_geometry method
        logger.debug("Mask in dataset assumed to represent subbasins.")
    ncells = np.sum(ds["mask"].values)
    scale_ratio = int(np.round(res / ds.raster.res[0]))

    if ncells < 4:
        raise ValueError(
            "(Sub)basin at original resolution should at least consist of two cells on "
            f"each axis and the total number of cells is {ncells}. "
            "Consider using a larger domain or higher spatial resolution. "
            "For subbasin models, consider a (higher) threshold on for example "
            "upstream area or stream order to snap the outlet."
        )
    elif ncells < 100 and scale_ratio > 1:
        logger.warning(
            f"(Sub)basin at original resolution is small and has {ncells} cells. "
            "This may results in errors later when upscaling flow directions. "
            "If so, consider using a larger domain or higher spatial resolution. "
            "For subbasin models, consider a (higher) threshold on for example "
            "upstream area or stream order to snap the outlet."
        )
    else:
        logger.info(f"(Sub)basin at original resolution has {ncells} cells.")

    if scale_ratio > 1:  # upscale flwdir
        if flwdir is None:
            # NOTE initialize with mask is FALSE
            flwdir = flw.flwdir_from_da(ds[flwdir_name], ftype=ftype, mask=False)
        if xy is not None:
            logger.debug("Burn subbasin outlet in upstream area data.")
            if isinstance(xy, gpd.GeoDataFrame):
                assert xy.crs == ds.raster.crs
                xy = xy.geometry.x, xy.geometry.y
            idxs_pit = flwdir.index(*xy)
            flwdir.add_pits(idxs=idxs_pit)
            uparea = ds[uparea_name].values
            uparea.flat[idxs_pit] = uparea.max() + 1.0
            ds[uparea_name].data = uparea
        logger.info(
            f"Upscale flow direction data: {scale_ratio:d}x, {upscale_method} method."
        )
        da_flw, flwdir_out = flw.upscale_flwdir(
            ds,
            flwdir=flwdir,
            scale_ratio=scale_ratio,
            method=upscale_method,
            uparea_name=uparea_name,
            flwdir_name=flwdir_name,
        )
        da_flw.raster.set_crs(ds.raster.crs)
        # make sure x_out and y_out get saved
        ds_out = da_flw.to_dataset().reset_coords(["x_out", "y_out"])
        dims = ds_out.raster.dims
        # find pits within basin mask
        idxs_pit0 = flwdir_out.idxs_pit
        outlon = ds_out["x_out"].values.ravel()
        outlat = ds_out["y_out"].values.ravel()
        sel = {
            ds.raster.x_dim: xr.Variable("yx", outlon[idxs_pit0]),
            ds.raster.y_dim: xr.Variable("yx", outlat[idxs_pit0]),
        }
        outbas_pit = ds.coords["mask"].sel(sel, method="nearest").values
        # derive basins
        if np.any(outbas_pit != 0):
            idxs_pit = idxs_pit0[outbas_pit != 0]
            basins = flwdir_out.basins(idxs=idxs_pit).astype(np.int32)
            ds_out.coords["mask"] = xr.Variable(
                dims=ds_out.raster.dims, data=basins != 0, attrs=dict(_FillValue=0)
            )
        else:
            # This is a patch for basins which are clipped based on bbox or wrong geom
            mask_int = ds["mask"].astype(np.int8)
            mask_int.raster.set_nodata(-1)  # change nodata value
            ds_out.coords["mask"] = mask_int.raster.reproject_like(
                da_flw, method="nearest"
            ).astype(bool)
            basins = ds_out["mask"].values.astype(np.int32)
            logger.warning(
                "The basin delineation might be wrong as no original resolution outlets"
                " are found in the upscaled map."
            )
        ds_out[basins_name] = xr.Variable(dims, basins, attrs=dict(_FillValue=0))
        # calculate upstream area using subgrid ucat cell areas
        outidx = np.where(
            ds_out["mask"], da_flw.coords["idx_out"].values, flwdir_out._mv
        )
        subare = flwdir.ucat_area(outidx, unit="km2")[1]
        uparea = flwdir_out.accuflux(subare)
        attrs = dict(_FillValue=-9999, unit="km2")
        ds_out[uparea_name] = xr.Variable(dims, uparea, attrs=attrs).astype(np.float32)
        # NOTE: subgrid cella area is currently not used in wflow
        ds_out["subare"] = xr.Variable(dims, subare, attrs=attrs).astype(np.float32)
        if "elevtn" in ds:
            subelv = ds["elevtn"].values.flat[outidx]
            subelv = np.where(outidx >= 0, subelv, -9999)
            attrs = dict(_FillValue=-9999, unit="m+REF")
            ds_out["subelv"] = xr.Variable(dims, subelv, attrs=attrs)
        # initiate masked flow dir
        flwdir_out = flw.flwdir_from_da(
            ds_out[flwdir_name], ftype=flwdir.ftype, mask=True
        )
    else:
        # NO upscaling : source resolution equals target resolution
        # NOTE (re-)initialize with mask is TRUE
        ftype = flwdir.ftype if flwdir is not None and ftype == "infer" else ftype
        flwdir = flw.flwdir_from_da(ds[flwdir_name], ftype=ftype, mask=True)
        flwdir_out = flwdir
        ds_out = xr.DataArray(
            name=flwdir_name,
            data=flwdir_out.to_array(),
            coords=ds.raster.coords,
            dims=ds.raster.dims,
            attrs=dict(
                long_name=f"{ftype} flow direction",
                _FillValue=flwdir_out._core._mv,
            ),
        ).to_dataset()
        dims = ds_out.raster.dims
        ds_out.coords["mask"] = xr.Variable(
            dims=dims, data=flwdir_out.mask.reshape(flwdir_out.shape)
        )
        # copy data variables from source if available
        for dvar in [basins_name, uparea_name, strord_name]:
            if dvar in ds.data_vars:
                ds_out[dvar] = xr.where(
                    ds_out["mask"],
                    ds[dvar],
                    ds[dvar].dtype.type(ds[dvar].raster.nodata),
                )
                ds_out[dvar].attrs.update(ds[dvar].attrs)
        # basins
        if basins_name not in ds_out.data_vars:
            if basins is None:
                basins = flwdir_out.basins(idxs=flwdir_out.idxs_pit).astype(np.int32)
            ds_out[basins_name] = xr.Variable(dims, basins, attrs=dict(_FillValue=0))
        else:
            # make sure dtype in ds_out is np.int32
            ds_out[basins_name] = ds_out[basins_name].astype(np.int32)
        # upstream area
        if uparea_name not in ds_out.data_vars:
            uparea = flwdir_out.upstream_area("km2")  # km2
            attrs = dict(_FillValue=-9999, unit="km2")
            ds_out[uparea_name] = xr.Variable(dims, uparea, attrs=attrs).astype(
                np.float32
            )
        # cell area
        # NOTE: subgrid cella area is currently not used in wflow
        subare = ds.raster.area_grid() / 1e6  # km2
        attrs = dict(_FillValue=-9999, unit="km2")
        ds_out["subare"] = xr.Variable(dims, subare, attrs=attrs).astype(np.float32)
    # logging
    npits = flwdir_out.idxs_pit.size
    xy_pit = flwdir_out.xy(flwdir_out.idxs_pit[:5])
    xy_pit_str = ", ".join([f"({x:.5f},{y:.5f})" for x, y in zip(*xy_pit)])
    # stream order
    if strord_name not in ds_out.data_vars:
        logger.info("Derive stream order.")
        strord = flwdir_out.stream_order()
        ds_out[strord_name] = xr.Variable(dims, strord)
        ds_out[strord_name].raster.set_nodata(255)

    # clip to basin extent
    ds_out = ds_out.raster.clip_mask(da_mask=ds_out[basins_name])
    # also mask idx_out coords if present
    if "idx_out" in ds_out:
        ds_out["idx_out"] = ds_out["idx_out"].where(
            ds_out[basins_name], ds_out["idx_out"].raster.nodata
        )

    ds_out.raster.set_crs(ds.raster.crs)
    logger.info(f"Map shape: {ds_out.raster.shape}; active cells: {flwdir_out.ncells}.")
    logger.info(f"Outlet coordinates ({len(xy_pit[0])}/{npits}): {xy_pit_str}.")
    if np.any(np.asarray(ds_out.raster.shape) == 1):
        raise ValueError(
            "The output extent at model resolution should at least consist of two "
            "cells on each axis. Consider using a larger domain or higher spatial "
            "resolution. For subbasin models, consider a (higher) threshold to snap "
            "the outlet."
        )

    # Convert flow direction from d8 to ldd format
    flwdir_data = ds_out["flwdir"].values.astype(np.uint8)  # force dtype
    # if d8 convert to ldd
    if core_d8.isvalid(flwdir_data):
        data = core_conversion.d8_to_ldd(flwdir_data)
        da_flwdir = xr.DataArray(
            name="flwdir",
            data=data,
            coords=ds_out.raster.coords,
            dims=ds_out.raster.dims,
            attrs=dict(
                long_name="ldd flow direction",
                _FillValue=core_ldd._mv,
            ),
        )
        ds_out["flwdir"] = da_flwdir
    return ds_out, flwdir_out


def topography(
    ds: xr.Dataset,
    ds_like: xr.Dataset,
    elevtn_name: str = "elevtn",
    lndslp_name: str = "lndslp",
    method: str = "average",
):
    """Return topography maps (see list below) at model resolution.

    Based on gridded elevation data input.

    The following topography maps are calculated:

    - elevtn : average elevation [m]
    - lndslp : average land surface slope [m/m]

    Parameters
    ----------
    ds : xarray.DataArray
        Dataset containing gridded flow direction and elevation data.
    ds_like : xarray.DataArray
        Dataset at model resolution.
    elevtn_name, lndslp_name : str, optional
        Name of elevation [m] and land surface slope [m/m] variables in ds
    method: str, optional
        Resample method used to reproject the input data, by default "average"

    Returns
    -------
    ds_out : xarray.DataArray
        Dataset containing gridded hydrography data
    flwdir_out : pyflwdir.FlwdirRaster
        Flow direction raster object.

    See Also
    --------
    pyflwdir.dem.slope
    """
    logger.info("Derive elevation and slope.")
    if lndslp_name not in ds.data_vars:
        logger.info(f"Slope map {lndslp_name} not found: derive from elevation map.")
        crs = ds[elevtn_name].raster.crs
        nodata = ds[elevtn_name].raster.nodata
        ds[lndslp_name] = xr.Variable(
            dims=ds.raster.dims,
            data=pyflwdir.dem.slope(
                elevtn=ds[elevtn_name].values,
                nodata=nodata,
                latlon=crs is not None and crs.to_epsg() == 4326,
                transform=ds[elevtn_name].raster.transform,
            ),
        )
        ds[lndslp_name].raster.set_nodata(nodata)
    # clip or reproject if non-identical grids
    ds_out = ds[[elevtn_name, lndslp_name]].raster.reproject_like(ds_like, method)
    # Mask possible negative values for slope
    ds_out["lndslp"] = np.maximum(ds_out["lndslp"], 0.0)
    ds_out[elevtn_name].attrs.update(unit="m")
    ds_out[lndslp_name].attrs.update(unit="m.m-1")
    return ds_out


def parse_region(
    data_catalog: DataCatalog,
    region: dict,
    hydrography_fn: str | xr.Dataset,
    resolution: float | int = 1 / 120.0,
    basin_index_fn: str | xr.Dataset | None = None,
) -> tuple[dict[str, gpd.GeoDataFrame], Optional[np.ndarray], xr.Dataset]:
    """Parse the region dictionary to get basin geometry and coordinates."""
    ds_org = data_catalog.get_rasterdataset(hydrography_fn)
    if ds_org is None:
        raise ValueError(f"hydrography_fn {hydrography_fn} not found in data catalog.")

    # Check on resolution
    hydrography_resolution = ds_org.raster.res[0]
    scale_ratio = resolution / hydrography_resolution

    if math.isclose(scale_ratio, 1):
        pass
    elif scale_ratio < 0.75:
        raise ValueError(
            f"Model resolution {resolution} should be larger than the "
            f"{hydrography_fn} resolution {hydrography_resolution}."
        )
    elif 0.75 < scale_ratio < 1.25:
        logger.warning(
            f"Model resolution {resolution} does not match the hydrography "
            f"resolution {hydrography_resolution}. This might lead to unexpected "
            f"results, using hydrography resolution instead: "
            f"{hydrography_resolution}."
        )
        resolution = hydrography_resolution
    else:
        if ds_org.raster.crs.is_geographic:
            if resolution > 1:  # 111 km
                raise ValueError(
                    f"The model resolution {resolution} should be smaller than 1 "
                    "degree (111km) for geographic coordinate systems. "
                    "Make sure you provided res in degree rather than in meters."
                )

    # get basin geometry and clip data
    kind = next(iter(region))
    xy = None
    if kind in ["basin", "subbasin"]:
        # parse_region_basin does not return xy, only geom...
        # should be fixed in core
        region_kwargs = _parse_region_value(
            region.pop(kind),
            data_catalog=data_catalog,
        )
        region_kwargs.update(region)
        if basin_index_fn is not None:
            bas_index = data_catalog.get_source(basin_index_fn)
        else:
            bas_index = None
        geom, xy = get_basin_geometry(
            ds=ds_org,
            kind=kind,
            basin_index=bas_index,
            **region_kwargs,
        )
    elif kind == "bbox":
        logger.warning(
            "Kind 'bbox' for the region is not recommended as it can lead "
            "to mistakes in the catchment delineation. Use carefully."
        )
        geom = parse_region_bbox(region)
    elif kind == "geom":
        logger.warning(
            "Kind 'geom' for the region is not recommended as it can lead "
            "to mistakes in the catchment delineation. Use carefully."
        )
        geom = parse_region_geom(region)
    else:
        raise ValueError(
            f"wflow region kind not understood or supported: {kind}. "
            "Use 'basin', 'subbasin', 'bbox' or 'geom'."
        )

    if geom is not None and geom.crs is None:
        raise ValueError("wflow region geometry has no CRS")

    # Set the basins geometry
    ds_org = ds_org.raster.clip_geom(geom, align=resolution, buffer=10)
    ds_org.coords["mask"] = ds_org.raster.geometry_mask(geom)

    geometries = {}
    if not math.isclose(scale_ratio, 1):
        geometries["basins_highres"] = geom

    return geometries, xy, ds_org
