"""
File Purpose: misc. tools related to checking xarray size,
and doing manipulations more efficiently with chunking for very large arrays.
"""
import numpy as np
import xarray as xr

from .xarray_accessors import pcAccessor
from .xarray_agg_stats import xarray_sum
from ..docs_tools import format_docstring
from ..math import product
from ..oop_tools import attach_attrs
from ..sentinels import UNSET
from ...defaults import DEFAULTS
from ...errors import (
    DimensionValueError,
    InputError,
    MemorySizeError,
)


### --------------------- Predict Size --------------------- ###

@pcAccessor.register('max_dim_sizes', totype='array')
def xarray_max_dim_sizes(array, *more_arrays, assert_broadcastable=False):
    '''return dict of {dim: max size across all arrays} for all dims in any of the arrays.

    assert_broadcastable: bool
        whether to raise DimensionValueError if any dim has different sizes in different arrays
        (other than size=1, which is considered broadcastable to any other size.)
        [TODO] more lenient definition which checks shared coordinates instead.
    '''
    arrays = (array,) + more_arrays
    result = {}
    for arr in arrays:
        for d, s in arr.sizes.items():
            if d in result:
                if result[d] != s:
                    if s == 1 or result[d] == 1:
                        result[d] = max(result[d], s)
                    elif assert_broadcastable:
                        errmsg = (f'cannot get max_dim_sizes due to different sizes for dim {d!r}: '
                                  f'{result[d]} vs {s}.')
                        raise DimensionValueError(errmsg)
                    else:
                        result[d] = max(result[d], s)
            else:
                result[d] = s
    return result

@pcAccessor.register('predict_result_size', totype='array')
def xarray_predict_result_size(array, *more_arrays, units='items'):
    '''return prediction of how large the product of these xarray objects would be.
    (or, any other math operation which broadcasts operations in the standard way..)

    units: str
        tells the units of the result (default: 'items'). Options:
        'items' --> number of items (elements) in the joined array.
        'bytes', 'kB', 'MB', 'GB' --> memory size of the data in joined array.
            (based on each array's dtype.itemsize and number of items.)
    '''
    arrays = (array,) + more_arrays
    VALID_UNITS = ('items', 'bytes', 'kB', 'MB', 'GB')
    if units not in VALID_UNITS:
        raise InputError(f'invalid units={units!r}. Expected one of {VALID_UNITS}.')
    result_dim_sizes = xarray_max_dim_sizes(*arrays, assert_broadcastable=True)
    n_items = product(result_dim_sizes.values())
    if units == 'items':
        return n_items
    # else, need to convert to bytes
    result_dtype = np.result_type(*arrays)
    n_bytes = n_items * result_dtype.itemsize
    bytes_per_unit = {'bytes': 1, 'kB': 1024, 'MB': 1024**2, 'GB': 1024**3}
    return n_bytes / bytes_per_unit[units]

@pcAccessor.register('result_size_check', totype='array')
def xarray_result_size_check(array, *more_arrays, GBmax=UNSET, safety=1, errmsg=None):
    '''ensure that the result of combining these xarrays (e.g., via product) is not too large.
    Too large if safety * predicted_size_GB > GBmax

    GBmax: UNSET, None, or number
        maximum size allowed, in Gigabytes.
        UNSET --> use DEFAULTS.RESULT_ARRAY_GBYTES_MAX
        None --> no limit (equivalent to "skip this check")
    safety: number
        safety factor; safety * predicted_size_GB is the actual limit.
        smaller is "safer" (effectively reduces limit)
    errmsg: None or str
        error message to use if the size is too large.
        will be formatted with .format(nGB=nGB, GBmax=GBmax, safety=safety,
                         dimsizes=(dict of predicted sizes of result dims)).

    raises MemorySizeError if array is too large.
    '''
    GBmax_is_default = (GBmax is UNSET)
    if GBmax is UNSET:
        GBmax = DEFAULTS.RESULT_ARRAY_GBYTES_MAX
    if GBmax is not None:
        nGB = xarray_predict_result_size(array, *more_arrays, units='GB')
        if safety * nGB > GBmax:
            dimsizes = xarray_max_dim_sizes(array, *more_arrays)
            if errmsg is None:
                errmsg = 'predicted size of result is too large: {nGB:.2f} GB > GBmax (={GBmax})'
                if safety != 1:
                    errmsg += ' / safety (={safety})'
                errmsg += '\nPredicted result dimsizes: {dimsizes}'
                if GBmax_is_default:
                    errmsg += '\nTo change limit, adjust DEFAULTS.RESULT_ARRAY_GBYTES_MAX.'
                errmsg += '\nTo skip this check, set the limit to None instead of a number.'
            errmsg = errmsg.format(nGB=nGB, GBmax=GBmax, safety=safety, dimsizes=dimsizes)
            raise MemorySizeError(errmsg)


