from typing import Union, Optional, List, Tuple
import logging
import os
from os.path import join, isdir, dirname, isfile
from pathlib import Path
import numpy as np
import xarray as xr
import xugrid as xu
import geopandas as gpd
from shapely.geometry import box, Polygon

from ..raster import GEO_MAP_COORD
from .model_api import Model
from .. import workflows

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


class MeshMixin(object):
    # placeholders
    # We cannot initialize an empty xu.UgridDataArray
    _API = {
        "mesh": Union[xu.UgridDataArray, xu.UgridDataset],
    }

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._mesh = None

    ## general setup methods
    def setup_mesh_from_raster(
        self,
        raster_fn: str,
        variables: Optional[list] = None,
        fill_method: Optional[str] = None,
        resampling_method: Optional[str] = "mean",
        all_touched: Optional[bool] = True,
    ) -> None:
        """
        This component adds data variable(s) from ``raster_fn`` to mesh object.

        Raster data is interpolated to the mesh grid using the ``resampling_method``.
        If raster is a dataset, all variables will be added unless ``variables`` list is specified.

        Adds model layers:

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

        Parameters
        ----------
        raster_fn: str
            Source name of raster data in data_catalog.
        variables: list, optional
            List of variables to add to mesh from raster_fn. By default all.
        fill_method : str, optional
            If specified, fills no data values using fill_nodata method. AVailable methods
            are {'linear', 'nearest', 'cubic', 'rio_idw'}.
        resampling_method: str, optional
            Method to sample from raster data to mesh. By default mean. Options include
            {'count', 'min', 'max', 'sum', 'mean', 'std', 'median', 'q##'}.
        all_touched : bool, optional
            If True, all pixels touched by geometries will used to define the sample.
            If False, only pixels whose center is within the geometry or that are
            selected by Bresenham's line algorithm will be used. By default True.
        """
        self.logger.info(f"Preparing mesh 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
        )
        if isinstance(ds, xr.DataArray):
            ds = ds.to_dataset()

        if fill_method is not None:
            ds = ds.raster.interpolate_na(method=fill_method)

        # Convert mesh grid as geodataframe for sampling
        # Reprojection happens to gdf inside of zonal_stats method
        ds_sample = ds.raster.zonal_stats(
            gdf=self.mesh_gdf, stats=resampling_method, all_touched=all_touched
        )
        # Rename variables
        rm_dict = {f"{var}_{resampling_method}": var for var in ds.data_vars}
        ds_sample = ds_sample.rename(rm_dict)
        # Convert to UgridDataset
        uds_sample = xu.UgridDataset(ds_sample, grids=self.mesh.ugrid.grid)

        self.set_mesh(uds_sample)

    def setup_mesh_from_raster_reclass(
        self,
        raster_fn: str,
        reclass_table_fn: str,
        reclass_variables: list,
        variable: Optional[str] = None,
        fill_nodata: Optional[str] = None,
        resampling_method: Optional[Union[str, list]] = "mean",
        all_touched: Optional[bool] = True,
        **kwargs,
    ) -> None:
        """
        This component adds data variable(s) to mesh object by reclassifying the data in ``raster_fn`` based on ``reclass_table_fn``.

        The reclassified raster data are subsequently interpolated to the mesh using ``resampling_method``.

        Adds model layers:

        * **reclass_variables** mesh: reclassified raster data interpolated to the model mesh

        Parameters
        ----------
        raster_fn: str
            Source name of raster data in data_catalog. Should be a DataArray. Else use **kwargs to select variables/time_tuple in
            :py:meth:`hydromt.data_catalog.get_rasterdataset` method
        reclass_table_fn: str
            Source name of reclassification table of raster_fn in data_catalog.
        reclass_variables: list
            List of reclass_variables from reclass_table_fn table to add to mesh. Index column should match values in raster_fn.
        variable: str, optional
            Name of raster 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'}.
        resampling_method: str, list, optional
            Method to sample from raster data to mesh. Can be a list per variable in ``reclass_variables`` or a
            single method for all. By default mean for all reclass_variables.
            Options include {'count', 'min', 'max', 'sum', 'mean', 'std', 'median', 'q##'}.
        all_touched : bool, optional
            If True, all pixels touched by geometries will used to define the sample.
            If False, only pixels whose center is within the geometry or that are
            selected by Bresenham's line algorithm will be used. By default True.
        """
        self.logger.info(
            f"Preparing mesh data by reclassifying the data in {raster_fn} based on {reclass_table_fn}."
        )
        # Read raster data and mapping 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 raster. "
                "Please select one using the 'variable' argument"
            )
        df_vars = self.data_catalog.get_dataframe(
            reclass_table_fn, variables=reclass_variables
        )

        if fill_nodata is not None:
            ds = ds.raster.interpolate_na(method=fill_nodata)

        # Mapping function
        ds_vars = da.raster.reclassify(reclass_table=df_vars, method="exact")

        # Convert mesh grid as geodataframe for sampling
        # Reprojection happens to gdf inside of zonal_stats method
        ds_sample = ds_vars.raster.zonal_stats(
            gdf=self.mesh_gdf,
            stats=np.unique(np.atleast_1d(resampling_method)),
            all_touched=all_touched,
        )
        # Rename variables
        if isinstance(resampling_method, str):
            resampling_method = np.repeat(resampling_method, len(reclass_variables))
        rm_dict = {
            f"{var}_{mtd}": var
            for var, mtd in zip(reclass_variables, resampling_method)
        }
        ds_sample = ds_sample.rename(rm_dict)
        ds_sample = ds_sample[reclass_variables]
        # Convert to UgridDataset
        uds_sample = xu.UgridDataset(ds_sample, grids=self.mesh.ugrid.grid)

        self.set_mesh(uds_sample)

    @property
    def mesh(self) -> Union[xu.UgridDataArray, xu.UgridDataset]:
        """Model static mesh data. Returns a xarray.Dataset."""
        # XU grid data type Xarray dataset with xu sampling.
        if self._mesh is None and self._read:
            self.read_mesh()
        return self._mesh

    def set_mesh(
        self,
        data: Union[xu.UgridDataArray, xu.UgridDataset],
        name: Optional[str] = None,
    ) -> None:
        """Add data to mesh.

        All layers of mesh have identical spatial coordinates in Ugrid conventions.

        Parameters
        ----------
        data: xugrid.UgridDataArray or xugrid.UgridDataset
            new layer to add to mesh
        name: str, optional
            Name of new object layer, this is used to overwrite the name of a UgridDataArray.
        """
        if not isinstance(data, (xu.UgridDataArray, xu.UgridDataset)):
            raise ValueError(
                "New mesh data in set_mesh should be of type xu.UgridDataArray or xu.UgridDataset"
            )
        if isinstance(data, xu.UgridDataArray):
            if name is not None:
                data = data.rename(name)
            elif data.name is None:
                raise ValueError(
                    f"Cannot set mesh from {str(type(data).__name__)} without a name."
                )
            data = data.to_dataset()
        if self._mesh is None:  # NOTE: mesh is initialized with None
            self._mesh = data
        else:
            for dvar in data.data_vars:
                if dvar in self._mesh:
                    self.logger.warning(f"Replacing mesh parameter: {dvar}")
                self._mesh[dvar] = data[dvar]

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

        key-word arguments are passed to :py:func:`xr.open_dataset`

        Parameters
        ----------
        fn : str, optional
            filename relative to model root, by default 'mesh/mesh.nc'
        """
        self._assert_read_mode
        for ds in self._read_nc(fn, **kwargs).values():
            uds = xu.UgridDataset(ds)
            if ds.rio.crs is not None:  # parse crs
                uds.ugrid.grid.set_crs(ds.raster.crs)
                uds = uds.drop_vars(GEO_MAP_COORD, errors="ignore")
            self.set_mesh(uds)

    def write_mesh(self, fn: str = "mesh/mesh.nc", **kwargs) -> None:
        """Write model grid data to netcdf file at <root>/<fn>

        key-word arguments are passed to :py:meth:`xarray.Dataset.ugrid.to_netcdf`

        Parameters
        ----------
        fn : str, optional
            filename relative to model root, by default 'grid/grid.nc'
        """
        if self._mesh is None:
            self.logger.debug("No mesh data found, skip writing.")
            return
        self._assert_write_mode
        # filename
        _fn = join(self.root, fn)
        if not isdir(dirname(_fn)):
            os.makedirs(dirname(_fn))
        self.logger.debug(f"Writing file {fn}")
        # ds_new = xu.UgridDataset(grid=ds_out.ugrid.grid) # bug in xugrid?
        ds_out = self.mesh.ugrid.to_dataset()
        if self.mesh.ugrid.grid.crs is not None:
            # save crs to spatial_ref coordinate
            ds_out = ds_out.rio.write_crs(self.mesh.ugrid.grid.crs)
        ds_out.to_netcdf(_fn, **kwargs)


class MeshModel(MeshMixin, Model):
    """Model class Mesh Model for mesh models in HydroMT"""

    _CLI_ARGS = {"region": "setup_mesh", "res": "setup_mesh"}
    _NAME = "mesh_model"

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

    ## general setup methods
    def setup_mesh(
        self,
        region: dict,
        res: Optional[float] = None,
        crs: int = None,
    ) -> xu.UgridDataset:
        """Creates an 2D unstructured mesh or reads an existing 2D mesh according UGRID conventions.
        An 2D unstructured mesh will be created as 2D rectangular grid from a geometry (geom_fn) or bbox. If an existing
        2D mesh is given, then no new mesh will be generated

        Note Only existing meshed with only 2D grid can be read.
        #FIXME: read existing 1D2D network file and extract 2D part.

        Adds/Updates model layers:

        * **mesh** mesh topology: add mesh topology to mesh object

        Parameters
        ----------
        region : dict
            Dictionary describing region of interest, e.g.:

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

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

            * {'mesh': 'path/to/2dmesh_file'}
        res: float
            Resolution used to generate 2D mesh [unit of the CRS], required if region is not based on 'mesh'.
        crs : EPSG code, int, optional
            Optional EPSG code of the model. If None using the one from region, and else 4326.

        Returns
        -------
        mesh2d : xu.UgridDataset
            Generated mesh2d.

        """
        self.logger.info(f"Preparing 2D mesh.")

        if "mesh" not in region:
            if not isinstance(res, (int, float)):
                raise ValueError("res argument required")
            kind, region = workflows.parse_region(region, logger=self.logger)
            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')
            else:
                raise ValueError(
                    f"Region for mesh must of kind [bbox, geom, mesh], kind {kind} not understood."
                )
            if crs is not None:
                geom = geom.to_crs(crs)
            # Generate grid based on res for region bbox
            xmin, ymin, xmax, ymax = geom.total_bounds
            # note we flood the number of faces within bounds
            ncol = int((xmax - xmin) // res)
            nrow = int((ymax - ymin) // res)
            dx, dy = res, -res
            faces = []
            for i in range(nrow):
                top = ymax + i * dy
                bottom = ymax + (i + 1) * dy
                for j in range(ncol):
                    left = xmin + j * dx
                    right = xmin + (j + 1) * dx
                    faces.append(box(left, bottom, right, top))
            grid = gpd.GeoDataFrame(geometry=faces, crs=geom.crs)
            # If needed clip to geom
            if kind != "bbox":
                # TODO: grid.intersects(geom) does not seem to work ?
                grid = grid.loc[
                    gpd.sjoin(
                        grid, geom, how="left", op="intersects"
                    ).index_right.notna()
                ].reset_index()
            # Create mesh from grid
            grid.index.name = "mesh2d_nFaces"
            mesh2d = xu.UgridDataset.from_geodataframe(grid)
            mesh2d.ugrid.grid.set_crs(grid.crs)

        else:
            mesh2d_fn = region["mesh"]
            if isinstance(mesh2d_fn, (str, Path)) and isfile(mesh2d_fn):
                self.logger.info(f"An existing 2D grid is used to prepare 2D mesh.")

                ds = xr.open_dataset(mesh2d_fn, mask_and_scale=False)
            elif isinstance(mesh2d_fn, xr.Dataset):
                ds = mesh2d_fn
            else:
                raise ValueError(
                    f"Region 'mesh' file {mesh2d_fn} not found, please check"
                )
            topologies = [
                k for k in ds.data_vars if ds[k].attrs.get("cf_role") == "mesh_topology"
            ]
            for topology in topologies:
                topodim = ds[topology].attrs["topology_dimension"]
                if topodim != 2:  # chek if 2d mesh file else throw error
                    raise NotImplementedError(
                        f"{mesh2d_fn} cannot be opened. Please check if the existing grid is "
                        f"an 2D mesh and not 1D2D mesh. This option is not yet available for 1D2D meshes."
                    )

            # Continues with a 2D grid
            mesh2d = xu.UgridDataset(ds)
            # Check crs and reproject to model crs
            if crs is None:
                crs = 4326
            if ds.rio.crs is not None:  # parse crs
                mesh2d.ugrid.grid.set_crs(ds.raster.crs)
            else:
                # Assume model crs
                self.logger.warning(
                    f"Mesh data from {mesh2d_fn} doesn't have a CRS. Assuming crs option {crs}"
                )
                mesh2d.ugrid.grid.set_crs(crs)
            mesh2d = mesh2d.drop_vars(GEO_MAP_COORD, errors="ignore")

        # Reproject to user crs option if needed
        if mesh2d.ugrid.grid.crs != crs and crs is not None:
            self.logger.info(f"Reprojecting mesh to crs {crs}")
            mesh2d.ugrid.grid.to_crs(self.crs)

        self.set_mesh(mesh2d)

        # This setup method returns region so that it can be wrapped for models which require
        # more information
        return mesh2d

    ## I/O
    def read(
        self,
        components: List = [
            "config",
            "mesh",
            "geoms",
            "forcing",
            "states",
            "results",
        ],
    ) -> 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', 'mesh', 'geoms', 'forcing', 'states', 'results']
        """
        super().read(components=components)

    def write(
        self,
        components: List = ["config", "mesh", "geoms", "forcing", "states"],
    ) -> 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', 'mesh', 'geoms', 'forcing', 'states']
        """
        super().write(components=components)

    # MeshModel specific methods

    # MeshModel properties
    @property
    def bounds(self) -> Tuple:
        """Returns model mesh bounds."""
        if self._mesh is not None:
            return self._mesh.ugrid.grid.bounds

    @property
    def region(self) -> gpd.GeoDataFrame:
        """Returns geometry of region of the model area of interest."""
        region = gpd.GeoDataFrame()
        if "region" in self.geoms:
            region = self.geoms["region"]
        elif self.mesh is not None:
            crs = self.mesh.ugrid.grid.crs
            if crs is 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

    @property
    def mesh_gdf(self) -> gpd.GeoDataFrame:
        """Returns geometry of mesh as a gpd.GeoDataFrame"""
        if self._mesh is not None:
            name = [n for n in self.mesh.data_vars][0]  # works better on a DataArray
            return self._mesh[name].ugrid.to_geodataframe()
