"""Utility functions."""

from __future__ import annotations

import logging
import os
import sys
import tarfile
import warnings
from functools import wraps
from inspect import getfullargspec
from pathlib import Path

import numpy as np
import requests
from scipy.integrate import trapezoid

from pyspectral.bandnames import BANDNAMES
from pyspectral.config import get_config

TQDM_LOADED = True
try:
    from tqdm import tqdm
except ImportError:
    TQDM_LOADED = False


LOG = logging.getLogger(__name__)

WAVE_LENGTH = "wavelength"
WAVE_NUMBER = "wavenumber"

INSTRUMENTS = {"Envisat": "aatsr",
               "GOES-16": "abi",
               "GOES-17": "abi",
               "GOES-18": "abi",
               "GOES-19": "abi",
               "FY-4A": "agri",
               "FY-4B": ["agri", "ghi"],
               "Himawari-8": "ahi",
               "Himawari-9": "ahi",
               "GEO-KOMPSAT-2A": "ami",
               "GEO-KOMPSAT-2B": "goci-2",
               "NOAA-10": "avhrr/1",
               "NOAA-6": "avhrr/1",
               "NOAA-8": "avhrr/1",
               "TIROS-N": "avhrr/1",
               "NOAA-11": "avhrr/2",
               "NOAA-12": "avhrr/2",
               "NOAA-14": "avhrr/2",
               "NOAA-7": "avhrr/2",
               "NOAA-9": "avhrr/2",
               "Metop-A": "avhrr/3",
               "Metop-B": "avhrr/3",
               "Metop-C": "avhrr/3",
               "NOAA-15": "avhrr/3",
               "NOAA-16": "avhrr/3",
               "NOAA-17": "avhrr/3",
               "NOAA-18": "avhrr/3",
               "NOAA-19": "avhrr/3",
               "HY-1C": "cocts",
               "Meteosat-12": "fci",
               "MTG-I1": "fci",
               "Metop-SG-A1": "metimage",
               "EOS-Aqua": "modis",
               "EOS-Terra": "modis",
               "Sentinel-2A": "msi",
               "Sentinel-2B": "msi",
               "Sentinel-2C": "msi",
               "Arctica-M-N1": "msu-gsa",
               "Electro-L-N2": "msu-gs",
               "Sentinel-3A": ["olci", "slstr"],
               "Sentinel-3B": ["olci", "slstr"],
               "Landsat-8": "oli_tirs",
               "Landsat-9": "oli_tirs",
               "Meteosat-10": "seviri",
               "Meteosat-11": "seviri",
               "Meteosat-8": "seviri",
               "Meteosat-9": "seviri",
               "NOAA-20": "viirs",
               "NOAA-21": "viirs",
               "Suomi-NPP": "viirs",
               "FY-3A": ["virr", "mersi-1"],
               "FY-3B": ["virr", "mersi-1"],
               "FY-3C": ["virr", "mersi-1"],
               "FY-3D": "mersi-2",
               "FY-3F": "mersi-3",
               "FY-3G": "mersi-rm",
               "DSCOVR": "epic"}


INSTRUMENT_TRANSLATION_DASH2SLASH = {"avhrr-1": "avhrr/1",
                                     "avhrr-2": "avhrr/2",
                                     "avhrr-3": "avhrr/3"}

HTTP_PYSPECTRAL_RSR = "https://zenodo.org/records/18928854/files/pyspectral_rsr_data.tgz"

RSR_DATA_VERSION_FILENAME = "PYSPECTRAL_RSR_VERSION"
RSR_DATA_VERSION = "v1.5.1"


ATM_CORRECTION_LUT_VERSION = {}
ATM_CORRECTION_LUT_VERSION["antarctic_aerosol"] = {"version": "v1.0.1",
                                                   "filename": "PYSPECTRAL_ATM_CORR_LUT_AA"}
ATM_CORRECTION_LUT_VERSION["continental_average_aerosol"] = {"version": "v1.0.1",
                                                             "filename": "PYSPECTRAL_ATM_CORR_LUT_CAA"}
ATM_CORRECTION_LUT_VERSION["continental_clean_aerosol"] = {"version": "v1.0.1",
                                                           "filename": "PYSPECTRAL_ATM_CORR_LUT_CCA"}
ATM_CORRECTION_LUT_VERSION["continental_polluted_aerosol"] = {"version": "v1.0.1",
                                                              "filename": "PYSPECTRAL_ATM_CORR_LUT_CPA"}
