"""
File Purpose: misc. tools related to xarray
"""
import datetime as datetime
import inspect
from collections.abc import Callable

import numpy as np
import xarray as xr

from .xarray_accessors import pcAccessor
from ..arrays import np_dtype_object_to_str
from ..docs_tools import format_docstring
from ..os_tools import next_unique_name
from ..sentinels import UNSET
from ..timing import ProgressUpdater
from ...defaults import DEFAULTS
from ...errors import (
    InputError, InputMissingError, InputConflictError,
    DimensionKeyError,
)


### --------------------- Copying --------------------- ###

@pcAccessor.register('copy_kw')
def xarray_copy_kw(array, dims=None, *, name=None, array_to_dataset=False):
    '''return dict of info suitable for creating a similar array or dataset.
    result includes dims, coords, and attrs (unchanged, copied).

    dims: None, str, or iterable of str
        if provided, only include these dims (and related coords) in the result.
        Useful if only interested in some of the dims,
            e.g. if array has x,y,z,t dims but only want to mimic dims and coords from x,t.
    name: None, bool, or str
        whether to also include name in result
        None --> True if array has name, else False
        True --> name = array.name.
        str --> name = name.
    array_to_dataset: bool
        if True, result will be suitable for creating a Dataset based on the input DataArray.
        (Equivalent to using dims=[], name=False, and deleting 'dims' key from result.)
    '''
    # handle array_to_dataset mode:
    if array_to_dataset:
        result = xarray_copy_kw(array, dims=[], name=False, array_to_dataset=False)
        del result['dims']
        return result
    # 'usual' usage
    result = dict()
    if isinstance(dims, str): dims = [dims]
    if dims is None: dims = array.dims
    elif any(d not in array.dims for d in dims):
        errmsg = (f'cannot copy_kw for dims={dims!r} because some dims not in array.dims={array.dims}.')
        raise DimensionKeyError(errmsg)
    result['dims'] = dims
    coords = dict()
    for cname, cval in array.coords.items():
        if len(cval.dims) == 0:  # scalar coord; always keep!
            coords[cname] = cval
        elif len(cval.dims) == 1 and cval.dims[0] in dims:  # 1D coord, and relevant
            coords[cname] = cval
        elif all(d in dims for d in cval.dims):  # 2D+ coord, and all relevant dims exist
            coords[cname] = (tuple(cval.dims), cval.data)
    result['coords'] = coords
    attrs = array.attrs.copy()
    result['attrs'] = attrs
    if name is None:
        name = hasattr(array, 'name')
    if name == True:
        name = array.name
    if name != False:
        result['name'] = name
    return result

@pcAccessor.register('with_data', totype='array')
def xarray_with_data(xarray, data, dims=None, *, name=None):
    '''return data as an xr.DataArray with dims, coords, name, and attrs copied from xarray.
    Similar to DataArray.copy(data=data), but here allows a bit more control.

    dims: None, str, or iterable of str
        if provided, only include these dims (and related coords) in the result.
        Useful if only interested in some of the dims,
            e.g. if array has x,y,z,t dims but only want to mimic dims and coords from x,t.
    name: None, bool, or str
        whether to also include name in result
        None --> True if array has name, else False
        True --> name = array.name.
        str --> name = name.
    '''
    if not isinstance(xarray, xr.DataArray):
        raise InputError(f'xarray_with_data expects args[0] to be a DataArray, not {type(xarray)}.')
    kw = xarray_copy_kw(xarray, dims=dims, name=name)
    return xr.DataArray(data, **kw)


### --------------------- Converting Dataset to Array --------------------- ###

@pcAccessor.register('as_array')
def xarray_as_array(array):
    '''return array if DataArray, else array[var] if array is Dataset with one var.
    if array is Dataset with multiple vars, crash with InputError.
    '''
    if isinstance(array, xr.DataArray):
        return array
    elif isinstance(array, xr.Dataset):
        if len(array) == 1:
            return array[list(array)[0]]
        else:
            errmsg = ('as_array(Dataset) only possible if Dataset has exactly 1 data_var, '
                     f'but got Dataset with {len(array)} vars: {list(array.data_vars)}.')
            raise InputError(errmsg)
    else:
        raise InputError(f'cannot convert object of type {type(array)} to DataArray.')


### --------------------- Merging --------------------- ###

@format_docstring(print_freq=DEFAULTS.PROGRESS_UPDATES_PRINT_FREQ,
                  ngroup=DEFAULTS.XARRAY_MERGE_NGROUP)