### --------------------- Multiply Then Sum Along --------------------- ###

_paramdocs_xarray_multiply_then_sum_along = {
    'loop': '''None, str, or (str, int, mode).
        Indicates looping instructions (if any) to use during xarray_multiply_then_sum_along.
        None --> no looping; simply multiply all arrays, then sum along dim.
            Conceptually simplest. But, internally builds unnecessarily large array,
            which might be larger than DEFAULTS.RESULT_ARRAY_GBYTES_MAX and thus trigger MemorySizeError.
        str --> loop over each point along the indicated dimension.
            If loop != dim, sum at each point and then concatenate final result.
            If loop == dim, keep a running sum during the loop; no need to concatenate final result.
        'int' or 'sum' --> equivalent to using loop == dim.
        (str, integer, mode) --> (loop dimension, chunk size, chunk size mode)
            loop over the indicated dimension in chunks of the indicated size,
            instead of only one point at a time. This is usually the most efficient option. Mode can be:
                'GB' --> size of chunk arrays in GB (approximately)
                    E.g., ('logT', 5, 'GB') picks chunk size for logT such that each chunk is 5 GB or less.
                'size' --> number of points in the dimension.
                    E.g., ('x', 5, 'size') does first 5 points along x, then next 5, etc.
                'n' --> number of chunks (roughly)
                    E.g., ('x', 5, 'n') does first 1/5 of n points, then next 1/5, etc.
                    Not guaranteed to be exactly 5 chunks, due to rounding.''',
    'errmsg': '''None or str
        error message to use if the predicted size of any intermediate array is too big.
        Will be hit by .format(nGB=nGB, GBmax=GBmax, dimsizes=(dict of predicted sizes of result dims)).
        None --> make reasonably informative default error message.''',
}

_xarray_multiply_then_sum_along_errmsg = \
    '''In xarray_multiply_then_sum_along, multiplying arrays (before summing) would be too large!
    The limit is set by DEFAULTS.RESULT_ARRAY_GBYTES_MAX = {GBmax} GB.
    Predicted size: {nGB:.2f} GB, corresponding to dims with sizes (number of items):
        {dimsizes}.

    See help(xarray_multiply_then_sum_along) or help(array.pc.multiply_then_sum_along) for details.

    Options to fix this problem include any combination of the following:
        - work with smaller arrays. E.g., use slices or chunks appropriately.
        - adjust `loop` kwarg to loop over an internal dimension.
            See help(xarray_multiply_then_sum_along) for details.
        - (simple but not reommended) increase DEFAULTS.RESULT_ARRAY_GBYTES_MAX,
            or, set it to None to disable this check entirely.
            (Warning: might crash machine if internal array is too big!!)'''

@pcAccessor.register('multiply_then_sum_along', totype='array')
@format_docstring(**_paramdocs_xarray_multiply_then_sum_along)
@attach_attrs(_paramdocs_dict=_paramdocs_xarray_multiply_then_sum_along,
              _errmsg_default=_xarray_multiply_then_sum_along_errmsg)