ATM_CORRECTION_LUT_VERSION["desert_aerosol"] = {"version": "v1.0.1",
                                                "filename": "PYSPECTRAL_ATM_CORR_LUT_DA"}
ATM_CORRECTION_LUT_VERSION["marine_clean_aerosol"] = {"version": "v1.0.1",
                                                      "filename": "PYSPECTRAL_ATM_CORR_LUT_MCA"}
ATM_CORRECTION_LUT_VERSION["marine_polluted_aerosol"] = {"version": "v1.0.1",
                                                         "filename": "PYSPECTRAL_ATM_CORR_LUT_MPA"}
ATM_CORRECTION_LUT_VERSION["marine_tropical_aerosol"] = {"version": "v1.0.1",
                                                         "filename": "PYSPECTRAL_ATM_CORR_LUT_MTA"}
ATM_CORRECTION_LUT_VERSION["rural_aerosol"] = {"version": "v1.0.1",
                                               "filename": "PYSPECTRAL_ATM_CORR_LUT_RA"}
ATM_CORRECTION_LUT_VERSION["urban_aerosol"] = {"version": "v1.0.1",
                                               "filename": "PYSPECTRAL_ATM_CORR_LUT_UA"}
ATM_CORRECTION_LUT_VERSION["rayleigh_only"] = {"version": "v1.0.1",
                                               "filename": "PYSPECTRAL_ATM_CORR_LUT_RO"}

#: Aerosol types available as downloadable LUTs for rayleigh correction
AEROSOL_TYPES = ["antarctic_aerosol", "continental_average_aerosol",
                 "continental_clean_aerosol", "continental_polluted_aerosol",
                 "desert_aerosol", "marine_clean_aerosol",
                 "marine_polluted_aerosol", "marine_tropical_aerosol",
                 "rayleigh_only", "rural_aerosol", "urban_aerosol"]

ATMOSPHERES = {"subarctic summer": 4, "subarctic winter": 5,
               "midlatitude summer": 6, "midlatitude winter": 7,
               "tropical": 8, "us-standard": 9}

HTTPS_RAYLEIGH_LUTS = {}
LUT_URL_PREFIX = "https://zenodo.org/record/1288441/files/pyspectral_atm_correction_luts"
for atype in AEROSOL_TYPES:
    name = {"rayleigh_only": "no_aerosol"}.get(atype, atype)
    url = "{prefix}_{name}.tgz".format(prefix=LUT_URL_PREFIX, name=name)
    HTTPS_RAYLEIGH_LUTS[atype] = url


def get_rayleigh_lut_dir(aerosol_type: str) -> Path:
    """Get the rayleigh LUT directory for the specified aerosol type."""
    conf = get_config()
    local_rayleigh_dir = Path(conf.get('rayleigh_dir'))  # type: ignore
    return local_rayleigh_dir / aerosol_type


def convert2wavenumber(rsr):
    """Convert Spectral Responses from wavelength to wavenumber space.

    Take rsr data set with all channels and detectors for an instrument
    each with a set of wavelengths and normalised responses and
    convert to wavenumbers and responses

    :rsr: Relative Spectral Response function (all bands)
    Returns:
      :retv: Relative Spectral Responses in wave number space
      :info: Dictionary with scale (to go convert to SI units) and unit

    """
    retv = {}
    for chname in rsr.keys():  # Go through bands/channels
        retv[chname] = _band2wavenumber(rsr[chname])

    unit = "cm-1"
    si_scale = 100.0
    return retv, {"unit": unit, "si_scale": si_scale}


def _band2wavenumber(ch_rsr):
    band_wavenumbers = {}
    for det in ch_rsr.keys():  # Go through detectors
        if "wavenumber" in ch_rsr[det].keys():
            # Make a copy. Data are already in wave number space
            band_wavenumbers[det] = ch_rsr[det].copy()
            LOG.debug("RSR data already in wavenumber space. No conversion needed.")
            continue
        band_wavenumbers[det] = _detector2wavenumber(ch_rsr[det])

    return band_wavenumbers


def _detector2wavenumber(det_rsr):
    det_wavenumbers = {}
    for key in det_rsr.keys():
        _wavelength2wavenumber(det_rsr, key, det_wavenumbers)
        _collect_responses(det_rsr, key, det_wavenumbers)

    return det_wavenumbers


