"""Convenience functions which operate on datasets."""

import warnings
from typing import Any, Literal

import cftime as cf
import numpy as np
import xarray as xr
from scipy.interpolate import NearestNDInterpolator

import ilamb3.regions as ilreg
from ilamb3.exceptions import NoSiteDimension, NoUncertainty


def get_dim_name(
    dset: xr.Dataset | xr.DataArray,
    dim: Literal["time", "lat", "lon", "depth", "site"],
) -> str:
    """
    Return the name of the `dim` dimension from the dataset.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The input dataset/dataarray.
    dim : str, one of {`time`, `lat`, `lon`, `depth`, `site`}
        The dimension to find in the dataset/dataarray.

    Returns
    -------
    str
        The name of the dimension.

    See Also
    --------
    get_coord_name : A variant when the coordinate is not a dimension.

    Notes
    -----
    This function is meant to handle the problem that not all data calls the dimensions
    the same things ('lat', 'Lat', 'latitude', etc). We could replace this with
    cf-xarray functionality. My concern is that we want this to work even if the
    datasets are not CF-compliant (e.g. raw model output).
    """
    dim_names = {
        "time": ["time", "TIME", "month"],
        "lat": ["lat", "latitude", "Latitude", "lat_", "Lat", "LATITUDE"],
        "lon": ["lon", "longitude", "Longitude", "lon_", "Lon", "LONGITUDE"],
        "depth": ["depth", "lev"],
    }
    # Assumption: the 'site' dimension is what is left over after all others are removed
    if dim == "site":
        try:
            get_dim_name(dset, "lat")
            get_dim_name(dset, "lon")
            raise NoSiteDimension("Dataset/dataarray is spatial")
        except KeyError:
            pass
        possible_names = list(
            set(dset.dims) - set([d for _, dims in dim_names.items() for d in dims])
        )
        if len(possible_names) == 1:
            return possible_names[0]
        msg = f"Ambiguity in locating a site dimension, found: {possible_names}"
        raise NoSiteDimension(msg)
    possible_names = dim_names[dim]
    dim_name = set(dset.dims).intersection(possible_names)
    if len(dim_name) != 1:
        if len(dim_name) == 0:
            msg = f"{dim} dimension not found: {dset.dims} not in [{','.join(possible_names)}]"
        else:
            msg = f"Ambiguity in determining {dim} dimension: multiple {dset.dims} in [{','.join(possible_names)}]"
        raise KeyError(msg)
    return str(dim_name.pop())


def get_coord_name(
    dset: xr.Dataset | xr.DataArray,
    coord: Literal["lat", "lon"],
) -> str:
    """
    Return the name of the `coord` coordinate from the dataset.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The input dataset/dataarray.
    coord : str, one of {`lat`, `lon`}
        The coordinate to find in the dataset/dataarray.

    Returns
    -------
    str
        The name of the coordinate.

    See Also
    --------
    get_dim_name : A variant when the coordinate is a dimension.
    """
    coord_names = {
        "lat": ["lat", "latitude", "Latitude", "lat_", "Lat", "LATITUDE"],
        "lon": ["lon", "longitude", "Longitude", "lon_", "Lon", "LONGITUDE"],
    }
    possible_names = coord_names[coord]
    coord_name = set(dset.coords).intersection(possible_names)
    if len(coord_name) != 1:
        msg = f"'{coord}' coordinate not found or ambiguous: {coord_name=}"
        raise KeyError(msg)
    return str(coord_name.pop())


def is_temporal(da: xr.DataArray) -> bool:
    """
    Return if the dataarray is temporal.

    Parameters
    ----------
    da : xr.DataArray
        The input dataarray.

    Returns
    -------
    bool
        True if time dimension is present, False otherwise.
    """
    try:
        get_dim_name(da, "time")
        return True
    except KeyError:
        pass
    return False


def is_spatial(da: xr.DataArray) -> bool:
    """
    Return if the dataarray is spatial.

    Parameters
    ----------
    da : xr.DataArray
        The input dataarray.

    Returns
    -------
    bool
        True if latitude and longitude dimensions are present, False otherwise.
    """
    try:
        get_dim_name(da, "lat")
        get_dim_name(da, "lon")
        return True
    except KeyError:
        pass
    # Ocean grids have 2D lat/lons sometimes stored in their coordinates
    if is_latlon2d(da):
        return True
    return False


def is_site(da: xr.DataArray) -> bool:
    """
    Return if the dataarray is a collection of sites.

    Parameters
    ----------
    da : xr.DataArray
        The input dataarray.

    Returns
    -------
    bool
        True if sites, False otherwise.
    """
    try:
        dim_lat = get_dim_name(da, "lat")
        dim_lon = get_dim_name(da, "lon")
    except KeyError:
        dim_lat = dim_lon = None
    try:
        coord_lat = get_coord_name(da, "lat")
        coord_lon = get_coord_name(da, "lon")
    except KeyError:
        return False
    if (
        dim_lat is None
        and coord_lat is not None
        and dim_lon is None
        and coord_lon is not None
    ):
        return True
    return False


