import abc
import numbers
from pathlib import Path
from typing import Any, Mapping, Optional, Self

import numpy as np
import xarray as xr
import xugrid as xu
from xarray.core.utils import is_scalar

import imod
from imod.common.interfaces.ipackagebase import IPackageBase
from imod.typing.grid import (
    GridDataArray,
    GridDataset,
    is_spatial_grid,
    merge_with_dictionary,
)

TRANSPORT_PACKAGES = ("adv", "dsp", "ssm", "mst", "ist", "src")
EXCHANGE_PACKAGES = ("gwfgwf", "gwfgwt", "gwtgwt")
UTIL_PACKAGES = ("ats", "hpc")


def _is_scalar_nan(da: GridDataArray):
    """
    Test if is_scalar_nan, carefully avoid loading grids in memory
    """
    scalar_data: bool = is_scalar(da)
    if scalar_data:
        stripped_value = da.to_numpy()[()]
        return isinstance(stripped_value, numbers.Real) and np.isnan(stripped_value)  # type: ignore[call-overload]
    return False


class PackageBase(IPackageBase, abc.ABC):
    """
    This class is used for storing a collection of Xarray DataArrays or UgridDataArrays
    in a dataset. A load-from-file method is also provided. Storing to file is done by calling
    object.dataset.to_netcdf(...)
    """

    # This method has been added to allow mock.patch to mock created objects
    # https://stackoverflow.com/questions/64737213/how-to-patch-the-new-method-of-a-class
    def __new__(cls, *_, **__):
        return super(PackageBase, cls).__new__(cls)

    def __init__(
        self, variables_to_merge: Mapping[str, GridDataArray | float | int | bool | str]
    ):
        # Merge variables, perform exact join to verify if coordinates values
        # are consistent amongst variables.
        self.__dataset = merge_with_dictionary(variables_to_merge, join="exact")

    @property
    def dataset(self) -> GridDataset:
        return self.__dataset

    @dataset.setter
    def dataset(self, value: GridDataset) -> None:
        self.__dataset = value

    def __getitem__(self, key):
        return self.dataset.__getitem__(key)

    def __setitem__(self, key, value):
        self.dataset.__setitem__(key, value)

    def to_netcdf(
        self, *args, mdal_compliant: bool = False, crs: Optional[Any] = None, **kwargs
    ) -> None:
        """
        Write dataset contents to a netCDF file. Custom encoding rules can be
        provided on package level by overriding the _netcdf_encoding in the
        package

        Parameters
        ----------
        *args:
            Will be passed on to ``xr.Dataset.to_netcdf`` or
            ``xu.UgridDataset.to_netcdf``.
        mdal_compliant: bool, optional
            Convert data with
            :func:`imod.prepare.spatial.mdal_compliant_ugrid2d` to MDAL
            compliant unstructured grids. Defaults to False.
        crs: Any, optional
            Anything accepted by rasterio.crs.CRS.from_user_input
            Requires ``rioxarray`` installed.
        **kwargs:
            Will be passed on to ``xr.Dataset.to_netcdf`` or
            ``xu.UgridDataset.to_netcdf``.

        """
        kwargs.update({"encoding": self._netcdf_encoding()})

        dataset = self.dataset
        if isinstance(dataset, xu.UgridDataset):
            if mdal_compliant:
                dataset = dataset.ugrid.to_dataset()
                mdal_dataset = imod.util.spatial.mdal_compliant_ugrid2d(
                    dataset, crs=crs
                )
                mdal_dataset.to_netcdf(*args, **kwargs)
            else:
                dataset.ugrid.to_netcdf(*args, **kwargs)
        else:
            if is_spatial_grid(dataset):
                dataset = imod.util.spatial.gdal_compliant_grid(dataset, crs=crs)
            dataset.to_netcdf(*args, **kwargs)

    def _netcdf_encoding(self) -> dict:
        """

        The encoding used in the to_netcdf method
        Override this to provide custom encoding rules

        """
        return {}

    @classmethod
    def _from_dataset(cls, ds: GridDataset) -> Self:
        """
        Create package from dataset. Note that no initialization validation is
        done.
        """
        instance = cls.__new__(cls)
        instance.dataset = ds
        return instance

    @classmethod
    def from_file(cls, path: str | Path, **kwargs) -> Self:
        """
        Loads an imod mf6 package from a file (currently only netcdf and zarr are supported).
        Note that it is expected that this file was saved with imod.mf6.Package.dataset.to_netcdf(),
        as the checks upon package initialization are not done again!

        Parameters
        ----------
        path : str, pathlib.Path
            Path to the file.
        **kwargs : keyword arguments
            Arbitrary keyword arguments forwarded to ``xarray.open_dataset()``, or
            ``xarray.open_zarr()``.
        Refer to the examples.

        Returns
        -------
        package : imod.mf6.Package
            Returns a package with data loaded from file.

        Examples
        --------

        To load a package from a file, e.g. a River package:

        >>> river = imod.mf6.River.from_file("river.nc")

        For large datasets, you likely want to process it in chunks. You can
        forward keyword arguments to ``xarray.open_dataset()`` or
        ``xarray.open_zarr()``:

        >>> river = imod.mf6.River.from_file("river.nc", chunks={"time": 1})

        Refer to the xarray documentation for the possible keyword arguments.
        """
        path = Path(path)
        if path.suffix in (".zip", ".zarr"):
            # TODO: seems like a bug? Remove str() call if fixed in xarray/zarr
            dataset = xr.open_zarr(str(path), **kwargs)
        else:
            dataset = xr.open_dataset(path, **kwargs)

        if dataset.ugrid_roles.topology:
            dataset = xu.UgridDataset(dataset)
            dataset = imod.util.spatial.from_mdal_compliant_ugrid2d(dataset)
            # Drop node dim, as we don't need in the UgridDataset (it will be
            # preserved in the ``grid`` property, so topology stays intact!)
            node_dim = dataset.grid.node_dimension
            dataset = dataset.drop_dims(node_dim, errors="ignore")

        # Replace NaNs by None
        for key, value in dataset.items():
            if _is_scalar_nan(value):
                dataset[key] = None

        return cls._from_dataset(dataset)