def xarray_merge(arrays, *, ngroup=UNSET, verbose=True, missing_name_ok=True, to_ds=None,
                 _updater=None, **kw_xr_merge):
    '''xr.merge() these xarray objects (probably DataArrays or Datasets),
    but maybe slightly more efficiently than just using xr.merge() directly.
    Also, has slightly more convenient default behaviors:
        - by default, allow name=None on inputs (treat all inputs with None name as same name)
        - by default, if all inputs are DataArray with same name, output DataArray not Dataset.

    ngroup: int>=2, None, or UNSET
        [EFF] max number of arrays per xr.merge() call.
        (Final results should be unaffected by ngroup, just affects internal strategy.)
        UNSET --> use DEFAULTS.XARRAY_MERGE_NGROUP (default: {ngroup}).
        None --> merge all arrays in a single xr.merge() call.
        int>=2 --> merge arrays in groups of size ngroup, then merge the merged groups.
            (This can be much faster than ngroup=None; see xarray_merge for details.)
    verbose: bool or int
        whether to print progress updates. (Ignored if N <= ngroup.)
        0: never print progress updates
        1: print updates every DEFAULTS.PROGRESS_UPDATES_PRINT_FREQ seconds (default: {print_freq})
        2: print updates every file.
    missing_name_ok: bool
        whether to allow name=None as a valid name for input arrays.
        True --> treat all inputs with name=None as having the same name (None).
            (if result is a Dataset this can lead to a key in result being None.)
    to_ds: None
        whether result should be a Dataset.
        None --> DataArray if and only if all inputs are DataArrays with same name.
        True --> always return Dataset.
        False --> always return a DataArray (crash if any inputs are not DataArrays with same name).
    _updater: None or ProgressUpdater
        if provided, use this (instead of a new ProgressUpdater) to print updates, ignoring `verbose`.

    --- Timing notes: ---
        Tt seems like merging more than ~16 arrays at once can cause significant slowdowns,
        when compared to merging groups of 16, then merging those results.
        Tests with xarray 2024.7.0, Python 3.9.13, with subsets of 256 arrays totaling roughly 2 GB:
            64 arrays with N=None --> ~42 seconds | N=2 --> ~16s | N=4, 8, 16 --> ~12s
            128 arrays with N=None --> (didn't test it) | N=16 --> ~30 seconds
            256 arrays with N=None --> crashed due to memory errors | N=16 --> ~102 seconds
    '''
    # [TODO] is this basically just becoming xarray.combine_by_coords()?
    #   should that method be used instead? (Maybe it would be faster, too?)

    # check for any "obvious" crashes early:
    if len(arrays) == 0:
        raise InputMissingError('xarray_merge expects at least 1 DataArray or Dataset.')
    _to_ds_input = to_ds  # helps with debugging
    if not to_ds:
        _all_arrays = all(isinstance(arr, xr.DataArray) for arr in arrays)
        names = [arr.name for arr in arrays]
        _all_samename = len(set(names)) == 1
        if to_ds is None:
            to_ds = (not (_all_arrays and _all_samename))
        else:  # to_ds == False
            if not _all_arrays:
                raise InputConflictError('to_ds=False only possible if all inputs are DataArrays.')
            if not _all_samename:
                errmsg = f'to_ds=False requires all inputs to have the same name. Got: {names!r}.'
                raise InputConflictError(errmsg)
    if not missing_name_ok:
        if any(isinstance(arr, xr.DataArray) and arr.name is None for arr in arrays):
            raise InputConflictError('missing_name_ok=False but at least 1 input has name=None.')
    # handle simple case now (only 1 input)
    if len(arrays) == 1:
        result = arrays[0]
        if to_ds and isinstance(result, xr.DataArray):
            return xr.Dataset({result.name: result})
        else:
            return result
    # more bookkeeping
    names = set(arr.name for arr in arrays if isinstance(arr, xr.DataArray))
    rename_Nones = False
    if None in names:
        arrays = [*arrays]  # <-- copy input list to ensure input isn't edited in-place
        rename_Nones = next_unique_name('__xarray_merge_None_placeholder__', names)
        for i, arr in enumerate(arrays):
            if isinstance(arr, xr.DataArray) and arr.name is None:
                arrays[i] = arr.rename(rename_Nones)
    if ngroup is UNSET:
        ngroup = DEFAULTS.XARRAY_MERGE_NGROUP
    # >>> actual merging <<<
    if ngroup is None or len(arrays) <= ngroup:
        result = xr.merge(arrays, **kw_xr_merge)
    else:
        # else, need to do grouping
        print_freq = None if verbose==1 else (0 if verbose>=2 else -1)
        updater = ProgressUpdater(print_freq=print_freq) if _updater is None else _updater
        _orig_arrays = arrays  # just to help with debugging / printouts.
        layer = 0
        while len(arrays) > ngroup:
            merged = []
            for i0 in range(0, len(arrays), ngroup):
                i1 = min(i0 + ngroup, len(arrays))
                updater.print(f'Merging i={i0} to i={i1-1}, of max(i)={len(arrays)-1}' +
                              '.' if layer==0 else f' (-- merging merged arrays; at layer {layer}).')
                merged_group = xr.merge(arrays[i0:i1], **kw_xr_merge)
                merged.append(merged_group)
            arrays = merged
            layer = layer + 1
        # final merge:
        updater.print(f'Final merge of ({len(arrays)}) merged arrays.')
        result = xr.merge(arrays, **kw_xr_merge)
        updater.finalize(f'Merging of {len(_orig_arrays)} arrays complete.')
    # bookkeeping
    if rename_Nones:
        result = result.rename({rename_Nones: None})
    if not to_ds:   # (result of xr.merge() is always a Dataset. Convert back to DataArray.)
        key0 = next(iter(result.keys()))
        result = result[key0]
    return result