def is_layered(da: xr.DataArray) -> bool:
    """
    Return if the dataarray is layered.

    Parameters
    ----------
    da : xr.DataArray
        The input dataarray.

    Returns
    -------
    bool
        True if a depth dimension is present, False otherwise.
    """
    try:
        get_dim_name(da, "depth")
        return True
    except KeyError:
        pass
    return False


def get_integer_dims(da: xr.DataArray) -> list[str]:
    """
    Return which dimensions are integers and therefore likely indices.
    """
    return [d for d in da.dims if d in da.coords and da[d].dtype.kind == "i"]


def is_latlon2d(da: xr.DataArray) -> bool:
    """
    Return if the dataarray has 2D latitudes and longitudes.

    Note
    ----
    To be considered 2d lat/lon, the dataset must have 1d indexing arrays of
    integer type and then 2d latitude/longitude coordinates in those indexing
    dimensions.
    """
    index_dims = get_integer_dims(da)
    if not index_dims:
        return False
    try:
        lat_coord = get_coord_name(da, "lat")
        lon_coord = get_coord_name(da, "lon")
    except KeyError:
        return False
    # Are the index dimensions of da also the dimensions of the the coordinates?
    if set(da[lon_coord].dims) == set(da[lat_coord].dims) == set(index_dims):
        return True
    return False


def latlon2d_to_1d(grid: xr.Dataset | xr.DataArray, da: xr.DataArray) -> xr.DataArray:
    """
    Return the da with 2D lat/lon interpolated at the grid resolution.
    """
    if not is_latlon2d(da):
        raise ValueError("Input dataarray is not a 2D lat/lon.")
    lat_dim = get_dim_name(grid, "lat")
    lon_dim = get_dim_name(grid, "lon")
    lat_coord = get_coord_name(da, "lat")
    lon_coord = get_coord_name(da, "lon")
    index_dims = get_integer_dims(da)
    if len(index_dims) != 2:
        raise ValueError(f"Expected 2 index dimension but found {index_dims=}")
    da[lon_coord] = xr.where(da[lon_coord] > 180, da[lon_coord] - 360, da[lon_coord])
    # Some models start the indices at 1 and not 0
    index_starting_number = int(da[index_dims[0]].min())
    if index_starting_number not in [0, 1]:
        raise ValueError(f"Unexpected {index_starting_number=}")
    # Create an interpolator that maps lat,lon --> i,j
    index_map = NearestNDInterpolator(
        np.array([da[lat_coord].values.flatten(), da[lon_coord].values.flatten()]).T,
        np.array(
            [
                (da[index_dims[0]] * xr.ones_like(da[index_dims[1]])).values.flatten(),
                (xr.ones_like(da[index_dims[0]]) * da[index_dims[1]]).values.flatten(),
            ]
        ).T,
    )
    # Then use the interpolator to find index arrays at the input regular grid.
    ids = index_map(
        (grid[lat_dim] * xr.ones_like(grid[lon_dim])).values.flatten(),
        (xr.ones_like(grid[lat_dim]) * grid[lon_dim]).values.flatten(),
    ).astype(int)
    ids -= index_starting_number

    # Create a new dataarray
    not_index_dims = [d for d in da.dims if d not in index_dims]
    coords = {c: da[c] for c in not_index_dims if c in da.coords}
    coords.update({"lat": grid["lat"], "lon": grid["lon"]})
    shp = (grid[lat_dim].size, grid[lon_dim].size)
    out = xr.DataArray(
        data=da.to_numpy()[..., ids[:, 0].reshape(shp), ids[:, 1].reshape(shp)],
        coords=coords,
        dims=not_index_dims + [lat_dim, lon_dim],
        attrs=da.attrs,
    )
    return out


def get_time_extent(
    dset: xr.Dataset | xr.DataArray, include_bounds: bool = True
) -> tuple[xr.DataArray, xr.DataArray]:
    """
    Return the time extent of the dataset/dataarray.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The input dataset.

    Returns
    -------
    tuple of xr.DataArray
        The minimum and maximum time.

    Notes
    -----
    The function will prefer the values in the 'bounds' array if present.
    """
    time_name = get_dim_name(dset, "time")
    time = dset[time_name]
    if "bounds" in time.attrs and include_bounds:
        if time.attrs["bounds"] in dset:
            time = dset[time.attrs["bounds"]]
    time.load()
    return time.min(), time.max()


