# -*- coding: utf-8 -*-
"""HydroMT GridModel class definition."""

import logging
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union

import geopandas as gpd
import numpy as np
import pandas as pd
import xarray as xr
from affine import Affine
from pyproj import CRS
from shapely.geometry import box

from .. import gis_utils, raster, workflows
from .model_api import Model

__all__ = ["GridModel"]
logger = logging.getLogger(__name__)


class GridMixin(object):
    # placeholders
    # xr.Dataset representation of all static parameter maps at the same resolution and
    # bounds - renamed from staticmaps
    _API = {"grid": xr.Dataset}

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._grid = None  # xr.Dataset()

    # generic grid methods
    def setup_grid_from_constant(
        self,
        constant: Union[int, float],
        name: str,
        dtype: Optional[str] = "float32",
        nodata: Optional[Union[int, float]] = None,
        mask_name: Optional[str] = "mask",
    ) -> List[str]:
        """HYDROMT CORE METHOD: Adds a grid based on a constant value.

        Parameters
        ----------
        constant: int, float
            Constant value to fill grid with.
        name: str
            Name of grid.
        dtype: str, optional
            Data type of grid. By default 'float32'.
        nodata: int, float, optional
            Nodata value. By default infered from dtype.
        mask_name: str, optional
            Name of mask in self.grid to use for masking raster_fn. By default 'mask'.
            Use None to disable masking.

        Returns
        -------
        list
            Names of added model grid layer.
        """
        da = workflows.grid.grid_from_constant(
            grid_like=self.grid,
            constant=constant,
            name=name,
            dtype=dtype,
            nodata=nodata,
            mask_name=mask_name,
        )
        # Add to grid
        self.set_grid(da)

        return [name]

    def setup_grid_from_rasterdataset(
        self,
        raster_fn: Union[str, Path, xr.DataArray, xr.Dataset],
        variables: Optional[List] = None,
        fill_method: Optional[str] = None,
        reproject_method: Optional[Union[List, str]] = "nearest",
        mask_name: Optional[str] = "mask",
        rename: Optional[Dict] = None,
    ) -> List[str]:
        """HYDROMT CORE METHOD: Add data variable(s) from ``raster_fn`` to grid object.

        If raster is a dataset, all variables will be added unless ``variables`` list
        is specified.

        Adds model layers:

        * **raster.name** grid: data from raster_fn

        Parameters
        ----------
        raster_fn: str, Path, xr.DataArray, xr.Dataset
            Data catalog key, path to raster file or raster xarray data object.
            If a path to a raster file is provided it will be added
            to the data_catalog with its name based on the file basename without
            extension.
        variables: list, optional
            List of variables to add to grid from raster_fn. By default all.
        fill_method : str, optional
            If specified, fills nodata values using fill_nodata method.
            Available methods are {'linear', 'nearest', 'cubic', 'rio_idw'}.
        reproject_method: list, str, optional
            See rasterio.warp.reproject for existing methods, by default 'nearest'.
            Can provide a list corresponding to ``variables``.
        mask_name: str, optional
            Name of mask in self.grid to use for masking raster_fn. By default 'mask'.
            Use None to disable masking.
        rename: dict, optional
            Dictionary to rename variable names in raster_fn before adding to grid
            {'name_in_raster_fn': 'name_in_grid'}. By default empty.

        Returns
        -------
        list
            Names of added model map layers
        """
        rename = rename or {}
        self.logger.info(f"Preparing grid data from raster source {raster_fn}")
        # Read raster data and select variables
        ds = self.data_catalog.get_rasterdataset(
            raster_fn,
            geom=self.region,
            buffer=2,
            variables=variables,
            single_var_as_array=False,
        )
        # Data resampling
        ds_out = workflows.grid.grid_from_rasterdataset(
            grid_like=self.grid,
            ds=ds,
            variables=variables,
            fill_method=fill_method,
            reproject_method=reproject_method,
            mask_name=mask_name,
            rename=rename,
        )
        # Add to grid
        self.set_grid(ds_out)

        return list(ds_out.data_vars.keys())

    def setup_grid_from_raster_reclass(
        self,
        raster_fn: Union[str, Path, xr.DataArray],
        reclass_table_fn: Union[str, Path, pd.DataFrame],
        reclass_variables: List,
        variable: Optional[str] = None,
        fill_method: Optional[str] = None,
        reproject_method: Optional[Union[List, str]] = "nearest",
        mask_name: Optional[str] = "mask",
        rename: Optional[Dict] = None,
        **kwargs,
    ) -> List[str]:
        """HYDROMT CORE METHOD: Add data variable(s) to grid object by reclassifying the data in ``raster_fn`` based on ``reclass_table_fn``.

        Adds model layers:

        * **reclass_variables** grid: reclassified raster data

        Parameters
        ----------
        raster_fn: str, Path, xr.DataArray
            Data catalog key, path to raster file or raster xarray data object.
            Should be a DataArray. Else use `variable` argument for selection.
        reclass_table_fn: str, Path, pd.DataFrame
            Data catalog key, path to tabular data file or tabular pandas dataframe
            object for the reclassification table of `raster_fn`.
        reclass_variables: list
            List of reclass_variables from reclass_table_fn table to add to maps.
            Index column should match values in `raster_fn`.
        variable: str, optional
            Name of raster_fn dataset variable to use. This is only required when
            reading datasets with multiple variables.
            By default None.
        fill_method : str, optional
            If specified, fills nodata values in `raster_fn` using fill_nodata method
            before reclassifying. Available methods are
            {'linear', 'nearest', 'cubic', 'rio_idw'}.
        reproject_method: str, optional
            See rasterio.warp.reproject for existing methods, by default "nearest".
            Can provide a list corresponding to ``reclass_variables``.
        mask_name: str, optional
            Name of mask in self.grid to use for masking raster_fn. By default 'mask'.
            Use None to disable masking.
        rename: dict, optional
            Dictionary to rename variable names in reclass_variables before adding to
            grid {'name_in_reclass_table': 'name_in_grid'}. By default empty.

        **kwargs : dict
            Additional keyword arguments to be passed to `get_rasterdataset`

        Returns
        -------
        list
            Names of added model grid layers
        """  # noqa: E501
        rename = rename or dict()
        self.logger.info(
            f"Preparing grid data by reclassifying the data in {raster_fn} based "
            f"on {reclass_table_fn}"
        )
        # Read raster data and remapping table
        da = self.data_catalog.get_rasterdataset(
            raster_fn, geom=self.region, buffer=2, variables=variable, **kwargs
        )
        if not isinstance(da, xr.DataArray):
            raise ValueError(
                f"raster_fn {raster_fn} should be a single variable. "
                "Please select one using the 'variable' argument"
            )
        df_vars = self.data_catalog.get_dataframe(
            reclass_table_fn, variables=reclass_variables
        )
        # Data resampling
        ds_vars = workflows.grid.grid_from_raster_reclass(
            grid_like=self.grid,
            da=da,
            reclass_table=df_vars,
            reclass_variables=reclass_variables,
            fill_method=fill_method,
            reproject_method=reproject_method,
            mask_name=mask_name,
            rename=rename,
        )
        # Add to maps
        self.set_grid(ds_vars.rename(rename))

        return list(ds_vars.data_vars.keys())

    def setup_grid_from_geodataframe(
        self,
        vector_fn: Union[str, Path, gpd.GeoDataFrame],
        variables: Optional[Union[List, str]] = None,
        nodata: Optional[Union[List, int, float]] = -1,
        rasterize_method: Optional[str] = "value",
        mask_name: Optional[str] = "mask",
        rename: Optional[Dict] = None,
        all_touched: Optional[bool] = True,
    ) -> List[str]:
        """HYDROMT CORE METHOD: Add data variable(s) to grid object by rasterizing the data from ``vector_fn``.

        Several type of rasterization are possible:
            * "fraction": the fraction of the grid cell covered by the vector
                shape is returned.
            * "area": the area of the grid cell covered by the vector shape is returned.
            * "value": the value from the variables columns of vector_fn are used.
                If this is used, variables must be specified.

        Parameters
        ----------
        vector_fn : str, Path, gpd.GeoDataFrame
            Data catalog key, path to vector file or a vector geopandas object.
        variables : List, str, optional
            List of variables to add to grid from vector_fn. Required if
            rasterize_method is "value", by default None.
        nodata : List, int, float, optional
            No data value to use for rasterization, by default -1. If a list is
            provided, it should have the same length has variables.
        rasterize_method : str, optional
            Method to rasterize the vector data. Either {"value", "fraction", "area"}.
            If "value", the value from the variables columns in vector_fn are used
            directly in the raster. If "fraction", the fraction of the grid
            cell covered by the vector file is returned. If "area", the area of the
            grid cell covered by the vector file is returned.
        mask_name: str, optional
            Name of mask in self.grid to use for masking raster_fn. By default 'mask'.
            Use None to disable masking.
        rename: dict, optional
            Dictionary to rename variable names in variables before adding to grid
            {'name_in_variables': 'name_in_grid'}. To rename with method fraction or
            area use {'vector_fn': 'name_in_grid'}. By default empty.
        all_touched : bool, optional
            If True (default), all pixels touched by geometries will be burned in.
            If false, only pixels whose center is within the polygon or that are
            selected by Bresenham's line algorithm will be burned in.

        Returns
        -------
        list
            Names of added model grid layers
        """  # noqa: E501
        rename = rename or dict()
        self.logger.info(f"Preparing grid data from vector '{vector_fn}'.")
        gdf = self.data_catalog.get_geodataframe(
            vector_fn, geom=self.region, dst_crs=self.crs
        )
        if gdf.empty:
            self.logger.warning(
                f"No shapes of {vector_fn} found within region,"
                " skipping setup_grid_from_vector."
            )
            return
        # Data resampling
        if vector_fn in rename.keys():
            # In case of choosing a new name with area or fraction method pass
            # the name directly
            rename = rename[vector_fn]
        ds = workflows.grid.grid_from_geodataframe(
            grid_like=self.grid,
            gdf=gdf,
            variables=variables,
            nodata=nodata,
            rasterize_method=rasterize_method,
            mask_name=mask_name,
            rename=rename,
            all_touched=all_touched,
        )
        # Add to grid
        self.set_grid(ds)

        return list(ds.data_vars.keys())

    @property
    def grid(self):
        """Model static gridded data as xarray.Dataset."""
        if self._grid is None:
            self._initialize_grid()
        return self._grid

    def _initialize_grid(self, skip_read=False) -> None:
        """Initialize grid object."""
        if self._grid is None:
            self._grid = xr.Dataset()
            if self._read and not skip_read:
                self.read_grid()

    def set_grid(
        self,
        data: Union[xr.DataArray, xr.Dataset, np.ndarray],
        name: Optional[str] = None,
    ):
        """Add data to grid.

        All layers of grid must have identical spatial coordinates.

        Parameters
        ----------
        data: xarray.DataArray or xarray.Dataset
            new map layer to add to grid
        name: str, optional
            Name of new map layer, this is used to overwrite the name of a DataArray
            and ignored if data is a Dataset
        """
        self._initialize_grid()
        # NOTE: variables in a dataset are not longer renamed as used to be the case in
        # set_staticmaps
        name_required = isinstance(data, np.ndarray) or (
            isinstance(data, xr.DataArray) and data.name is None
        )
        if name is None and name_required:
            raise ValueError(f"Unable to set {type(data).__name__} data without a name")
        if isinstance(data, np.ndarray):
            if data.shape != self.grid.raster.shape:
                raise ValueError("Shape of data and grid maps do not match")
            data = xr.DataArray(dims=self.grid.raster.dims, data=data, name=name)
        if isinstance(data, xr.DataArray):
            if name is not None:  # rename
                data.name = name
            data = data.to_dataset()
        elif not isinstance(data, xr.Dataset):
            raise ValueError(f"cannot set data of type {type(data).__name__}")

        if len(self._grid) == 0:  # empty grid
            self._grid = data
        else:
            for dvar in data.data_vars:
                if dvar in self._grid:
                    if self._read:
                        self.logger.warning(f"Replacing grid map: {dvar}")
                self._grid[dvar] = data[dvar]

    def read_grid(self, fn: str = "grid/grid.nc", **kwargs) -> None:
        """Read model grid data at <root>/<fn> and add to grid property.

        key-word arguments are passed to :py:meth:`~hydromt.models.Model.read_nc`

        Parameters
        ----------
        fn : str, optional
            filename relative to model root, by default 'grid/grid.nc'
        **kwargs : dict
            Additional keyword arguments to be passed to the `read_nc` method.
        """
        self._assert_read_mode()
        self._initialize_grid(skip_read=True)

        # Load grid data in r+ mode to allow overwritting netcdf files
        if self._read and self._write:
            kwargs["load"] = True
        loaded_nc_files = self.read_nc(fn, single_var_as_array=False, **kwargs)
        for ds in loaded_nc_files.values():
            self.set_grid(ds)

    def write_grid(
        self,
        fn: str = "grid/grid.nc",
        gdal_compliant: bool = False,
        rename_dims: bool = False,
        force_sn: bool = False,
        **kwargs,
    ) -> None:
        """Write model grid data to netcdf file at <root>/<fn>.

        key-word arguments are passed to :py:meth:`~hydromt.models.Model.write_nc`

        Parameters
        ----------
        fn : str, optional
            filename relative to model root, by default 'grid/grid.nc'
        **kwargs : dict
            Additional keyword arguments to be passed to the `write_nc` method.
        gdal_compliant : bool, optional
            If True, write grid data in a way that is compatible with GDAL,
            by default False
        rename_dims: bool, optional
            If True and gdal_compliant, rename x_dim and y_dim to standard names
            depending on the CRS (x/y for projected and lat/lon for geographic).
        force_sn: bool, optional
            If True and gdal_compliant, forces the dataset to have
            South -> North orientation.
        """
        if len(self.grid) == 0:
            self.logger.debug("No grid data found, skip writing.")
        else:
            self._assert_write_mode()
            # write_nc requires dict - use dummy 'grid' key
            self.write_nc(
                {"grid": self.grid},
                fn,
                gdal_compliant=gdal_compliant,
                rename_dims=rename_dims,
                force_sn=force_sn,
                **kwargs,
            )