def _wavelength2wavenumber(det_rsr, key, det_wavenumbers):
    if key == "wavelength":
        # micro meters to cm
        wnum = 1. / (1e-4 * det_rsr[key])
        det_wavenumbers["wavenumber"] = wnum[::-1]


def _collect_responses(det_rsr, key, det_wavenumbers):
    if key == "response":
        # Flip the response array:
        if isinstance(det_rsr[key], dict):
            det_wavenumbers[key] = {}
            for name in det_rsr[key].keys():
                resp = det_rsr[key][name]
                det_wavenumbers[key][name] = resp[::-1]
        else:
            resp = det_rsr[key]
            det_wavenumbers[key] = resp[::-1]


def get_central_wave(wav, resp, weight=1.0):
    """Calculate the central wavelength or the central wavenumber.

    Calculate the central wavelength or the central wavenumber, depending on
    which parameters is input.  On default the weighting funcion is
    f(lambda)=1.0, but it is possible to add a custom weight, e.g. f(lambda) =
    1./lambda**4 for Rayleigh scattering calculations

    """
    # info: {"unit": unit, "si_scale": si_scale}
    # To get the wavelenght/wavenumber in SI units (m or m-1):
    # wav = wav * info["si_scale"]

    # res = np.trapz(resp*wav, wav) / np.trapz(resp, wav)
    # Check if it is a wavelength or a wavenumber and convert to microns or cm-1:
    # This should perhaps be user defined!?
    # if info["unit"].find("-1") > 0:
    # Wavenumber:
    #     res *=
    return trapezoid(resp * wav * weight, wav) / trapezoid(resp * weight, wav)


def get_bandname_from_wavelength(sensor, wavelength, rsr, epsilon=0.1, multiple_bands=False):
    """Get the bandname from h5 rsr provided the approximate wavelength."""
    chdist_min = 2.0
    chfound = []
    for channel in rsr:
        chdist = abs(
            rsr[channel]["det-1"]["central_wavelength"] - wavelength)
        if chdist < chdist_min and chdist < epsilon:
            chfound.append(BANDNAMES.get(sensor, BANDNAMES["generic"]).get(channel, channel))

    if len(chfound) == 1:
        return chfound[0]
    if len(chfound) > 1:
        bstrlist = ["band={}".format(b) for b in chfound]
        if not multiple_bands:
            raise AttributeError("More than one band found with that wavelength! {}".format(str(bstrlist)))
        LOG.debug("More than one band found with requested wavelength: %s", str(bstrlist))
        return chfound
    return None


def sort_data(x_vals, y_vals):
    """Sort the data so that x is monotonically increasing and contains no duplicates."""
    # Sort data
    # (This is needed in particular for EOS-Terra responses, as there are duplicates)
    idxs = np.argsort(x_vals)
    x_vals = x_vals[idxs]
    y_vals = y_vals[idxs]

    # De-duplicate data
    mask = np.r_[True, (np.diff(x_vals) > 0)]
    if not mask.all():
        numof_duplicates = np.repeat(mask, np.equal(mask, False)).shape[0]
        LOG.debug("Number of duplicates in the response function: %d - removing them",
                  numof_duplicates)
    x_vals = x_vals[mask]
    y_vals = y_vals[mask]

    return x_vals, y_vals


def convert2hdf5(ClassIn, platform_name, bandnames, scale=1e-06, detectors=None):
    """Retrieve original RSR data and convert to internal hdf5 format.

    *scale* is the number which has to be multiplied to the wavelength data in
    order to get it in the SI unit meter

    """
    import h5py

    instruments = [ClassIn(bandnames[i], platform_name) for i in range(len(bandnames))]
    instr_name = instruments[0].instrument.replace("/", "")
    filename = os.path.join(instruments[0].output_dir,
                            "rsr_{0}_{1}.h5".format(instr_name,
                                                    platform_name))

    with h5py.File(filename, "w") as h5f:
        _write_global_attrs(h5f, instruments[0], platform_name, bandnames)
        _write_channels(h5f, instruments, bandnames, scale, detectors)


def _write_global_attrs(h5f, instrument, platform_name, bandnames):
    h5f.attrs["description"] = ("Relative Spectral Responses for " +
                                instrument.instrument.upper())
    h5f.attrs["platform_name"] = platform_name
    h5f.attrs["band_names"] = bandnames