def compute_time_measures(dset: xr.Dataset | xr.DataArray) -> xr.DataArray:
    """
    Return the length of each time interval.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The input dataset.

    Returns
    -------
    xr.DataArray
        The time measures, the length of the time intervals, in days.

    Notes
    -----
    In order to integrate in time, we need the time measures. While this function is
    written for greatest flexibility, the most accurate time measures will be computed
    when a dataset is passed in where the 'bounds' on the 'time' dimension are labeled
    and part of the dataset.
    """

    def _measure1d(time, time_name):  # numpydoc ignore=GL08
        if time.size == 1:
            msg = "Cannot estimate time measures from single value without bounds"
            raise ValueError(msg)
        # compute time measures as differences (timedelta64[ns]) converted to days
        delt = time.diff(dim=time_name).to_numpy().astype(float) * 1e-9 / 3600 / 24
        # pad the array to match the original time array size
        delt = np.hstack([delt[0], delt, delt[-1]])
        # compute the average measure in days for each time step
        msr = xr.DataArray(
            0.5 * (delt[:-1] + delt[1:]), coords=[time], dims=[time_name]
        )
        msr.attrs["units"] = "d"
        return msr

    time_name = get_dim_name(dset, "time")
    time = dset[time_name]
    timeb_name = time.attrs["bounds"] if "bounds" in time.attrs else None
    if timeb_name is None or timeb_name not in dset:
        return _measure1d(time, time_name)
    # compute from the bounds
    delt = dset[timeb_name]
    nbnd = delt.dims[-1]
    delt = delt.diff(nbnd).squeeze().compute()
    measure = delt.astype("float") * 1e-9 / 86400  # .diff [ns] to [d]
    measure.attrs["units"] = "d"
    return measure


def has_bounds(dset: xr.Dataset | xr.DataArray, dim: str, bounds: str | None) -> bool:
    """
    Do the bounds for the input dataset and dimension exist?
    """
    if bounds is None:
        return False
    if isinstance(dset, xr.DataArray):
        return False
    if dim not in dset.dims:
        return False
    if "bounds" not in dset[dim].attrs:
        return False
    if bounds != dset[dim].attrs["bounds"]:
        raise ValueError(
            f"Given {bounds=} is not the same as the dimension attribute {dset[dim].attrs['bounds']}"
        )
    if bounds in dset:
        return True
    return False


def get_ancillary_variable(ds: xr.Dataset, varname: str) -> xr.Dataset | None:
    """
    Return the sub-dataset with the ancillary variables found in ds[varname].
    """
    var = ds[varname]
    if "ancillary_variables" not in var.attrs:
        return None
    anc_vars = [v for v in var.attrs["ancillary_variables"].split() if v in ds]
    if not anc_vars:
        raise ValueError(
            f"Ancillary variables specified {var.attrs['ancillary_variables']=} but at least one was not found in the dataset {ds=}."
        )
    return ds[anc_vars]


def get_bounds_variable(ds: xr.Dataset, varname: str) -> xr.DataArray | None:
    """
    Return the dataarray defined as the bounds of the input variable.
    """
    var = ds[varname]
    if "bounds" not in var.attrs:
        return None
    bnd_var = var.attrs["bounds"]
    if bnd_var not in ds:
        raise ValueError(
            f"A bound variable is specified {var.attrs['bounds']=} but was not found in the dataset {ds=}."
        )
    return ds[bnd_var]


def compute_cell_measures(dset: xr.Dataset | xr.DataArray) -> xr.DataArray:
    """
    Return the area of each spatial cell.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The input dataset.

    Returns
    -------
    xr.DataArray
        The cell areas.

    Notes
    -----
    It would be better to get these from the model data itself, but they are not always
    provided, particularly in reference data.
    """
    earth_radius = 6.371e6  # [m]
    lat_name = get_dim_name(dset, "lat")
    lon_name = get_dim_name(dset, "lon")
    lat = dset[lat_name]
    lon = dset[lon_name]
    latb_name = lat.attrs["bounds"] if "bounds" in lat.attrs else None
    lonb_name = lon.attrs["bounds"] if "bounds" in lon.attrs else None
    # we prefer to compute your cell areas from the lat/lon bounds if they are
    # part of the dataset...
    if has_bounds(dset, lat_name, latb_name) & has_bounds(dset, lon_name, lonb_name):
        delx = dset[lonb_name] * np.pi / 180
        dely = np.sin(dset[latb_name] * np.pi / 180)
        other_dims = delx.dims[-1]
        delx = earth_radius * delx.diff(other_dims).squeeze()
        dely = earth_radius * dely.diff(other_dims).squeeze()  # type: ignore
        msr = dely * delx
        msr.attrs["units"] = "m2"
        return msr
    # ...and if they aren't, we assume the lat/lon we have is a cell centroid
    # and compute the area.
    lon = lon.values
    lat = lat.values
    delx = 0.5 * (lon[:-1] + lon[1:])
    dely = 0.5 * (lat[:-1] + lat[1:])
    delx = np.vstack(
        [
            np.hstack([lon[0] - 0.5 * (lon[1] - lon[0]), delx]),
            np.hstack([delx, lon[-1] + 0.5 * (lon[-1] - lon[-2])]),
        ]
    ).T
    dely = np.vstack(
        [
            np.hstack([lat[0] - 0.5 * (lat[1] - lat[0]), dely]),
            np.hstack([dely, lat[-1] + 0.5 * (lat[-1] - lat[-2])]),
        ]
    ).T
    delx = delx * np.pi / 180
    dely = np.sin(dely * np.pi / 180)
    delx = earth_radius * np.diff(delx, axis=1).squeeze()
    dely = earth_radius * np.diff(dely, axis=1).squeeze()
    delx = xr.DataArray(
        data=np.abs(delx), dims=[lon_name], coords={lon_name: dset[lon_name]}
    )
    dely = xr.DataArray(
        data=np.abs(dely), dims=[lat_name], coords={lat_name: dset[lat_name]}
    )
    msr = dely * delx
    msr.attrs["units"] = "m2"
    return msr