class GridModel(GridMixin, Model):

    """Model class Grid Model for gridded models in HydroMT."""

    _CLI_ARGS = {"region": "setup_grid"}
    _NAME = "grid_model"

    def __init__(
        self,
        root: str = None,
        mode: str = "w",
        config_fn: str = None,
        data_libs: List[str] = None,
        logger=logger,
    ):
        """Initialize a GridModel for distributed models with a regular grid."""
        super().__init__(
            root=root,
            mode=mode,
            config_fn=config_fn,
            data_libs=data_libs,
            logger=logger,
        )

    ## GENERIC SETUP METHODS
    def setup_grid(
        self,
        region: dict,
        res: Optional[float] = None,
        crs: int = None,
        rotated: bool = False,
        hydrography_fn: Optional[str] = None,
        basin_index_fn: Optional[str] = None,
        add_mask: bool = True,
        align: bool = True,
        dec_origin: int = 0,
        dec_rotation: int = 3,
    ) -> xr.DataArray:
        """HYDROMT CORE METHOD: Create a 2D regular grid or reads an existing grid.

        A 2D regular grid will be created from a geometry (geom_fn) or bbox. If an
        existing grid is given, then no new grid will be generated.

        Adds/Updates model layers (if add_mask):
        * **mask** grid mask: add grid mask to grid object

        Parameters
        ----------
        region : dict
            Dictionary describing region of interest, e.g.:
            * {'bbox': [xmin, ymin, xmax, ymax]}
            * {'geom': 'path/to/polygon_geometry'}
            * {'grid': 'path/to/grid_file'}
            * {'basin': [x, y]}

            Region must be of kind [grid, bbox, geom, basin, subbasin, interbasin].
        res: float
            Resolution used to generate 2D grid [unit of the CRS], required if region
            is not based on 'grid'.
        crs : EPSG code, int, str optional
            EPSG code of the model or "utm" to let hydromt find the closest projected
        rotated : bool, optional
            if True, a minimum rotated rectangular grid is fitted around the region,
            by default False. Only  applies if region is of kind 'bbox', 'geom'
        hydrography_fn : str
            Name of data source for hydrography data. Required if region is of kind
                'basin', 'subbasin' or 'interbasin'.

            * Required variables: ['flwdir'] and any other 'snapping' variable required
                to define the region.

            * Optional variables: ['basins'] if the `region` is based on a
                (sub)(inter)basins without a 'bounds' argument.

        basin_index_fn : str
            Name of data source with basin (bounding box) geometries associated with
            the 'basins' layer of `hydrography_fn`. Only required if the `region` is
            based on a (sub)(inter)basins without a 'bounds' argument.
        add_mask : bool, optional
            Add mask variable to grid object, by default True.
        align : bool, optional
            If True (default), align target transform to resolution.
        dec_origin : int, optional
            number of decimals to round the origin coordinates, by default 0
        dec_rotation : int, optional
            number of decimals to round the rotation angle, by default 3

        Returns
        -------
        grid : xr.DataArray
            Generated grid mask.
        """
        self.logger.info("Preparing 2D grid.")

        kind = next(iter(region))  # first key of region
        if kind in ["bbox", "geom", "basin", "subbasin", "interbasin"]:
            # Do not parse_region for grid as we want to allow for more (file) formats
            kind, region = workflows.parse_region(
                region, data_catalog=self.data_catalog, logger=self.logger
            )
        elif kind != "grid":
            raise ValueError(
                f"Region for grid must be of kind [grid, bbox, geom, basin, subbasin,"
                f" interbasin], kind {kind} not understood."
            )

        # Derive xcoords, ycoords and geom for the different kind options
        if kind in ["bbox", "geom"]:
            if not isinstance(res, (int, float)):
                raise ValueError("res argument required for kind 'bbox', 'geom'")
            if kind == "bbox":
                bbox = region["bbox"]
                geom = gpd.GeoDataFrame(geometry=[box(*bbox)], crs=4326)
            elif kind == "geom":
                geom = region["geom"]
                if geom.crs is None:
                    raise ValueError('Model region "geom" has no CRS')
            if crs is not None:
                crs = gis_utils.parse_crs(crs, bbox=geom.total_bounds)
                geom = geom.to_crs(crs)
            # Generate grid based on res for region bbox
            # TODO add warning on res value if crs is projected or not?
            if not rotated:
                xmin, ymin, xmax, ymax = geom.total_bounds
                res = abs(res)
                if align:
                    xmin = round(xmin / res) * res
                    ymin = round(ymin / res) * res
                    xmax = round(xmax / res) * res
                    ymax = round(ymax / res) * res
                xcoords = np.linspace(
                    xmin + res / 2,
                    xmax - res / 2,
                    num=round((xmax - xmin) / res),
                    endpoint=True,
                )
                ycoords = np.flip(
                    np.linspace(
                        ymin + res / 2,
                        ymax - res / 2,
                        num=round((ymax - ymin) / res),
                        endpoint=True,
                    )
                )
            else:  # rotated
                geomu = geom.unary_union
                x0, y0, mmax, nmax, rot = workflows.grid.rotated_grid(
                    geomu, res, dec_origin=dec_origin, dec_rotation=dec_rotation
                )
                transform = (
                    Affine.translation(x0, y0)
                    * Affine.rotation(rot)
                    * Affine.scale(res, res)
                )

        elif kind in ["basin", "subbasin", "interbasin"]:
            # retrieve global hydrography data (lazy!)
            ds_hyd = self.data_catalog.get_rasterdataset(hydrography_fn)
            if "bounds" not in region:
                region.update(basin_index=self.data_catalog.get_source(basin_index_fn))
            # get basin geometry
            geom, xy = workflows.get_basin_geometry(
                ds=ds_hyd,
                kind=kind,
                logger=self.logger,
                **region,
            )
            # get ds_hyd again but clipped to geom, one variable is enough
            da_hyd = self.data_catalog.get_rasterdataset(
                hydrography_fn, geom=geom, variables=["flwdir"]
            )
            if not isinstance(res, (int, float)):
                self.logger.info(
                    "res argument not defined, using resolution of "
                    f"hydrography_fn {da_hyd.raster.res}"
                )
                res = da_hyd.raster.res
            # Reproject da_hyd based on crs and grid and align, method is not important
            # only coords will be used
            # TODO add warning on res value if crs is projected or not?
            if res != da_hyd.raster.res and crs != da_hyd.raster.crs:
                da_hyd = da_hyd.raster.reproject(
                    dst_crs=crs,
                    dst_res=res,
                    align=align,
                )
            # Get xycoords, geom
            xcoords = da_hyd.raster.xcoords.values
            ycoords = da_hyd.raster.ycoords.values
            if geom.crs != da_hyd.raster.crs:
                crs = da_hyd.raster.crs
                geom = geom.to_crs(crs)
        elif kind == "grid":
            # Support more formats for grid input (netcdf, zarr, io.open_raster)
            fn = region[kind]
            if isinstance(fn, (xr.DataArray, xr.Dataset)):
                da_like = fn
            else:
                da_like = self.data_catalog.get_rasterdataset(fn)
            # Get xycoords, geom
            xcoords = da_like.raster.xcoords.values
            ycoords = da_like.raster.ycoords.values
            geom = da_like.raster.box
            if crs is not None or res is not None:
                self.logger.warning(
                    "For region kind 'grid', the gris crs/res are used and not"
                    f" user-defined crs {crs} or res {res}"
                )
            crs = da_like.raster.crs

        # Instantiate grid object
        # Generate grid using hydromt full method
        if not rotated:
            coords = {"y": ycoords, "x": xcoords}
            grid = raster.full(
                coords=coords,
                nodata=1,
                dtype=np.uint8,
                name="mask",
                attrs={},
                crs=geom.crs,
                lazy=False,
            )
        else:
            grid = raster.full_from_transform(
                transform,
                shape=(mmax, nmax),
                nodata=1,
                dtype=np.uint8,
                name="mask",
                attrs={},
                crs=geom.crs,
                lazy=False,
            )
        # Create geometry_mask with geom
        if add_mask:
            grid = grid.raster.geometry_mask(geom, all_touched=True)
            grid.name = "mask"
        # Remove mask variable mask from grid if not add_mask
        else:
            grid = grid.to_dataset()
            grid = grid.drop_vars("mask")

        # Add region and grid to model
        self.set_geoms(geom, "region")
        self.set_grid(grid)

        return grid

    ## I/O
    def read(
        self,
        components: List = None,
    ) -> None:
        """Read the complete model schematization and configuration from model files.

        Parameters
        ----------
        components : List, optional
            List of model components to read, each should have an associated
            read_<component> method. By default ['config', 'maps', 'grid',
            'geoms', 'tables', 'forcing', 'states', 'results']
        """
        components = components or [
            "config",
            "grid",
            "geoms",
            "tables",
            "forcing",
            "states",
            "results",
        ]
        super().read(components=components)

    def write(
        self,
        components: List = None,
    ) -> None:
        """Write the complete model schematization and configuration to model files.

        Parameters
        ----------
        components : List, optional
            List of model components to write, each should have an
            associated write_<component> method. By default
            ['config', 'maps', 'grid', 'geoms', 'tables', 'forcing', 'states']
        """
        components = components or [
            "config",
            "maps",
            "grid",
            "geoms",
            "tables",
            "forcing",
            "states",
        ]
        super().write(components=components)

    # Properties for subclass GridModel
    @property
    def res(self) -> Tuple[float, float]:
        """Returns the resolution of the model grid."""
        if len(self.grid) > 0:
            return self.grid.raster.res

    @property
    def transform(self):
        """Returns spatial transform of the model grid."""
        if len(self.grid) > 0:
            return self.grid.raster.transform

    @property
    def crs(self) -> Union[CRS, None]:
        """Returns coordinate reference system embedded in the model grid."""
        if self.grid.raster.crs is not None:
            return CRS(self.grid.raster.crs)

    @property
    def bounds(self) -> List[float]:
        """Returns the bounding box of the model grid."""
        if len(self.grid) > 0:
            return self.grid.raster.bounds

    @property
    def region(self) -> gpd.GeoDataFrame:
        """Returns the geometry of the model area of interest."""
        region = gpd.GeoDataFrame()
        if "region" in self.geoms:
            region = self.geoms["region"]
        elif len(self.grid) > 0:
            crs = self.crs
            if crs is not None and hasattr(crs, "to_epsg"):
                crs = crs.to_epsg()  # not all CRS have an EPSG code
            region = gpd.GeoDataFrame(geometry=[box(*self.bounds)], crs=crs)
        return region

    def set_crs(self, crs: CRS) -> None:
        """Set coordinate reference system of the model grid."""
        if len(self.grid) > 0:
            self.grid.raster.set_crs(crs)