def _write_channels(h5f, instruments, bandnames, scale, detectors):
    for i, chname in enumerate(bandnames):
        sensor = instruments[i]
        grp = h5f.create_group(chname)

        # If multiple detectors, assume all have same wavelength range in SRF.
        if detectors is not None:
            _write_multidetector_data(grp, sensor, scale, detectors)
        else:
            _write_single_detector_data(grp, sensor, scale, detectors)


def _write_single_detector_data(grp, sensor, scale, detectors):
    wvl = sensor.rsr["wavelength"][~np.isnan(sensor.rsr["wavelength"])]
    # Store detector central wavelength
    rsp = sensor.rsr["response"][~np.isnan(sensor.rsr["wavelength"])]
    grp.attrs["central_wavelength"] = get_central_wave(wvl, rsp)

    # Store resonse data
    arr = sensor.rsr["response"]
    dset = grp.create_dataset("response", arr.shape, dtype="f")
    dset[...] = arr

    _save_wavelengths(grp, arr, scale)


def _save_wavelengths(grp, arr, scale):
    dset = grp.create_dataset("wavelength", arr.shape, dtype="f")
    dset.attrs["unit"] = "m"
    dset.attrs["scale"] = scale
    dset[...] = arr


def _write_multidetector_data(grp, sensor, scale, detectors):
    # Get wavelength
    wvl = sensor.rsr[detectors[0]]["wavelength"][~np.isnan(sensor.rsr[detectors[0]]["wavelength"])]
    grp.attrs["number_of_detectors"] = len(detectors)

    for cur_det in detectors:
        det_grp = grp.create_group(cur_det)
        rsp = sensor.rsr[cur_det]["response"][~np.isnan(sensor.rsr[cur_det]["wavelength"])]
        det_grp.attrs["central_wavelength"] = get_central_wave(wvl, rsp)
        arr = sensor.rsr[cur_det]["response"]
        dset = det_grp.create_dataset("response", arr.shape, dtype="f")
        dset[...] = arr

    _save_wavelengths(grp, arr, scale)


def download_rsr(dest_dir: str | Path | None = None, dry_run: bool = False) -> None:
    """Download the relative spectral response functions.

    Download the pre-compiled HDF5 formatted relative spectral response
    functions from the internet as tarballs, extracts them, then deletes
    the tarball.

    See :func:`pyspectral.rsr_reader.check_and_download` for a "smart" version
    of this process that only downloads the necessary files.

    Args:
        dest_dir: Path to put the temporary tarball and extracted RSR files.
        dry_run: If True, don't actually download files, only log what
            URLs would be downloaded. Defaults to False.

    """
    config = get_config()
    local_rsr_dir = Path(config.get("rsr_dir"))  # type: ignore
    dest_path = Path(dest_dir) if dest_dir is not None else local_rsr_dir

    LOG.info(f"Download RSR files and store in directory {dest_path}")
    filename = dest_path / "pyspectral_rsr_data.tgz"
    LOG.debug(f"RSR URL: {HTTP_PYSPECTRAL_RSR}")
    LOG.debug(f"Destination = {dest_path}")
    if dry_run:
        return

    _download_tarball_and_extract(HTTP_PYSPECTRAL_RSR, filename, dest_path)


def download_luts(aerosol_types=None, dry_run=False, aerosol_type=None):
    """Download the luts from internet.

    See :func:`pyspectral.rayleigh.check_and_download` for a "smart" version
    of this process that only downloads the necessary files.

    Args:
        aerosol_types (Iterable): Aerosol types to download the LUTs for.
            Defaults to all aerosol types. See :data:`AEROSOL_TYPES` for the
            full list.
        dry_run (bool): If True, don't actually download files, only log what
            URLs would be downloaded. Defaults to False.
        aerosol_type (str): Deprecated.

    """
    aerosol_types = _get_aerosol_types(aerosol_types, aerosol_type)
    for subname in aerosol_types:
        LOG.debug("Aerosol type: %s", subname)
        lut_tarball_url = HTTPS_RAYLEIGH_LUTS[subname]
        LOG.debug("Atmospheric LUT URL = %s", lut_tarball_url)

        subdir_path = get_rayleigh_lut_dir(subname)
        LOG.debug(f"Create directory: {subdir_path}")
        if not dry_run:
            subdir_path.mkdir(exist_ok=True, parents=True)
        if dry_run:
            continue

        local_tarball_pathname = subdir_path / "pyspectral_rayleigh_correction_luts.tgz"
        _download_tarball_and_extract(lut_tarball_url, local_tarball_pathname, subdir_path)