def coarsen_dataset(dset: xr.Dataset, varname: str, res: float = 0.5) -> xr.Dataset:
    """
    Return the convervative spatially coarsened dataset.

    Parameters
    ----------
    dset : xr.Dataset
        The input dataset.
    varname : str
        The name of the dataarray to be coarsened.
    res : float, optional
        The target resolution in degrees.

    Returns
    -------
    xr.Dataset
        The coarsened dataset.
    """
    assert varname in dset
    # This coarsening does not permit you to give the lats/lons, only the resolution
    lat_name = get_dim_name(dset, "lat")
    lon_name = get_dim_name(dset, "lon")
    ds_res = np.sqrt(
        dset[lat_name].diff(lat_name).mean().values ** 2
        + dset[lon_name].diff(lon_name).mean().values ** 2
    )
    if res < ds_res:
        raise ValueError(
            f"Cannot coarsen as the dataset resolution was already coarser than the target resolution: {res} < {ds_res}"
        )
    fine_per_coarse = int(round(res / ds_res))
    # In order for our coarsening to be conservative, we will need to use cell measures
    if "cell_measures" not in dset:
        dset["cell_measures"] = compute_cell_measures(dset)
    # The data to be coarsened may contain nan's. If a variable is nan for all
    # dimensions not part of the measures (usually just 'time') then make the
    # measure also nan.
    msr = dset["cell_measures"]
    var = dset[varname]
    other_dims = set(var.dims).difference(msr.dims)
    dset["cell_measures"] = msr.where(~(var.isnull().all(dim=other_dims)), drop=True)
    # Sometimes variables have nan's in some but not all slices. For noe we will
    # detect this and warn. A fix will require extending the measures into those
    # other dimensions and then masking as data requires.
    if ((~var.isnull().all(dim=other_dims)) & (var.isnull().any(dim=other_dims))).sum():
        warnings.warn(
            f"Non-uniformly distributed nan's found in {other_dims} dimensions. Coarsened result may not be conservative."
        )

    # The 'weighted' accessor does not work with coarsen and so we will do a
    # weighted average here manually.
    dset[varname] *= dset["cell_measures"]
    dset_coarse = dset.coarsen(
        {"lat": fine_per_coarse, "lon": fine_per_coarse}, boundary="pad"
    ).sum()
    dset[varname] /= dset["cell_measures"]
    dset_coarse[varname] /= dset_coarse["cell_measures"]
    return dset_coarse


def integrate_time(
    dset: xr.Dataset | xr.DataArray,
    varname: str | None = None,
    mean: bool = False,
) -> xr.DataArray:
    """
    Return the time integral or mean of the dataset.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The input dataset/dataarray.
    varname : str, optional
        The variable to integrate, must be given if a dataset is passed in.
    mean : bool, optional
        Enable to divide the integral by the integral of the measures, returning the
        mean in a functional sense.

    Returns
    -------
    integral
        The integral or mean as a DataArray.

    Notes
    -----
    This interface is useful in our analysis as many times we want to report the total
    of a quantity (total mass of carbon) and other times we want the mean value (e.g.
    temperature). This allows the analysis code to read the same where a flag can be
    passed to change the behavior. We could consider replacing with xarray.integrate.
    However, as of `v2023.6.0`, this does not handle the `pint` units correctly, and can
    only be done in a single dimension at a time, leaving the spatial analog to be hand
    coded. It also uses trapezoidal rule which should return the same integration, but
    could have small differences depending on how endpoints are interpretted.
    """
    time_name = get_dim_name(dset, "time")
    if isinstance(dset, xr.Dataset):
        assert varname is not None
        var = dset[varname]
        msr = (
            dset["time_measures"]
            if "time_measures" in dset
            else compute_time_measures(dset)
        )
    else:
        var = dset
        msr = compute_time_measures(dset)
    if mean:
        return var.weighted(msr.fillna(0)).mean(dim=time_name)
    var = var.pint.quantify()
    msr = msr.pint.quantify()
    out = var.weighted(msr.fillna(0)).sum(dim=time_name)
    out = out.pint.dequantify()
    return out