def xarray_multiply_then_sum_along(array, *more_arrays, dim, loop=('int', 0.45, 'GB'), errmsg=None):
    '''returns result of product of arrays then .sum(dim=dim).
    Numerically equivalent to (arr0 * arr1 * ...).sum(dim=dim).

    But, internally does extra checks, and can be more efficient for large arrays:
        - ensures intermediate steps and final result never lead to "way too large" array,
            i.e. never bigger than DEFAULTS.RESULT_ARRAY_GBYTES_MAX.
            If it would be too big, crash with helpful MemorySizeError, instead.
        - may loop over a dimension, depending on `loop`, to avoid creating large intermediate arrays.

    array, more_arrays: xarray.DataArray objects
        the arrays to be multiplied
    dim: str
        the dimension to sum along after multiplying.
        (ignored if not present in any of the arrays.)
    loop: {loop}
    errmsg: {errmsg}
    '''
    # setup
    if errmsg is None:
        errmsg = xarray_multiply_then_sum_along._errmsg_default
    arrays = (array, *more_arrays)
    if loop is None:  # simplest, but also most likely to be too big
        xarray_result_size_check(*arrays, errmsg=errmsg)
        result = xarray_sum(product(arrays), dim=dim, missing_dims='ignore')
        return result
    else:
        ## bookkeeping ##
        if isinstance(loop, str):
            loop = (loop, 1, 'size')
        loopdimstr, loopsize, mode = loop
        if loopdimstr in ('int', 'sum'):
            x = dim   # x == loop dim.
        else:
            x = loopdimstr
        _valid_modes = ('size', 'n', 'GB')
        if mode not in _valid_modes:
            raise InputError(f"invalid loop mode: {mode!r}; expected one of {_valid_modes}")
        # size
        dim_sizes = xarray_max_dim_sizes(*arrays)
        if mode == 'size':
            size = loopsize    # size == size (number of points)
        elif mode == 'n':
            size = dim_sizes[x] // loopsize  # e.g. dimsize=110, n=25 --> size=4.
        elif mode == 'GB':
            full_GB = xarray_predict_result_size(*arrays, units='GB')
            # e.g. loopsize = 0.5 GB, full_GB = 2.3 --> size = int(dimsize / 5)
            size = int(dim_sizes[x] * min(1, loopsize / full_GB))  # min() in case loopsize > full_GB.
            if DEFAULTS.DEBUG >= 2:
                print(f'xarray_multiply_then_sum_along chose size={size} (of full={dim_sizes[x]}) '
                      f'for loopdim={x!r}, to make {full_GB * size/dim_sizes[x]:.3f} GB per chunk.')
        if size < 1:
            size = 1
        # slicers (based on size)
        if size == dim_sizes[x]:
            slicers = None  # no need to loop after all!
        elif size == 1:
            slicers = list(range(dim_sizes[x]))  # loop over each point
        else:
            slicers = []
            for i in range(0, dim_sizes[x], size):
                slicers.append(slice(i, min(i+size, dim_sizes[x]), None))
        # print final looping decisions:
        if DEFAULTS.DEBUG >= 3:
            print(f'xarray_multiply_then_sum_along using x={x!r}, size={size}, (due to mode={mode!r})')
        ## looping ##
        if slicers is None:  # not looping after all!
            xarray_result_size_check(*arrays, errmsg=errmsg)
            result = xarray_sum(product(arrays), dim=dim, missing_dims='ignore')
        elif x == dim:   # looping over the dimension being integrated along!
            for i, idx in enumerate(slicers):
                arrays_i = [a.isel({x: idx}, drop=True) if x in a.dims else a for a in arrays]
                if i == 0:
                    xarray_result_size_check(*arrays_i, errmsg=errmsg)
                    result = xarray_sum(product(arrays_i), dim=dim, missing_dims='ignore')
                else:
                    result += xarray_sum(product(arrays_i), dim=dim, missing_dims='ignore')
        else:  # looping over any other dimension (must concat at the end)
            result = []
            for i, idx in enumerate(slicers):
                arrays_i = [a.isel({x: idx}, drop=True) if x in a.dims else a for a in arrays]
                if i == 0:
                    xarray_result_size_check(*arrays_i, errmsg=errmsg)
                result.append(xarray_sum(product(arrays_i), dim=dim, missing_dims='ignore'))
            result = xr.concat(result, dim=x)  # concat along loop dimension
        return result