def _get_aerosol_types(aerosol_types, aerosol_type):
    if aerosol_type is not None:
        warnings.warn("'aerosol_type' is deprecated, use 'aerosol_types' instead.", UserWarning,
                      stacklevel=3)
        if isinstance(aerosol_type, (list, tuple, set)):
            aerosol_types = aerosol_type
        else:
            aerosol_types = [aerosol_type]
    elif aerosol_types is None:
        aerosol_types = list(HTTPS_RAYLEIGH_LUTS.keys())
    return aerosol_types


HEADERS = {
    "User-Agent": "pyspectral (+https://github.com/pytroll/pyspectral)",
    "Accept": "*/*",
}


def _download_tarball_and_extract(tarball_url: str, local_pathname: Path, extract_dir: Path) -> None:
    chunk_size = 1024 * 1024  # 1 MB

    response = requests.get(
        tarball_url,
        headers=HEADERS,
        stream=True,
        allow_redirects=True,
    )
    response.raise_for_status()

    total_size = int(response.headers.get("content-length", 0))

    with open(local_pathname, "wb") as handle:
        for data in _tqdm_or_iter(response.iter_content(chunk_size=chunk_size),
                                  total=(int(total_size / chunk_size + 0.5)),
                                  unit="kB"):
            if data:
                handle.write(data)

    tar_kwargs = {} if sys.version_info < (3, 12) else {"filter": "data"}
    with tarfile.open(local_pathname) as tar:
        tar.extractall(extract_dir, **tar_kwargs)  # type: ignore

    os.remove(local_pathname)


def _tqdm_or_iter(an_iterable, **tqdm_kwargs):
    """Wrap an iterable with tqdm if it is available, otherwise return the iterable."""
    if TQDM_LOADED:
        return tqdm(iterable=an_iterable, **tqdm_kwargs)
    else:
        return an_iterable


def debug_on():
    """Turn debugging logging on."""
    logging_on(logging.DEBUG)


_is_logging_on = False


def logging_on(level=logging.WARNING):
    """Turn logging on."""
    global _is_logging_on

    if not _is_logging_on:
        console = logging.StreamHandler()
        console.setFormatter(logging.Formatter("[%(levelname)s: %(asctime)s :"
                                               " %(name)s] %(message)s",
                                               "%Y-%m-%d %H:%M:%S"))
        console.setLevel(level)
        logging.getLogger("").addHandler(console)
        _is_logging_on = True

    log = logging.getLogger("")
    log.setLevel(level)
    for h in log.handlers:
        h.setLevel(level)


class NullHandler(logging.Handler):
    """Empty handler."""

    def emit(self, record):
        """Record a message."""


def logging_off():
    """Turn logging off."""
    logging.getLogger("").handlers = [NullHandler()]


def get_logger(name):
    """Return logger with null handle."""
    log = logging.getLogger(name)
    if not log.handlers:
        log.addHandler(NullHandler())
    return log


def get_fullwidth_halfmax(rsp, wvl):
    """Compute the full-width half maximum (bandwidth) for a given spectral response.

    Ths FWHM is defined as the wavelength range over which the SRF is greater than half
    of the maximum response. Often, this is referred to as the bandwidth of a channel.
    """
    half_max = 0.5
    indices = np.where(rsp >= half_max)[0]
    if len(indices) < 2:
        raise ValueError("The FWHM cannot be computed.")
    else:
        fwhm = wvl[indices[-1]] - wvl[indices[0]]
    return fwhm


def get_bounds_integrated_energy(rsp, wvl, ener_perc_lim=1):
    """Compute wavelengths over which the integrated energy is greater than a threshold.

    This measure is typically used to assess out of band performance and to ensure that
    the spectral response function is not sensitive to specific absorption or emission
    features. On occasion, it is also used to define or measure the bandwidth of a sensor.

    The ``ener_perc_lim`` variable defines the threshold percentage of integrated energy.
    This function returns the minimum and maximum wavelengths over which the integrated
    energy is greater than (minimum) and less than (maximum) the threshold.
    """
    crs = np.cumsum(rsp)
    crs = crs / np.max(crs) * 100.
    mini_ind = np.where(crs >= ener_perc_lim)[0]
    maxi_ind = np.where(crs <= (100 - ener_perc_lim))[0]

    return wvl[mini_ind][0], wvl[maxi_ind][-1]