def accumulate_time(
    dset: xr.Dataset | xr.DataArray,
    varname: str | None = None,
) -> xr.DataArray:
    """
    Return the time accumulation of the dataset.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The input dataset/dataarray.
    varname : str, optional
        The variable to integrate, must be given if a dataset is passed in.

    Returns
    -------
    xr.DataArray
        The accumulation.
    """
    time_name = get_dim_name(dset, "time")
    if isinstance(dset, xr.Dataset):
        assert varname is not None
        var = dset[varname]
        msr = (
            dset["time_measures"]
            if "time_measures" in dset
            else compute_time_measures(dset)
        )
    else:
        var = dset
        msr = compute_time_measures(dset)
    var = (var.pint.quantify() * msr.pint.quantify()).pint.dequantify()
    var = var.cumsum(dim=time_name)
    return var


def std_time(
    dset: xr.Dataset | xr.DataArray, varname: str | None = None
) -> xr.DataArray:
    """
    Return the standard deviation of a variable in time.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The input dataset/dataarray.
    varname : str, optional
        The variable, must be given if a dataset is passed in.

    Returns
    -------
    xr.DataArray
        The weighted standard deviation.
    """
    time_name = get_dim_name(dset, "time")
    if isinstance(dset, xr.Dataset):
        var = dset[varname]
        msr = (
            dset["time_measures"]
            if "time_measures" in dset
            else compute_time_measures(dset)
        )
    else:
        var = dset
        msr = compute_time_measures(dset)
    return var.weighted(msr.fillna(0)).std(dim=time_name)


def integrate_space(
    dset: xr.DataArray | xr.Dataset,
    varname: str,
    region: None | str = None,
    mean: bool = False,
    weight: xr.DataArray | None = None,
) -> xr.DataArray:
    """
    Return the space integral or mean of the dataset.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The input dataset/dataarray.
    varname : str, optional
        The variable to integrate, must be given if a dataset is passed in.
    region : str, optional
        The region label, one of `ilamb3.Regions.regions` or `None` to indicate that the
        whole spatial domain should be used.
    mean : bool, optional
        Enable to divide the integral by the integral of the measures, returning the
        mean in a functional sense.
    weight : xr.DataArray, optional
        Optional weight for the spatial integral. Used when mass weighting.

    Returns
    -------
    xr.Dataset
        The integral or mean.

    Notes
    -----
    This interface is useful in our analysis as many times we want to report the total
    of a quantity (total mass of carbon) and other times we want the mean value (e.g.
    temperature). This allows the analysis code to read the same where a flag can be
    passed to change the behavior. We could consider replacing with xarray.integrate.
    However, as of `v2023.6.0`, this does not handle the `pint` units correctly, and can
    only be done in a single dimension at a time.
    """
    dset = dset.pint.dequantify()
    if region is not None:
        regions = ilreg.Regions()
        dset = regions.restrict_to_region(dset, region)
    if not isinstance(dset, xr.Dataset):
        dset = dset.to_dataset(name=varname)
    var = dset[varname]
    space = [get_dim_name(var, "lat"), get_dim_name(var, "lon")]
    msr = (
        dset["cell_measures"]
        if "cell_measures" in dset
        and set(dset["cell_measures"].dims).issubset(var.dims)
        else compute_cell_measures(var)
    )
    if weight is not None:
        assert isinstance(weight, xr.DataArray)
        msr = (msr.pint.quantify() * weight.pint.quantify()).pint.dequantify()
    out = var.weighted(msr.fillna(0))
    if mean:
        out = out.mean(dim=space)
        out.attrs["units"] = var.attrs["units"]
    else:
        out = out.sum(dim=space)
        out.attrs["units"] = f"({var.attrs['units']})*({msr.attrs['units']})"
    return out


