# noqa: D100
import datetime
import operator
import os
import warnings
from copy import deepcopy
from pathlib import PosixPath
from typing import Optional, Union

import numpy as np
import xarray as xr
import xesmf as xe

from .config import parse_config

# TODO: Implement logging, warnings, etc.
# TODO: Change all paths to PosixPath objects, including in the catalog?
# TODO: Add an option to call xesmf.util.grid_2d or xesmf.util.grid_global
# TODO: Implement support for an OBS2SIM kind of interpolation


__all__ = ["regrid_dataset", "create_mask"]


@parse_config
def regrid_dataset(
    ds: xr.Dataset,
    ds_grid: xr.Dataset,
    weights_location: Union[str, PosixPath],
    *,
    regridder_kwargs: Optional[dict] = None,
    intermediate_grids: Optional[dict] = None,
    to_level: str = "regridded",
) -> xr.Dataset:
    """Regrid a dataset according to weights and a reference grid.

    Based on an intake_esm catalog, this function performs regridding on Zarr files.

    Parameters
    ----------
    ds : xarray.Dataset
        Dataset to regrid. The Dataset needs to have lat/lon coordinates.
        Supports a 'mask' variable compatible with ESMF standards.
    weights_location : Union[str, PosixPath]
        Path to the folder where weight file is saved.
    ds_grid : xr.Dataset
        Destination grid. The Dataset needs to have lat/lon coordinates.
        Supports a 'mask' variable compatible with ESMF standards.
    regridder_kwargs : dict
        Arguments to send xe.Regridder(). If it contains `skipna`, that
        one is passed to the regridder call directly.
    intermediate_grids : dict
        This argument is used to do a regridding in many steps, regridding to regular
        grids before regridding to the final ds_grid.
        This is useful when there is a large jump in resolution between ds and ds grid.
        The format is a nested dictionary shown in Notes.
        If None, no intermediary grid is used, there is only a regrid from ds to ds_grid.
    to_level : str
        The processing level to assign to the output.
        Defaults to 'regridded'

    Returns
    -------
    xarray.Dataset
        Regridded dataset

    Notes
    -----
    intermediate_grids =
      {'name_of_inter_grid_1': {'cf_grid_2d': {arguments for util.cf_grid_2d },'regridder_kwargs':{arguments for xe.Regridder}},
        'name_of_inter_grid_2': dictionary_as_above}

    See Also
    --------
    xesmf.regridder, xesmf.util.cf_grid_2d
    """
    regridder_kwargs = regridder_kwargs or {}

    ds_grids = []  # list of target grids
    reg_arguments = []  # list of accompanying arguments for xe.Regridder()
    if intermediate_grids:
        for name_inter, dict_inter in intermediate_grids.items():
            reg_arguments.append(dict_inter["regridder_kwargs"])
            ds_grids.append(xe.util.cf_grid_2d(**dict_inter["cf_grid_2d"]))

    ds_grids.append(ds_grid)  # add final ds_grid
    reg_arguments.append(regridder_kwargs)  # add final regridder_kwargs

    out = None

    # Whether regridding is required
    if ds["lon"].equals(ds_grid["lon"]) & ds["lat"].equals(ds_grid["lat"]):
        out = ds
        if "mask" in out:
            out = out.where(out.mask == 1)
            out = out.drop_vars(["mask"])

    else:
        for i, (ds_grid, regridder_kwargs) in enumerate(zip(ds_grids, reg_arguments)):
            # if this is not the first iteration (out != None),
            # get result from last iteration (out) as input
            ds = out or ds

            kwargs = deepcopy(regridder_kwargs)
            # if weights_location does no exist, create it
            if not os.path.exists(weights_location):
                os.makedirs(weights_location)
            id = ds.attrs["cat:id"] if "cat:id" in ds.attrs else "weights"
            # give unique name to weights file
            weights_filename = os.path.join(
                weights_location,
                f"{id}_regrid{i}"
                f"{'_'.join(kwargs[k] for k in kwargs if isinstance(kwargs[k], str))}.nc",
            )

            # TODO: Support for conservative regridding (use xESMF to add corner information), Locstreams, etc.

            # Re-use existing weight file if possible
            if os.path.isfile(weights_filename) and not (
                ("reuse_weights" in kwargs) and (kwargs["reuse_weights"] is False)
            ):
                kwargs["weights"] = weights_filename
                kwargs["reuse_weights"] = True
            skipna = regridder_kwargs.pop("skipna", False)
            regridder = _regridder(
                ds_in=ds, ds_grid=ds_grid, filename=weights_filename, **regridder_kwargs
            )

            # The regridder (when fed Datasets) doesn't like if 'mask' is present.
            if "mask" in ds:
                ds = ds.drop_vars(["mask"])
            out = regridder(ds, keep_attrs=True, skipna=skipna)

            # double-check that grid_mapping information is transferred
            gridmap_out = any(
                "grid_mapping" in ds_grid[da].attrs for da in ds_grid.data_vars
            )
            if gridmap_out:
                gridmap = np.unique(
                    [
                        ds_grid[da].attrs["grid_mapping"]
                        for da in ds_grid.data_vars
                        if "grid_mapping" in ds_grid[da].attrs
                        and ds_grid[da].attrs["grid_mapping"] in ds_grid
                    ]
                )
                if len(gridmap) != 1:
                    warnings.warn(
                        "Could not determine and transfer grid_mapping information."
                    )
                else:
                    # Add the grid_mapping attribute
                    for v in out.data_vars:
                        out[v].attrs["grid_mapping"] = gridmap[0]
                    # Add the grid_mapping coordinate
                    if gridmap[0] not in out:
                        out = out.assign_coords({gridmap[0]: ds_grid[gridmap[0]]})
                    # Regridder seems to seriously mess up the rotated dimensions
                    for d in out.lon.dims:
                        out[d] = ds_grid[d]
                        if d not in out.coords:
                            out = out.assign_coords({d: ds_grid[d]})
            else:
                gridmap = np.unique(
                    [
                        ds[da].attrs["grid_mapping"]
                        for da in ds.data_vars
                        if "grid_mapping" in ds[da].attrs
                    ]
                )
                # Remove the original grid_mapping attribute
                for v in out.data_vars:
                    if "grid_mapping" in out[v].attrs:
                        out[v].attrs.pop("grid_mapping")
                # Remove the original grid_mapping coordinate if it is still in the output
                out = out.drop_vars(set(gridmap).intersection(out.variables))

            # History
            kwargs_for_hist = deepcopy(regridder_kwargs)
            kwargs_for_hist.setdefault("method", regridder.method)
            if intermediate_grids and i < len(intermediate_grids):
                name_inter = list(intermediate_grids.keys())[i]
                cf_grid_2d_args = intermediate_grids[name_inter]["cf_grid_2d"]
                new_history = (
                    f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "
                    f"regridded with regridder arguments {kwargs_for_hist} to a xesmf"
                    f" cf_grid_2d with arguments {cf_grid_2d_args}  - xESMF v{xe.__version__}"
                )
            else:
                new_history = (
                    f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] "
                    f"regridded with arguments {kwargs_for_hist} - xESMF v{xe.__version__}"
                )
            history = (
                f"{new_history}\n{out.attrs['history']}"
                if "history" in out.attrs
                else new_history
            )
            out.attrs["history"] = history

    out = out.drop_vars("latitude_longitude", errors="ignore")
    # Attrs
    out.attrs["cat:processing_level"] = to_level
    out.attrs["cat:domain"] = (
        ds_grid.attrs["cat:domain"] if "cat:domain" in ds_grid.attrs else None
    )
    return out