def get_wave_range(in_chan, threshold=0.15):
    """Return central, min and max wavelength in an RSR greater than threshold.

    An RSR function will generally start near zero, increase to a maximum and
    then drop back to near zero. This function takes advantage of this to find
    the first and last points where the RSR is greater than a threshold. These
    points are then defined as the minimum and maximum wavelengths for a
    given channel, and can be used, for example, in Satpy reader YAML files.

    """
    cwl = get_central_wave(in_chan["wavelength"], in_chan["response"])

    wvls = in_chan["wavelength"]
    rsr = in_chan["response"]

    pts = (rsr > threshold).nonzero()
    min_wvl = wvls[pts[0][0]]
    max_wvl = wvls[pts[0][-1]]

    return [min_wvl, cwl, max_wvl]


def convert2str(value):
    """Convert a value to string.

    Args:
        value: Either a str, bytes or 1-element numpy array

    """
    value = bytes2string(value)
    return np2str(value)


def np2str(value):
    """Convert an ``numpy.bytes_`` to str.

    Note: ``numpy.string_`` was deprecated in numpy 2.0 in favor of
    ``numpy.bytes_``.

    Args:
        value (ndarray): scalar or 1-element numpy array to convert
    Raises:
        ValueError: if value is array larger than 1-element or it is not of
                    type `numpy.bytes_` or it is not a numpy array

    """
    if isinstance(value, str):
        return value

    if hasattr(value, "dtype") and \
            issubclass(value.dtype.type, (np.str_, np.bytes_, np.object_)) \
            and value.size == 1:
        value = value.item()
        # python 3 - was scalar numpy array of bytes
        # otherwise python 2 - scalar numpy array of "str"
        if not isinstance(value, str):
            return value.decode()
        return value

    raise ValueError("Array is not a string type or is larger than 1")


def bytes2string(var):
    """Decode a bytes variable and return a string."""
    if isinstance(var, bytes):
        return var.decode("utf-8")
    return var


def _replace_inst_name(instruments):
    """Replace an instrument name if it's not available on a platform."""
    if isinstance(instruments, list):
        return instruments[0]
    else:
        return instruments


def check_and_adjust_instrument_name(platform_name, instrument):
    """Check instrument name and try fix if inconsistent.

    It checks against the possible listed instrument names for each platform.
    It also makes an adjustment replacing names like avhrr/1 with avhrr1,
    removing the "/".
    """
    instr = INSTRUMENTS.get(platform_name, instrument.lower())
    if isinstance(instr, list):
        for inst in instr:
            goodinst = are_instruments_identical(inst, instrument.lower())
            if goodinst:
                break
    else:
        goodinst = are_instruments_identical(instr, instrument.lower())

    if not goodinst:
        instrument = _replace_inst_name(instr)
        LOG.warning("Inconsistent instrument/satellite input - instrument set to %s",
                    instrument)

    return instrument.lower().replace("/", "").replace("-", "")


def are_instruments_identical(name1, name2):
    """Given two instrument names check if they are both describing the same instrument.

    Takes care of the case of AVHRR where the internal pyspectral naming
    (following WMO Oscar) is is with a slash as in "avhrr/1", but where a
    naming using a dash instead is equally accepted, as in "avhrr-1".
    """
    if not isinstance(name1, str) or not isinstance(name2, str):
        raise ValueError("Instrument names must be strings.")

    if name1 == name2:
        return True
    if name1 == INSTRUMENT_TRANSLATION_DASH2SLASH.get(name2):
        return True
    return False


def use_map_blocks_on(argument_to_run_map_blocks_on):
    """Use map blocks on a given argument.

    This decorator assumes only one of the arguments of the decorated function is chunked.
    """
    def decorator(f):
        argspec = getfullargspec(f)
        argument_index = argspec.args.index(argument_to_run_map_blocks_on)

        @wraps(f)
        def wrapper(*args, **kwargs):
            array = args[argument_index]
            chunks = getattr(array, "chunks", None)
            if chunks is None:
                return f(*args, **kwargs)
            import dask.array as da
            import xarray as xr
            if isinstance(array, da.Array):
                return da.map_blocks(f, *args, **kwargs)
            elif isinstance(array, xr.DataArray):
                new_array = array.copy()
                new_args = list(args)
                new_args[argument_index] = array.data
                new_data = da.map_blocks(f, *new_args, **kwargs)
                new_array.data = new_data
                return new_array
            else:
                raise NotImplementedError(f"Don't know how to map_blocks on {type(array)}")
        return wrapper
    return decorator
