from functools import partial

import numpy as np


def _np_grouped_op(group_idx, array, op, axis=-1, size=None, fill_value=None, dtype=None, out=None):
    """
    most of this code is from shoyer's gist
    https://gist.github.com/shoyer/f538ac78ae904c936844
    """
    # assumes input is sorted, which I do in core._prepare_for_flox
    aux = group_idx

    flag = np.concatenate(([True], aux[1:] != aux[:-1]))
    uniques = aux[flag]
    (inv_idx,) = flag.nonzero()

    if size is None:
        size = np.max(uniques) + 1
    if dtype is None:
        dtype = array.dtype

    if out is None:
        out = np.full(array.shape[:-1] + (size,), fill_value=fill_value, dtype=dtype)

    if (len(uniques) == size) and (uniques == np.arange(size)).all():
        # The previous version of this if condition
        #     ((uniques[1:] - uniques[:-1]) == 1).all():
        # does not work when group_idx is [1, 2] for e.g.
        # This happens  during binning
        op.reduceat(array, inv_idx, axis=axis, dtype=dtype, out=out)
    else:
        out[..., uniques] = op.reduceat(array, inv_idx, axis=axis, dtype=dtype)

    return out


def _nan_grouped_op(group_idx, array, func, fillna, *args, **kwargs):
    result = func(group_idx, np.where(np.isnan(array), fillna, array), *args, **kwargs)
    # np.nanmax([np.nan, np.nan]) = np.nan
    # To recover this behaviour, we need to search for the fillna value
    # (either np.inf or -np.inf), and replace with NaN
    # Our choice of fillna does the right thing for sum, prod
    if fillna in (np.inf, -np.inf):
        allnangroups = result == fillna
        if allnangroups.any():
            result[allnangroups] = kwargs["fill_value"]
    return result


sum = partial(_np_grouped_op, op=np.add)
nansum = partial(_nan_grouped_op, func=sum, fillna=0)
prod = partial(_np_grouped_op, op=np.multiply)
nanprod = partial(_nan_grouped_op, func=prod, fillna=1)
max = partial(_np_grouped_op, op=np.maximum)
nanmax = partial(_nan_grouped_op, func=max, fillna=-np.inf)
min = partial(_np_grouped_op, op=np.minimum)
nanmin = partial(_nan_grouped_op, func=min, fillna=np.inf)
# TODO: all, any


def sum_of_squares(group_idx, array, *, axis=-1, size=None, fill_value=None, dtype=None):

    return sum(
        group_idx,
        array**2,
        axis=axis,
        size=size,
        fill_value=fill_value,
        dtype=dtype,
    )


def nansum_of_squares(group_idx, array, *, axis=-1, size=None, fill_value=None, dtype=None):
    return sum_of_squares(
        group_idx,
        np.where(np.isnan(array), 0, array),
        size=size,
        fill_value=fill_value,
        axis=axis,
        dtype=dtype,
    )


def nanlen(group_idx, array, *args, **kwargs):
    return sum(group_idx, (~np.isnan(array)).astype(int), *args, **kwargs)


def mean(group_idx, array, *, axis=-1, size=None, fill_value=None, dtype=None):
    if fill_value is None:
        fill_value = 0
    out = np.full(array.shape[:-1] + (size,), fill_value=fill_value, dtype=dtype)
    sum(group_idx, array, axis=axis, size=size, dtype=dtype, out=out)
    out /= nanlen(group_idx, array, size=size, axis=axis, fill_value=0)
    return out


def nanmean(group_idx, array, *, axis=-1, size=None, fill_value=None, dtype=None):
    if fill_value is None:
        fill_value = 0
    out = np.full(array.shape[:-1] + (size,), fill_value=fill_value, dtype=dtype)
    nansum(group_idx, array, size=size, axis=axis, dtype=dtype, out=out)
    out /= nanlen(group_idx, array, size=size, axis=axis, fill_value=0)
    return out