def sel(dset: xr.Dataset, coord: str, cmin: Any, cmax: Any) -> xr.Dataset:
    """
    Return a selection of the dataset.

    Parameters
    ----------
    dset : xr.Dataset
        The input dataset.
    coord : str
        The coordinate to slice.
    cmin :  Any
        The minimum coordinate value.
    cmax : Any
        The maximum coordinate value.

    Returns
    -------
    xr.Dataset
        The selected dataset.

    Notes
    -----
    The behavior of xarray.sel does not work for us here. We want to pick a slice of the
    dataset but where the value lies in between the coord bounds. Then we clip the min
    and max to be the limits of the slice.
    """

    def _get_interval(dset, dim, value, side):  # numpydoc ignore=GL08
        coord = dset[dim]
        if "bounds" in coord.attrs:
            if coord.attrs["bounds"] in dset:
                coord = dset[coord.attrs["bounds"]]
                coord.load()
                ind = ((coord[:, 0] <= value) & (coord[:, 1] >= value)).to_numpy()
                ind = np.where(ind)[0]
                if len(ind) not in [1, 2]:
                    raise ValueError(
                        f"We could not find an appropriate interval for {value} in \n{coord}"
                    )
                if len(ind) == 2:
                    if side == "low":
                        return ind[1]
                return ind[0]
        raise NotImplementedError(f"Must have a bounds {coord}")

    dset = dset.isel(
        {
            coord: slice(
                _get_interval(dset, coord, cmin, "low"),
                _get_interval(dset, coord, cmax, "high") + 1,
            )
        }
    )
    # adjust the bounds and coord values
    bnds = dset[coord].attrs["bounds"]
    dset[bnds][0, 0] = cmin
    dset[bnds][-1, 1] = cmax
    dim = dset[coord].to_numpy()
    dim[0] = dset[bnds][0, 0].values + 0.5 * (
        dset[bnds][0, 1].values - dset[bnds][0, 0].values
    )
    dim[-1] = dset[bnds][-1, 0].values + 0.5 * (
        dset[bnds][-1, 1].values - dset[bnds][-1, 0].values
    )
    attrs = dset[coord].attrs
    dset[coord] = dim
    dset[coord].attrs = attrs
    return dset


def integrate_depth(
    dset: xr.Dataset | xr.DataArray,
    varname: str | None = None,
    mean: bool = False,
) -> xr.DataArray:
    """
    Return the depth integral or mean of the dataset.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The dataset of dataarray to integrate.
    varname : str, optional
        If dset is a dataset, the variable name to integrate.
    mean : bool, optional
        Enable to take a depth mean.

    Returns
    -------
    xr.DataArray
        The depth integral or sum.
    """
    if isinstance(dset, xr.DataArray):
        varname = dset.name
        dset = dset.to_dataset(name=varname)
    else:
        assert varname is not None
    var = dset[varname]

    # do we have a depth dimension
    try:
        depth_name = get_dim_name(dset, "depth")
    except KeyError:
        raise ValueError("Cannot integrate in depth without a depth dimension.")

    # does depth have bounds?
    if (
        "bounds" not in dset[depth_name].attrs
        or dset[depth_name].attrs["bounds"] not in dset
    ):
        dset = dset.cf.add_bounds(depth_name)

    # compute measures
    msr = dset[dset[depth_name].attrs["bounds"]]
    msr = msr.diff(dim=msr.dims[-1])
    msr.attrs["units"] = (
        dset[depth_name].attrs["units"] if "units" in dset[depth_name].attrs else "m"
    )

    # integrate
    out = var.weighted(msr.fillna(0))
    if mean:
        out = out.mean(dim=depth_name)
        out.attrs["units"] = var.attrs["units"]
    else:
        out = out.sum(dim=depth_name)
        out.attrs["units"] = f"({var.attrs['units']})*({msr.attrs['units']})"
    return out


def scale_by_water_density(da: xr.DataArray, target: str) -> xr.DataArray:
    """
    Conditionally scale the dataarray by density if needed in the conversion.

    Parameters
    ----------
    da : xr.DataArray
        The pint quantified input data array.
    target : str
        The target conversion unit as a string.

    Returns
    -------
    xr.DataArray
        The potentially water density scaled data array.

    Notes
    -----
    Most modeled hydrologic quantities tend to be output as a mass flux rate. However, a
    linear speed is often the preferred unit. For example a conversion such as `kg m-2
    s-1` to `mm d-1`. This is possible if scaled by the density of water which we do
    here conditionally if needed. Used by our `convert()` routine.
    """
    da = da.pint.quantify()
    ureg = da.pint.registry
    water_density = 998.2071 * ureg.kilogram / ureg.meter**3
    src = 1.0 * da.pint.units
    tar = ureg(target)
    if (src / tar).check("[mass] / [length]**3"):
        return da / water_density
    if (tar / src).check("[mass] / [length]**3"):
        return da * water_density
    return da