@pcAccessor.register('merge')
@format_docstring(xarray_merge_doc=xarray_merge.__doc__, sub_ntab=1)
def xarray_merge_from(array, other_arrays, *, ngroup=UNSET, verbose=True, **kw_xarray_merge):
    '''xr.merge() this xarray object (DataArray or Dataset) with other xarray objects.
    Internally, uses xarray_merge which might be a bit more efficient than xr.merge().

    xarray_merge docs copied below for convenience:
    -----------------------------------------------
        {xarray_merge_doc}
    '''
    all_arrays = [array, *other_arrays]
    return xarray_merge(all_arrays, ngroup=ngroup, **kw_xr_merge)


### --------------------- Vars from dataset --------------------- ###

@pcAccessor.register('vars_lookup', totype='dataset')
def xarray_vars_lookup(ds, lookup={}, **lookup_as_kw):
    '''return dict of {key from lookup: ds[lookup[key]]}
    result vals will be None for keys not found in ds.

    lookup: dict of {key: None, varname, or list of varnames}
        look for these varnames in ds. if list, use first varname found.
        if None, use key as varname.
    '''
    if lookup and lookup_as_kw:
        raise InputConflictError('provide lookup or lookup_as_kw, not both!')
    if lookup_as_kw:
        lookup = lookup_as_kw
    result = {}
    for key, find in lookup.items():
        result[key] = None
        if find is None:
            find = key
        if isinstance(find, str):
            find = [find]
        for varname in find:
            if varname in ds:
                result[key] = ds[varname]
                break
    return result

@pcAccessor.register('vars_lookup_with_defaults', totype='dataset')
def xarray_vars_lookup_with_defaults(ds, provided, lookup={}, defaults={}):
    '''return dict of {key from lookup: value for that key}.
    
    provided: dict of {key: None, str, or other value}
        provided values or lookup instructions for each key.
        None --> use instructions from lookup if key in lookup,
                 else value from defaults if key in defaults.
        str --> varname; result uses ds[varname]. (crash if not possible.)
        callable --> f of one arg; result uses f(ds).
        other value --> result uses this value directly.
    lookup: dict of {key: None, str, or list of str}
        'default' lookup instructions. Used when provided[key] is None.
    defaults: dict of {key: value or callable}
        default values for each key.
        Used when provided[key] is None and ds[lookup[key]] not found.
        callable --> called with ds as arg to get value.

    Examples:
        xarray_vars_lookup_with_defaults(ds, {}, {'n': None}, {'n': 100})
        # {'n': ds['n']} if ds['n'] exists, else {'n': 100}
        
        xarray_vars_lookup_with_defaults(ds, {}, {'n': 'var1'}, {'n': lambda ds: ds.size})
        # {'n': ds['var1']} if ds['var1'] exists, else {'n': ds.size}

        xarray_vars_lookup_with_defaults(ds, {'n': 'var1'}, {'n': 'var2'}, {'n': 7})
        # {'n': ds['var1']} if ds['var1'] exists, else crash.
        # (lookup & defaults are fully ignored, if non-None value in provided.)
    '''
    result = {}
    all_keys = set(provided) | set(lookup) | set(defaults)
    for key in all_keys:
        prov = provided.get(key, None)
        if prov is None:
            find = lookup.get(key, None)
            if find is None:
                find = key
            if isinstance(find, str):
                find = [find]
            for varname in find:
                if varname in ds:
                    result[key] = ds[varname]
                    break
            else:
                if key in defaults:
                    val = defaults[key]
                    if callable(val):
                        val = val(ds)
                    result[key] = val
                else:
                    errmsg = (f'no value for key={key!r}; not provided directly, '
                              f'varnames {find} not in ds, and no default provided.')
                    raise InputError(errmsg)
        elif isinstance(prov, str):
            if prov in ds:
                result[key] = ds[prov]
            else:
                errmsg = f'provided varname={prov!r} (for key={key!r}) not found in ds.'
                raise InputError(errmsg)
        elif callable(prov):
            result[key] = prov(ds)
        else:
            result[key] = prov
    return result