@parse_config
def create_mask(ds: Union[xr.Dataset, xr.DataArray], mask_args: dict) -> xr.DataArray:
    """Create a 0-1 mask based on incoming arguments.

    Parameters
    ----------
    ds : xr.Dataset or xr.DataArray
        Dataset or DataArray to be evaluated
    mask_args : dict
        Instructions to build the mask (required fields listed in the Notes).

    Note
    ----
    'mask' fields:
        variable: str, optional
            Variable on which to base the mask, if ds_mask is not a DataArray.
        where_operator: str, optional
            Conditional operator such as '>'
        where_threshold: str, optional
            Value threshold to be used in conjunction with where_operator.
        mask_nans: bool
            Whether to apply a mask on NaNs.

    Returns
    -------
    xr.DataArray
        Mask array.
    """
    # Prepare the mask for the destination grid
    ops = {
        "<": operator.lt,
        "<=": operator.le,
        "==": operator.eq,
        "!=": operator.ne,
        ">=": operator.ge,
        ">": operator.gt,
    }

    def cmp(arg1, op, arg2):
        operation = ops.get(op)
        return operation(arg1, arg2)

    mask_args = mask_args or {}
    if isinstance(ds, xr.DataArray):
        mask = ds
    elif isinstance(ds, xr.Dataset) and "variable" in mask_args:
        mask = ds[mask_args["variable"]]
    else:
        raise ValueError("Could not determine what to base the mask on.")

    if "time" in mask.dims:
        mask = mask.isel(time=0)

    if "where_operator" in mask_args:
        mask = xr.where(
            cmp(mask, mask_args["where_operator"], mask_args["where_threshold"]), 1, 0
        )
    else:
        mask = xr.ones_like(mask)
    if ("mask_nans" in mask_args) & (mask_args["mask_nans"] is True):
        mask = mask.where(np.isreal(mask), other=0)

    # Attributes
    if "where_operator" in mask_args:
        mask.attrs[
            "where_threshold"
        ] = f"{mask_args['variable']} {mask_args['where_operator']} {mask_args['where_threshold']}"
    mask.attrs["mask_nans"] = f"{mask_args['mask_nans']}"

    return mask


def _regridder(
    ds_in: xr.Dataset,
    ds_grid: xr.Dataset,
    filename: str,
    *,
    method: Optional[str] = "bilinear",
    unmapped_to_nan: Optional[bool] = True,
    **kwargs,
) -> xe.frontend.Regridder:
    """Call to xesmf Regridder with a few default arguments.

    Parameters
    ----------
    ds_in : xr.Dataset
        Incoming grid. The Dataset needs to have lat/lon coordinates.
    ds_grid : xr.Dataset
        Destination grid. The Dataset needs to have lat/lon coordinates.
    filename : str
        Path to the NetCDF file with weights information.
    method : str, optional
        Interpolation method.
    unmapped_to_nan : bool, optional
        Arguments to send xe.Regridder().
    regridder_kwargs : dict
        Arguments to send xe.Regridder().

    Returns
    -------
    xe.frontend.Regridder
        Regridder object
    """
    regridder = xe.Regridder(
        ds_in=ds_in,
        ds_out=ds_grid,
        method=method,
        unmapped_to_nan=unmapped_to_nan,
        **kwargs,
    )
    if ~os.path.isfile(filename):
        regridder.to_netcdf(filename)

    return regridder