def convert(
    dset: xr.Dataset | xr.DataArray,
    unit: str,
    varname: str | None = None,
) -> xr.Dataset | xr.DataArray:
    """
    Convert the units of the dataarray.

    Parameters
    ----------
    dset : xr.Dataset or xr.DataArray
        The dataset (specify varname) or dataarray who units you wish to convert.
    unit : str
        The unit to which we will convert.
    varname : str, optional
        If dset is a dataset, give the variable name to convert.

    Returns
    -------
    xr.Dataset or xr.DataArray
        The converted dataset.
    """
    if isinstance(dset, xr.DataArray):
        da = dset
    else:
        assert varname is not None
        da = dset[varname]
    da = da.pint.quantify()
    da = scale_by_water_density(da, unit)
    da = da.pint.to(unit)
    if isinstance(dset, xr.DataArray):
        return da.pint.dequantify()
    dset[varname] = da.pint.dequantify()
    return dset


def coarsen_annual(dset: xr.Dataset) -> xr.Dataset:
    """
    Return the dataset coarsened to annual.

    Parameters
    ----------
    dset : xr.Dataset
        The input dataset/dataarray.

    Returns
    -------
    xr.Dataset
        The coarsened dataset.
    """
    # Can't sum time objects so find them and remove
    time_name = get_dim_name(dset, "time")
    if "bounds" in dset[time_name].attrs:
        bounds = dset[time_name].attrs["bounds"]
        if bounds in dset:
            dset = dset.drop_vars(bounds)
    msr = compute_time_measures(dset).pint.dequantify()
    ann = (dset * msr).groupby(f"{time_name}.year").sum() / msr.groupby(
        f"{time_name}.year"
    ).sum()
    return ann


def shift_lon(dset: xr.Dataset) -> xr.Dataset:
    """
    Return the dataset with longitudes shifted to [-180,180].

    Parameters
    ----------
    dset : xr.Dataset
        The input dataset.

    Returns
    -------
    xr.Dataset
        The longitude-shifted dataset.
    """
    lon_name = get_coord_name(dset, "lon")
    if (dset[lon_name].min() >= 0) * (dset[lon_name].max() <= 360):
        dset[lon_name] = (dset[lon_name] + 180) % 360 - 180
        if "bounds" in dset[lon_name].attrs and dset[lon_name].attrs["bounds"] in dset:
            dset[dset[lon_name].attrs["bounds"]] = (
                dset[dset[lon_name].attrs["bounds"]] + 180
            ) % 360 - 180
        dset = dset.sortby(lon_name)
    return dset


def get_scalar_uncertainty(ds: xr.Dataset, varname: str) -> xr.DataArray:
    """
    Get a scalar uncertainty from the variable.

    Note
    ----
    Uncertainty is indicated either by the presence of the `ancillary_variables`
    attribute or the `bounds` attribute. In the case of the latter, we will
    return the harmonic mean as a scalar measure of uncertainty.
    """

    def _get_bound_dim(da: xr.DataArray) -> str | None:
        possible = set(da.dims).difference(da.coords)
        if len(possible) > 1:
            raise ValueError(
                f"Ambiguity in determinging the `bounds` dimension, found: {possible} in\n\n{da}"
            )
        if len(possible) == 0:
            return None
        return possible.pop()

    out = None
    # Uncertainty could be an ancillary variable with the same dims
    ds_anc = get_ancillary_variable(ds, varname)
    if ds_anc is not None and len(ds_anc) == 1:
        # We found an ancillary variable, but does it have the same dims
        out = (
            ds_anc[next(iter(ds_anc))]
            if set(ds[varname].dims) != set(ds_anc.dims)
            else None
        )

    # Uncertainty could be defined using intervals
    if out is None:
        try:
            out = get_interval_uncertainty(ds, varname)
        except NoUncertainty:
            pass

    if out is None:
        raise NoUncertainty()

    # Cleanup if an inverval style uncertainty
    bnd_dim = _get_bound_dim(out)
    if bnd_dim is not None:
        var = ds[varname]
        out = np.sqrt(((var - out) ** 2).sum(dim=bnd_dim))
    if "units" not in out.attrs:
        out.attrs["units"] = var.attrs["units"]
    out.name = "uncert"
    return out