### --------------------- Math Checks --------------------- ###

@pcAccessor.register('where_finite')
def xarray_where_finite(array):
    '''returns array, masked with NaNs anywhere that the values are not finite.'''
    return array.where(np.isfinite)


### --------------------- Type Casting --------------------- ###

@pcAccessor.register('astypes', totype='dataset')
def xarray_astypes(ds, var_to_dtype):
    '''return copy of ds with var.astype(dtype) for each var, dtype in var_to_dtype.items()'''
    to_assign = {dvar: ds[dvar].astype(dtype) for dvar, dtype in var_to_dtype.items()}
    return ds.assign(**to_assign)

@pcAccessor.register('convert_types')
def xarray_convert_types(array, oldtype_to_newtype):
    '''return copy of array or dataset, using var.astype(newtype) for any var with oldtype,
    for oldtype, newtype in oldtype_to_newtype.items().
    '''
    if isinstance(array, xr.Dataset):
        ds = array
        to_convert = dict()
        for var, val in ds.items():
            for oldtype, newtype in oldtype_to_newtype.items():
                if val.dtype == oldtype:
                    to_convert[var] = newtype
                    break
        return xarray_astypes(ds, to_convert)
    else:
        for oldtype, newtype in oldtype_to_newtype.items():
            if array.dtype == oldtype:
                return array.astype(newtype)
        else:  # no match found during loop
            return array

@pcAccessor.register('dtype_object_to_str', totype='array')
def xarray_dtype_object_to_str(array):
    '''convert array with dtype=object to dtype=str, by doing str(x) for every x in array.'''
    return array.copy(data=np_dtype_object_to_str(array.values))


### --------------------- Converting coords to strings --------------------- ###

@pcAccessor.register('object_coords_to_str')
def xarray_object_coords_to_str(array, *, maxlen=None, ndim_min=1):
    '''return copy of array with coords (of dtype=object) converted to string.
    maxlen: None or int>=5
        if int, truncate longer strings to this length-3, and add '...' to end.
    ndim_min: int
        minimum number of dimensions for a coord to be converted.
        e.g. ndim_min=1 --> coords with ndim=0 will not be altered
    '''
    to_assign = {}
    for cname, cc in array.coords.items():
        if cc.dtype == object and cc.ndim >= ndim_min:
            if cc.ndim >= 2:
                raise NotImplementedError('object_coords_to_str does not yet support ndim>=2.')
            cvals = [str(v) for v in cc.values]
            if maxlen is not None:
                cvals = [v[:maxlen-3] + '...' if len(v) > maxlen else v for v in cvals]
            to_assign[cname] = cc.copy(data=cvals)
    return array.assign_coords(to_assign)



### --------------------- Recording "HISTORY" attr --------------------- ###

@pcAccessor.register('update_call_history', totype='dataset')
def xarray_update_call_history(ds: xr.Dataset, local_vars: dict, func: Callable) -> None:
    """
    Add a history entry to a dataset. Edits the input dataset in-place, then returns it.

    ds.attrs["HISTORY"] stores entries which tell a history of functions and their inputs,
        which were called and recorded in the dataset via update_call_history().
    The entries are strings, intended primarily for human-readability.

    Also adds ds.attrs["date created"] if it does not exist,
        else adds or updates ds.attrs["date modified"] to today's date.

    Parameters
    ----------
    ds : `xarray.Dataset`
        Dataset to update.
    local_vars : `dict`
        Local variables from the calling function.
    func : `Callable`
        Function or method being recorded in the history.
    """
    string_vals = []
    for arg, value in local_vars.items():
        if arg in inspect.signature(func).parameters:
            if isinstance(value, (xr.Dataset, xr.DataArray, np.ndarray)):
                if isinstance(value, np.ndarray) and value.shape in [(), (1,)]:
                    string_vals.append(f"{arg}={value.tolist()}")
                elif isinstance(value, xr.DataArray) and value.size == 1:
                    string_vals.append(f"{arg}={value.values.tolist()}")
                else:
                    string_vals.append(f"{arg}={arg}")
            else:
                string_vals.append(f"{arg}={value}")

    history_entry = f"{func.__name__}({', '.join(string_vals)})"
    if "HISTORY" in ds.attrs:
        if isinstance(ds.attrs["HISTORY"], list):
            ds.attrs["HISTORY"].append(history_entry)
        else:
            ds.attrs["HISTORY"] = [ds.attrs["HISTORY"], history_entry]
    else:
        ds.attrs["HISTORY"] = [history_entry]

    today = datetime.datetime.now(tz=datetime.timezone.utc)
    if "date created" in ds.attrs:
        ds.attrs["date modified"] = today.strftime("%d-%b-%Y")
    else:
        ds.attrs["date created"] = today.strftime("%d-%b-%Y")
    return ds
