# -*- coding: utf-8 -*-

from __future__ import division


__author__ = """John Kirkham"""
__email__ = "kirkhamj@janelia.hhmi.org"


from ._version import get_versions
__version__ = get_versions()['version']
del get_versions


import numpy

import dask.array

from . import _compat
from . import _pycompat
from . import _utils


def center_of_mass(input, labels=None, index=None):
    """
    Calculate the center of mass of the values of an array at labels.

    Parameters
    ----------
    input : ndarray
        Data from which to calculate center-of-mass. The masses can either
        be positive or negative.
    labels : ndarray, optional
        Labels for objects in `input`, as generated by `ndimage.label`.
        Only used with `index`.  Dimensions must be the same as `input`.
    index : int or sequence of ints, optional
        Labels for which to calculate centers-of-mass. If not specified,
        all labels greater than zero are used.  Only used with `labels`.

    Returns
    -------
    center_of_mass : tuple, or list of tuples
        Coordinates of centers-of-mass.
    """

    input, labels, index = _utils._norm_input_labels_index(
        input, labels, index
    )

    # SciPy transposes these for some reason.
    # So we do the same thing here.
    # This only matters if index is some array.
    index = index.T

    input_i = _compat._indices(
        input.shape, dtype=numpy.int64, chunks=input.chunks
    )
    input_i_wt = input[None] * input_i

    lbl_mtch = _utils._get_label_matches(labels, index)

    input_i_wt_mtch = dask.array.where(
        lbl_mtch[index.ndim * (slice(None),) + (None,)],
        input_i_wt[index.ndim * (None,)],
        input_i_wt.dtype.type(0)
    )

    input_mtch = dask.array.where(
        lbl_mtch, input[index.ndim * (None,)], input.dtype.type(0)
    )

    input_i_wt_mtch = input_i_wt_mtch.astype(numpy.float64)
    input_mtch = input_mtch.astype(numpy.float64)

    com_lbl = input_i_wt_mtch.sum(
        axis=tuple(_pycompat.irange(1 + index.ndim, input_i_wt_mtch.ndim))
    )
    input_mtch_sum = input_mtch.sum(
        axis=tuple(_pycompat.irange(index.ndim, input_mtch.ndim))
    )
    com_lbl /= input_mtch_sum[index.ndim * (slice(None),) + (None,)]

    return com_lbl

@@ -2254,8 +2254,1215 @@
 com_lbl%0A
+%0A%0Adef sum(input, labels=None, index=None):%0A    %22%22%22%0A    Calculate the sum of the values of the array.%0A%0A    Parameters%0A    ----------%0A    input : array_like%0A        Values of %60input%60 inside the regions defined by %60labels%60%0A        are summed together.%0A    labels : array_like of ints, optional%0A        Assign labels to the values of the array. Has to have the same shape as%0A        %60input%60.%0A    index : array_like, optional%0A        A single label number or a sequence of label numbers of%0A        the objects to be measured.%0A%0A    Returns%0A    -------%0A    sum : ndarray or scalar%0A        An array of the sums of values of %60input%60 inside the regions defined%0A        by %60labels%60 with the same shape as %60index%60. If 'index' is None or scalar,%0A        a scalar is returned.%0A    %22%22%22%0A%0A    input, labels, index = _utils._norm_input_labels_index(%0A        input, labels, index%0A    )%0A%0A    lbl_mtch = _utils._get_label_matches(labels, index)%0A%0A    input_mtch = dask.array.where(%0A        lbl_mtch, input%5Bindex.ndim * (None,)%5D, input.dtype.type(0)%0A    )%0A%0A    input_mtch = input_mtch.astype(numpy.float64)%0A%0A    sum_lbl = input_mtch.sum(%0A        axis=tuple(_pycompat.irange(index.ndim, input_mtch.ndim))%0A    )%0A%0A    return sum_lbl%0A

Add sum function for label images