def get_interval_uncertainty(ds: xr.Dataset, varname: str) -> xr.DataArray:
    """
    Get the uncertainty interval from the variable.

    Note
    ----
    Uncertainty is indicated either by the presence of the `ancillary_variables`
    attribute or the `bounds` attribute.
    """

    def _get_bound_dim(da: xr.DataArray) -> str | None:
        possible = set(da.dims).difference(da.coords)
        if len(possible) > 1:
            raise ValueError(
                f"Ambiguity in determinging the `bounds` dimension, found: {possible} in\n\n{da}"
            )
        if len(possible) == 0:
            return None
        return possible.pop()

    # What ancillary/bounds variables are found?
    ds_anc = get_ancillary_variable(ds, varname)
    da_bnd = get_bounds_variable(ds, varname)

    # Handle if we find ancillary variables
    if ds_anc is not None:
        if len(ds_anc) == 1:
            # Only 1 ancillary variable present, but does it represent a bounds?
            da_bnd = ds[next(iter(ds_anc))]
            bnd_dim = _get_bound_dim(da_bnd)
            if bnd_dim is None:
                raise NoUncertainty("Only scalar uncertainty present.")
        elif len(ds_anc) == 2:
            # Assume that the 2 variables found form a bounds and create it
            da_bnd = xr.concat(
                [da for _, da in ds_anc.items()], dim="uncert_bnds"
            ).transpose()
        else:
            raise ValueError(
                f"Uncertainty can only be encoded with 1 or 2 ancillary variables, I found {ds_anc=}."
            )

    # If we get here and this is still none, we have no uncertainty
    if da_bnd is None:
        raise NoUncertainty()

    # Cleanup name and units
    if "units" not in da_bnd.attrs:
        da_bnd.attrs["units"] = ds[varname].attrs["units"]
    da_bnd.name = "uncert"
    return da_bnd


def fix_missing_bounds_attrs(ds: xr.Dataset) -> xr.Dataset:
    """
    Add a `bounds` attribute if variable appears to exist.
    """

    def _fix(ds: xr.Dataset, dim_name: str) -> xr.Dataset:
        dim = ds[dim_name]
        if "bounds" in dim.attrs:
            if dim.attrs["bounds"] in ds:
                return ds  # all good, just return
            else:
                ds[dim_name].attrs.pop("bounds")
                return ds
        # is there a likely bounds variable and the attr is missing?
        for possible_bnds in [f"{dim_name}_bnds", f"{dim_name}_bounds"]:
            if possible_bnds in ds:
                ds[dim_name].attrs["bounds"] = possible_bnds
        return ds

    for d in ds.dims:
        ds = _fix(ds, d)
    return ds


def get_mean_time_frequency(ds: xr.Dataset) -> float:
    """
    Return the mean time frequency of the dataset.
    """
    tm = compute_time_measures(ds)
    return float(tm.mean())


def compute_monthly_mean(ds: xr.Dataset) -> xr.Dataset:
    """
    Return the monthly mean of the input dataset in a `noleap` calendar.
    """
    if not is_temporal(ds):
        raise ValueError("Input dataset has no temporal dimension.")
    dt = get_mean_time_frequency(ds)
    if dt > 31:
        raise ValueError(f"Input dataset is already coarser than monthly {dt=}.")
    ds = ds.resample(time=xr.groupers.TimeResampler("MS")).mean()
    return ds


def compute_annual_mean(ds: xr.Dataset) -> xr.Dataset:
    """
    Return the annual mean of the input dataset.
    """
    if not is_temporal(ds):
        raise ValueError("Input dataset has no temporal dimension.")
    dt = get_mean_time_frequency(ds)
    if dt > 366:
        raise ValueError(f"Input dataset is already coarser than annual {dt=}.")
    ds = ds.resample(time=xr.groupers.TimeResampler("YS")).mean()
    return ds


def shift_time_by_years(ds: xr.Dataset, years: int) -> xr.Dataset:
    """
    Return the dataset shifted by the given number of years.
    """
    time_dim = get_dim_name(ds, "time")
    cls = ds[time_dim].values[0].__class__
    ds[time_dim] = [
        cls(
            t.dt.year + years, t.dt.month, t.dt.day, t.dt.hour, t.dt.minute, t.dt.second
        )
        for t in ds[time_dim]
    ]
    return ds


def convert_year_to_datetime(ds: xr.Dataset) -> xr.Dataset:
    """
    Return the dataset with a integer 'years' dimension converted to a
    CF-compliant time.
    """
    if "year" not in ds.coords:
        raise ValueError("Function expects a 'year' coordinate with integer years.")
    # We are about to reassign the 'time' dimension, let's remove everything
    # with 'time' in it. We are assuming that it is no longer relevant.
    if "time" in ds.coords:
        ds = ds.drop_vars([v for v in ds if "time" in ds[v].dims]).drop_dims("time")
    years = ds["year"]
    ds = ds.rename_dims(year="time").rename_vars(year="time")
    ds["time"] = [cf.DatetimeNoLeap(y, 7, 1) for y in years]
    ds["time_bnds"] = (
        ("time", "nb"),
        np.array(
            [
                [cf.DatetimeNoLeap(y, 1, 1) for y in years],
                [cf.DatetimeNoLeap(y + 1, 1, 1) for y in years],
            ]
        ).T,
    )
    ds["time"].attrs["bounds"] = "time_bnds"
    return ds
