"""Implementation for the RasterDatasetAdapter."""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, cast

import numpy as np
import pyproj
import xarray as xr
from pyproj import CRS

from hydromt._utils import (
    _has_no_data,
    _shift_dataset_time,
    _single_var_as_array,
)
from hydromt.data_catalog.adapters.adapter_utils import _slice_temporal_dimension
from hydromt.data_catalog.adapters.data_adapter_base import DataAdapterBase
from hydromt.error import NoDataException, NoDataStrategy, exec_nodata_strat
from hydromt.gis.raster import GEO_MAP_COORD
from hydromt.gis.raster_utils import _meridian_offset
from hydromt.typing import (
    Data,
    Geom,
    SourceMetadata,
    TimeRange,
    Variables,
)

logger = logging.getLogger(__name__)

__all__ = ["RasterDatasetAdapter"]


class RasterDatasetAdapter(DataAdapterBase):
    """Implementation for the RasterDatasetAdapter."""

    def transform(
        self,
        ds: xr.Dataset,
        metadata: SourceMetadata,
        *,
        mask: Optional[Geom] = None,
        align: Optional[bool] = None,
        variables: Optional[Variables] = None,
        time_range: Optional[TimeRange] = None,
        handle_nodata: NoDataStrategy = NoDataStrategy.RAISE,
        single_var_as_array: bool = True,
        buffer: int = 0,
    ) -> xr.Dataset | xr.DataArray | None:
        """Filter and harmonize the input RasterDataset.

        Parameters
        ----------
        ds : xr.Dataset
            input RasterDataset
        metadata : SourceMetadata
            source metadata
        mask : Optional[gpd.GeoDataFrame], optional
            mask to filter by geometry, by default None
        variables : Optional[List[str]], optional
            variable filter, by default None
        time_range : Optional[TimeRange], optional
            filter start and end times, by default None
        handle_nodata : NoDataStrategy, optional
            how to handle no data being present in the result, by default NoDataStrategy.RAISE
        single_var_as_array : bool, optional
            whether to return a xr.DataArray if only a single variable is present, by default True
        buffer : int, optional
            Buffer around the `bbox` or `geom` area of interest expressed in resolution multiplicity,
            by default 0

        Returns
        -------
        Optional[xr.Dataset]
            The filtered and harmonized RasterDataset, or None if no data was available

        Raises
        ------
        ValueError
            if not all variables are found in the data
        NoDataException
            if no data in left after slicing and handle_nodata is NoDataStrategy.RAISE
        """
        # rename variables and parse data and attrs
        ds = self._rename_vars(ds)
        ds = self._validate_spatial_dims(ds)
        ds = self._set_crs(ds, metadata.crs)
        ds = self._set_nodata(ds, metadata)
        ds = _shift_dataset_time(dt=self.unit_add.get("time", 0), ds=ds)
        # slice data
        ds = RasterDatasetAdapter._slice_data(
            ds,
            variables,
            mask,
            align,
            time_range,
            buffer=buffer,
            handle_nodata=handle_nodata,
        )
        if ds is None:
            return None  # if handle_nodata ignore

        # uniformize data
        ds = self._apply_unit_conversions(ds)
        ds = self._set_metadata(ds, metadata)
        # return array if single var and single_var_as_array
        return _single_var_as_array(ds, single_var_as_array, variables)

    def _rename_vars(self, ds: Data) -> Data:
        rm = {k: v for k, v in self.rename.items() if k in ds}
        ds = ds.rename(rm)
        return ds

    def _validate_spatial_dims(self, ds: Data) -> Data:
        if GEO_MAP_COORD in ds.data_vars:
            ds = ds.set_coords(GEO_MAP_COORD)
        try:
            ds.raster.set_spatial_dims()
            # transpose dims to get y and x dim last
            x_dim = ds.raster.x_dim
            y_dim = ds.raster.y_dim
            ds = ds.transpose(..., y_dim, x_dim)
        except ValueError:
            raise ValueError(
                f"RasterDataset: No valid spatial coords found in data {self.path}"
            )
        return ds

    def _set_crs(self, ds: Data, crs: Optional[CRS]) -> Data:
        # set crs
        if ds.raster.crs is None and crs is not None:
            ds.raster.set_crs(crs)
        elif ds.raster.crs is None:
            raise ValueError("RasterDataset: CRS not defined in data catalog or data.")
        elif crs is not None and ds.raster.crs != pyproj.CRS.from_user_input(crs):
            logger.warning(
                "RasterDataset: CRS from data catalog does not match CRS "
                " of data. The original CRS will be used. Please check your catalog."
            )
        return ds

    @staticmethod
    def _slice_data(
        ds: Data,
        variables: Optional[List] = None,
        mask: Optional[Geom] = None,
        align: Optional[float] = None,
        time_range: Optional[TimeRange] = None,
        buffer: int = 0,
        handle_nodata: NoDataStrategy = NoDataStrategy.RAISE,
    ) -> Optional[xr.Dataset]:
        """Filter the RasterDataset.

        Parameters
        ----------
        ds : xr.Dataset
            input Dataset.
        variables : Optional[List[str]], optional
            variable filter, by default None
        mask : Optional[gpd.GeoDataFrame], optional
            mask to filter by geometry, by default None
        align : Optional[float], optional
            resolution to align the bounding box, by default None
        time_range : Optional[TimeRange], optional
            filter start and end times, by default None
        buffer : int, optional
            Buffer around the `bbox` or `geom` area of interest expressed in resolution multiplicity,
            by default 0
        handle_nodata : NoDataStrategy, optional
            how to handle no data being present in the result, by default NoDataStrategy.RAISE

        Returns
        -------
        Optional[Union[xr.Dataset, xr.DataArray]]
            The filtered and harmonized RasterDataset, or None if no data was available

        Raises
        ------
        ValueError
            if not all variables are found in the data
        NoDataException
            if no data in left after slicing and handle_nodata is NoDataStrategy.RAISE
        """
        if isinstance(ds, xr.DataArray):  # xr.DataArray has no variables
            if ds.name is None:
                # dummy name, required to create dataset
                # renamed to variable in _single_var_as_array
                ds.name = "data"
            ds = ds.to_dataset()
        elif variables is not None:  # xr.Dataset has variables
            variables = cast(List, np.atleast_1d(variables).tolist())
            if len(variables) > 1 or len(ds.data_vars) > 1:
                mvars = [var not in ds.data_vars for var in variables]
                if any(mvars):
                    raise NoDataException(f"RasterDataset: variables not found {mvars}")
                ds = ds[variables]

        if time_range is not None:
            ds = _slice_temporal_dimension(ds, time_range, handle_nodata=handle_nodata)
            if ds is None:
                return None

        if mask is not None:
            ds = RasterDatasetAdapter._slice_spatial_dimensions(
                ds, mask, buffer=buffer, align=align, handle_nodata=handle_nodata
            )
            if ds is None:
                return None

        return ds

    @staticmethod
    def _slice_spatial_dimensions(
        ds: Data,
        mask: Optional[Geom] = None,
        align: Optional[float] = None,
        buffer: int = 0,
        handle_nodata: NoDataStrategy = NoDataStrategy.RAISE,
    ):
        # make sure bbox is in data crs
        bbox = None
        crs = ds.raster.crs
        epsg = crs.to_epsg()  # this could return None
        if mask is not None:
            bbox = mask.to_crs(crs).total_bounds
        # work with 4326 data that is defined at 0-360 degrees longtitude
        w, _, e, _ = ds.raster.bounds
        if epsg == 4326 and np.isclose(e - w, 360):  # allow for rounding errors
            ds = _meridian_offset(ds, bbox)

        # clip with bbox
        if bbox is not None:
            # check if bbox is fully covered
            bbox_str = ", ".join([f"{c:.3f}" for c in bbox])

            def _lt_or_close(a: float, b: float) -> bool:
                return np.isclose(a, b) or a < b

            w, s, e, n = ds.raster.bounds

            if not any(
                map(_lt_or_close, (w, s, bbox[2], bbox[3]), (bbox[0], bbox[1], e, n))
            ):
                logger.warning(
                    f"Dataset [{w}, {s}, {e}, {n}] does not fully cover bbox [{bbox_str}]"
                )

            logger.debug(f"Clip to [{bbox_str}] (epsg:{epsg}))")
            ds = ds.raster.clip_bbox(bbox, align=align, buffer=buffer)
            if _has_no_data(ds):
                exec_nodata_strat("No data left after bbox clipping.", handle_nodata)
                return None

        return ds

    def _apply_unit_conversions(self, ds: Data):
        unit_names = list(self.unit_mult.keys()) + list(self.unit_add.keys())
        unit_names = [k for k in unit_names if k in ds.data_vars]
        if len(unit_names) > 0:
            logger.debug(f"Convert units for {len(unit_names)} variables.")
        for name in list(set(unit_names)):  # unique
            m = self.unit_mult.get(name, 1)
            a = self.unit_add.get(name, 0)
            da = ds[name]
            attrs = da.attrs.copy()
            nodata_isnan = da.raster.nodata is None or np.isnan(da.raster.nodata)
            # nodata value is explicitly set to NaN in case no nodata value is provided
            nodata = np.nan if nodata_isnan else da.raster.nodata
            data_bool = ~np.isnan(da) if nodata_isnan else da != nodata
            ds[name] = xr.where(data_bool, da * m + a, nodata)
            ds[name].attrs.update(attrs)  # set original attributes
            ds[name].raster.set_nodata(nodata)  # reset nodata in case of change

        return ds

    def _set_nodata(self, ds, metadata: SourceMetadata):
        """Parse and apply nodata values from the data catalog."""
        # set nodata value
        if nodata := metadata.nodata is not None:
            if not isinstance(nodata, dict):
                no_data_values: Dict[str, Any] = {
                    k: metadata.nodata for k in ds.data_vars.keys()
                }
            else:
                no_data_values: Dict[str, Any] = nodata
            for k in ds.data_vars:
                mv = no_data_values.get(k, None)
                if mv is not None and ds[k].raster.nodata is None:
                    ds[k].raster.set_nodata(mv)
        return ds

    def _set_metadata(self, ds, metadata: SourceMetadata):
        # unit attributes
        attrs = metadata.attrs
        for k in attrs:
            ds[k].attrs.update(attrs[k])
        # set meta data
        # exclude extent and attrs that are dict
        ds.attrs.update(
            metadata.model_dump(exclude=["attrs", "extent"], exclude_unset=True)
        )
        return ds
