import functools
import operator
import pickle
from collections import OrderedDict
from distutils.version import LooseVersion

import numpy as np
import pandas as pd
import pytest
from numpy.testing import assert_array_equal

import xarray as xr
from xarray.core.computation import (
    _UFuncSignature, apply_ufunc, broadcast_compat_data, collect_dict_values,
    join_dict_keys, ordered_set_intersection, ordered_set_union, result_name,
    unified_dim_sizes)

from . import has_dask, raises_regex, requires_dask


def assert_identical(a, b):
    if hasattr(a, 'identical'):
        msg = 'not identical:\n%r\n%r' % (a, b)
        assert a.identical(b), msg
    else:
        assert_array_equal(a, b)


def test_signature_properties():
    sig = _UFuncSignature([['x'], ['x', 'y']], [['z']])
    assert sig.input_core_dims == (('x',), ('x', 'y'))
    assert sig.output_core_dims == (('z',),)
    assert sig.all_input_core_dims == frozenset(['x', 'y'])
    assert sig.all_output_core_dims == frozenset(['z'])
    assert sig.num_inputs == 2
    assert sig.num_outputs == 1
    assert str(sig) == '(x),(x,y)->(z)'
    assert sig.to_gufunc_string() == '(dim0),(dim0,dim1)->(dim2)'
    # dimension names matter
    assert _UFuncSignature([['x']]) != _UFuncSignature([['y']])


def test_result_name():

    class Named:
        def __init__(self, name=None):
            self.name = name

    assert result_name([1, 2]) is None
    assert result_name([Named()]) is None
    assert result_name([Named('foo'), 2]) == 'foo'
    assert result_name([Named('foo'), Named('bar')]) is None
    assert result_name([Named('foo'), Named()]) is None


def test_ordered_set_union():
    assert list(ordered_set_union([[1, 2]])) == [1, 2]
    assert list(ordered_set_union([[1, 2], [2, 1]])) == [1, 2]
    assert list(ordered_set_union([[0], [1, 2], [1, 3]])) == [0, 1, 2, 3]


def test_ordered_set_intersection():
    assert list(ordered_set_intersection([[1, 2]])) == [1, 2]
    assert list(ordered_set_intersection([[1, 2], [2, 1]])) == [1, 2]
    assert list(ordered_set_intersection([[1, 2], [1, 3]])) == [1]
    assert list(ordered_set_intersection([[1, 2], [2]])) == [2]


def test_join_dict_keys():
    dicts = [OrderedDict.fromkeys(keys) for keys in [['x', 'y'], ['y', 'z']]]
    assert list(join_dict_keys(dicts, 'left')) == ['x', 'y']
    assert list(join_dict_keys(dicts, 'right')) == ['y', 'z']
    assert list(join_dict_keys(dicts, 'inner')) == ['y']
    assert list(join_dict_keys(dicts, 'outer')) == ['x', 'y', 'z']
    with pytest.raises(ValueError):
        join_dict_keys(dicts, 'exact')
    with pytest.raises(KeyError):
        join_dict_keys(dicts, 'foobar')


def test_collect_dict_values():
    dicts = [{'x': 1, 'y': 2, 'z': 3}, {'z': 4}, 5]
    expected = [[1, 0, 5], [2, 0, 5], [3, 4, 5]]
    collected = collect_dict_values(dicts, ['x', 'y', 'z'], fill_value=0)
    assert collected == expected


def identity(x):
    return x


def test_apply_identity():
    array = np.arange(10)
    variable = xr.Variable('x', array)
    data_array = xr.DataArray(variable, [('x', -array)])
    dataset = xr.Dataset({'y': variable}, {'x': -array})

    apply_identity = functools.partial(apply_ufunc, identity)

    assert_identical(array, apply_identity(array))
    assert_identical(variable, apply_identity(variable))
    assert_identical(data_array, apply_identity(data_array))
    assert_identical(data_array, apply_identity(data_array.groupby('x')))
    assert_identical(dataset, apply_identity(dataset))
    assert_identical(dataset, apply_identity(dataset.groupby('x')))


def add(a, b):
    return apply_ufunc(operator.add, a, b)


def test_apply_two_inputs():
    array = np.array([1, 2, 3])
    variable = xr.Variable('x', array)
    data_array = xr.DataArray(variable, [('x', -array)])
    dataset = xr.Dataset({'y': variable}, {'x': -array})

    zero_array = np.zeros_like(array)
    zero_variable = xr.Variable('x', zero_array)
    zero_data_array = xr.DataArray(zero_variable, [('x', -array)])
    zero_dataset = xr.Dataset({'y': zero_variable}, {'x': -array})

    assert_identical(array, add(array, zero_array))
    assert_identical(array, add(zero_array, array))

    assert_identical(variable, add(variable, zero_array))
    assert_identical(variable, add(variable, zero_variable))
    assert_identical(variable, add(zero_array, variable))
    assert_identical(variable, add(zero_variable, variable))

    assert_identical(data_array, add(data_array, zero_array))
    assert_identical(data_array, add(data_array, zero_variable))
    assert_identical(data_array, add(data_array, zero_data_array))
    assert_identical(data_array, add(zero_array, data_array))
    assert_identical(data_array, add(zero_variable, data_array))
    assert_identical(data_array, add(zero_data_array, data_array))

    assert_identical(dataset, add(dataset, zero_array))
    assert_identical(dataset, add(dataset, zero_variable))
    assert_identical(dataset, add(dataset, zero_data_array))
    assert_identical(dataset, add(dataset, zero_dataset))
    assert_identical(dataset, add(zero_array, dataset))
    assert_identical(dataset, add(zero_variable, dataset))
    assert_identical(dataset, add(zero_data_array, dataset))
    assert_identical(dataset, add(zero_dataset, dataset))

    assert_identical(data_array, add(data_array.groupby('x'), zero_data_array))
    assert_identical(data_array, add(zero_data_array, data_array.groupby('x')))

    assert_identical(dataset, add(data_array.groupby('x'), zero_dataset))
    assert_identical(dataset, add(zero_dataset, data_array.groupby('x')))

    assert_identical(dataset, add(dataset.groupby('x'), zero_data_array))
    assert_identical(dataset, add(dataset.groupby('x'), zero_dataset))
    assert_identical(dataset, add(zero_data_array, dataset.groupby('x')))
    assert_identical(dataset, add(zero_dataset, dataset.groupby('x')))


def test_apply_1d_and_0d():
    array = np.array([1, 2, 3])
    variable = xr.Variable('x', array)
    data_array = xr.DataArray(variable, [('x', -array)])
    dataset = xr.Dataset({'y': variable}, {'x': -array})

    zero_array = 0
    zero_variable = xr.Variable((), zero_array)
    zero_data_array = xr.DataArray(zero_variable)
    zero_dataset = xr.Dataset({'y': zero_variable})

    assert_identical(array, add(array, zero_array))
    assert_identical(array, add(zero_array, array))

    assert_identical(variable, add(variable, zero_array))
    assert_identical(variable, add(variable, zero_variable))
    assert_identical(variable, add(zero_array, variable))
    assert_identical(variable, add(zero_variable, variable))

    assert_identical(data_array, add(data_array, zero_array))
    assert_identical(data_array, add(data_array, zero_variable))
    assert_identical(data_array, add(data_array, zero_data_array))
    assert_identical(data_array, add(zero_array, data_array))
    assert_identical(data_array, add(zero_variable, data_array))
    assert_identical(data_array, add(zero_data_array, data_array))

    assert_identical(dataset, add(dataset, zero_array))
    assert_identical(dataset, add(dataset, zero_variable))
    assert_identical(dataset, add(dataset, zero_data_array))
    assert_identical(dataset, add(dataset, zero_dataset))
    assert_identical(dataset, add(zero_array, dataset))
    assert_identical(dataset, add(zero_variable, dataset))
    assert_identical(dataset, add(zero_data_array, dataset))
    assert_identical(dataset, add(zero_dataset, dataset))

    assert_identical(data_array, add(data_array.groupby('x'), zero_data_array))
    assert_identical(data_array, add(zero_data_array, data_array.groupby('x')))

    assert_identical(dataset, add(data_array.groupby('x'), zero_dataset))
    assert_identical(dataset, add(zero_dataset, data_array.groupby('x')))

    assert_identical(dataset, add(dataset.groupby('x'), zero_data_array))
    assert_identical(dataset, add(dataset.groupby('x'), zero_dataset))
    assert_identical(dataset, add(zero_data_array, dataset.groupby('x')))
    assert_identical(dataset, add(zero_dataset, dataset.groupby('x')))


def test_apply_two_outputs():
    array = np.arange(5)
    variable = xr.Variable('x', array)
    data_array = xr.DataArray(variable, [('x', -array)])
    dataset = xr.Dataset({'y': variable}, {'x': -array})

    def twice(obj):
        def func(x):
            return (x, x)
        return apply_ufunc(func, obj, output_core_dims=[[], []])

    out0, out1 = twice(array)
    assert_identical(out0, array)
    assert_identical(out1, array)

    out0, out1 = twice(variable)
    assert_identical(out0, variable)
    assert_identical(out1, variable)

    out0, out1 = twice(data_array)
    assert_identical(out0, data_array)
    assert_identical(out1, data_array)

    out0, out1 = twice(dataset)
    assert_identical(out0, dataset)
    assert_identical(out1, dataset)

    out0, out1 = twice(data_array.groupby('x'))
    assert_identical(out0, data_array)
    assert_identical(out1, data_array)

    out0, out1 = twice(dataset.groupby('x'))
    assert_identical(out0, dataset)
    assert_identical(out1, dataset)


def test_apply_input_core_dimension():

    def first_element(obj, dim):
        def func(x):
            return x[..., 0]
        return apply_ufunc(func, obj, input_core_dims=[[dim]])

    array = np.array([[1, 2], [3, 4]])
    variable = xr.Variable(['x', 'y'], array)
    data_array = xr.DataArray(variable, {'x': ['a', 'b'], 'y': [-1, -2]})
    dataset = xr.Dataset({'data': data_array})

    expected_variable_x = xr.Variable(['y'], [1, 2])
    expected_data_array_x = xr.DataArray(expected_variable_x, {'y': [-1, -2]})
    expected_dataset_x = xr.Dataset({'data': expected_data_array_x})

    expected_variable_y = xr.Variable(['x'], [1, 3])
    expected_data_array_y = xr.DataArray(expected_variable_y,
                                         {'x': ['a', 'b']})
    expected_dataset_y = xr.Dataset({'data': expected_data_array_y})

    assert_identical(expected_variable_x, first_element(variable, 'x'))
    assert_identical(expected_variable_y, first_element(variable, 'y'))

    assert_identical(expected_data_array_x, first_element(data_array, 'x'))
    assert_identical(expected_data_array_y, first_element(data_array, 'y'))

    assert_identical(expected_dataset_x, first_element(dataset, 'x'))
    assert_identical(expected_dataset_y, first_element(dataset, 'y'))

    assert_identical(expected_data_array_x,
                     first_element(data_array.groupby('y'), 'x'))
    assert_identical(expected_dataset_x,
                     first_element(dataset.groupby('y'), 'x'))

    def multiply(*args):
        val = args[0]
        for arg in args[1:]:
            val = val * arg
        return val

    # regression test for GH:2341
    with pytest.raises(ValueError):
        apply_ufunc(multiply, data_array, data_array['y'].values,
                    input_core_dims=[['y']], output_core_dims=[['y']])
    expected = xr.DataArray(multiply(data_array, data_array['y']),
                            dims=['x', 'y'], coords=data_array.coords)
    actual = apply_ufunc(multiply, data_array, data_array['y'].values,
                         input_core_dims=[['y'], []], output_core_dims=[['y']])
    assert_identical(expected, actual)


def test_apply_output_core_dimension():

    def stack_negative(obj):
        def func(x):
            return np.stack([x, -x], axis=-1)
        result = apply_ufunc(func, obj, output_core_dims=[['sign']])
        if isinstance(result, (xr.Dataset, xr.DataArray)):
            result.coords['sign'] = [1, -1]
        return result

    array = np.array([[1, 2], [3, 4]])
    variable = xr.Variable(['x', 'y'], array)
    data_array = xr.DataArray(variable, {'x': ['a', 'b'], 'y': [-1, -2]})
    dataset = xr.Dataset({'data': data_array})

    stacked_array = np.array([[[1, -1], [2, -2]], [[3, -3], [4, -4]]])
    stacked_variable = xr.Variable(['x', 'y', 'sign'], stacked_array)
    stacked_coords = {'x': ['a', 'b'], 'y': [-1, -2], 'sign': [1, -1]}
    stacked_data_array = xr.DataArray(stacked_variable, stacked_coords)
    stacked_dataset = xr.Dataset({'data': stacked_data_array})

    assert_identical(stacked_array, stack_negative(array))
    assert_identical(stacked_variable, stack_negative(variable))
    assert_identical(stacked_data_array, stack_negative(data_array))
    assert_identical(stacked_dataset, stack_negative(dataset))
    assert_identical(stacked_data_array,
                     stack_negative(data_array.groupby('x')))
    assert_identical(stacked_dataset,
                     stack_negative(dataset.groupby('x')))

    def original_and_stack_negative(obj):
        def func(x):
            return (x, np.stack([x, -x], axis=-1))
        result = apply_ufunc(func, obj, output_core_dims=[[], ['sign']])
        if isinstance(result[1], (xr.Dataset, xr.DataArray)):
            result[1].coords['sign'] = [1, -1]
        return result

    out0, out1 = original_and_stack_negative(array)
    assert_identical(array, out0)
    assert_identical(stacked_array, out1)

    out0, out1 = original_and_stack_negative(variable)
    assert_identical(variable, out0)
    assert_identical(stacked_variable, out1)

    out0, out1 = original_and_stack_negative(data_array)
    assert_identical(data_array, out0)
    assert_identical(stacked_data_array, out1)

    out0, out1 = original_and_stack_negative(dataset)
    assert_identical(dataset, out0)
    assert_identical(stacked_dataset, out1)

    out0, out1 = original_and_stack_negative(data_array.groupby('x'))
    assert_identical(data_array, out0)
    assert_identical(stacked_data_array, out1)

    out0, out1 = original_and_stack_negative(dataset.groupby('x'))
    assert_identical(dataset, out0)
    assert_identical(stacked_dataset, out1)


def test_apply_exclude():

    def concatenate(objects, dim='x'):
        def func(*x):
            return np.concatenate(x, axis=-1)
        result = apply_ufunc(func, *objects,
                             input_core_dims=[[dim]] * len(objects),
                             output_core_dims=[[dim]],
                             exclude_dims={dim})
        if isinstance(result, (xr.Dataset, xr.DataArray)):
            # note: this will fail if dim is not a coordinate on any input
            new_coord = np.concatenate([obj.coords[dim] for obj in objects])
            result.coords[dim] = new_coord
        return result

    arrays = [np.array([1]), np.array([2, 3])]
    variables = [xr.Variable('x', a) for a in arrays]
    data_arrays = [xr.DataArray(v, {'x': c, 'y': ('x', range(len(c)))})
                   for v, c in zip(variables, [['a'], ['b', 'c']])]
    datasets = [xr.Dataset({'data': data_array}) for data_array in data_arrays]

    expected_array = np.array([1, 2, 3])
    expected_variable = xr.Variable('x', expected_array)
    expected_data_array = xr.DataArray(expected_variable, [('x', list('abc'))])
    expected_dataset = xr.Dataset({'data': expected_data_array})

    assert_identical(expected_array, concatenate(arrays))
    assert_identical(expected_variable, concatenate(variables))
    assert_identical(expected_data_array, concatenate(data_arrays))
    assert_identical(expected_dataset, concatenate(datasets))

    # must also be a core dimension
    with pytest.raises(ValueError):
        apply_ufunc(identity, variables[0], exclude_dims={'x'})


def test_apply_groupby_add():
    array = np.arange(5)
    variable = xr.Variable('x', array)
    coords = {'x': -array, 'y': ('x', [0, 0, 1, 1, 2])}
    data_array = xr.DataArray(variable, coords, dims='x')
    dataset = xr.Dataset({'z': variable}, coords)

    other_variable = xr.Variable('y', [0, 10])
    other_data_array = xr.DataArray(other_variable, dims='y')
    other_dataset = xr.Dataset({'z': other_variable})

    expected_variable = xr.Variable('x', [0, 1, 12, 13, np.nan])
    expected_data_array = xr.DataArray(expected_variable, coords, dims='x')
    expected_dataset = xr.Dataset({'z': expected_variable}, coords)

    assert_identical(expected_data_array,
                     add(data_array.groupby('y'), other_data_array))
    assert_identical(expected_dataset,
                     add(data_array.groupby('y'), other_dataset))
    assert_identical(expected_dataset,
                     add(dataset.groupby('y'), other_data_array))
    assert_identical(expected_dataset,
                     add(dataset.groupby('y'), other_dataset))

    # cannot be performed with xarray.Variable objects that share a dimension
    with pytest.raises(ValueError):
        add(data_array.groupby('y'), other_variable)

    # if they are all grouped the same way
    with pytest.raises(ValueError):
        add(data_array.groupby('y'), data_array[:4].groupby('y'))
    with pytest.raises(ValueError):
        add(data_array.groupby('y'), data_array[1:].groupby('y'))
    with pytest.raises(ValueError):
        add(data_array.groupby('y'), other_data_array.groupby('y'))
    with pytest.raises(ValueError):
        add(data_array.groupby('y'), data_array.groupby('x'))


def test_unified_dim_sizes():
    assert unified_dim_sizes([xr.Variable((), 0)]) == OrderedDict()
    assert (unified_dim_sizes([xr.Variable('x', [1]),
                               xr.Variable('x', [1])]) ==
            OrderedDict([('x', 1)]))
    assert (unified_dim_sizes([xr.Variable('x', [1]),
                               xr.Variable('y', [1, 2])]) ==
            OrderedDict([('x', 1), ('y', 2)]))
    assert (unified_dim_sizes([xr.Variable(('x', 'z'), [[1]]),
                               xr.Variable(('y', 'z'), [[1, 2], [3, 4]])],
                              exclude_dims={'z'}) ==
            OrderedDict([('x', 1), ('y', 2)]))

    # duplicate dimensions
    with pytest.raises(ValueError):
        unified_dim_sizes([xr.Variable(('x', 'x'), [[1]])])

    # mismatched lengths
    with pytest.raises(ValueError):
        unified_dim_sizes(
            [xr.Variable('x', [1]), xr.Variable('x', [1, 2])])


def test_broadcast_compat_data_1d():
    data = np.arange(5)
    var = xr.Variable('x', data)

    assert_identical(data, broadcast_compat_data(var, ('x',), ()))
    assert_identical(data, broadcast_compat_data(var, (), ('x',)))
    assert_identical(data[:], broadcast_compat_data(var, ('w',), ('x',)))
    assert_identical(data[:, None],
                     broadcast_compat_data(var, ('w', 'x', 'y'), ()))

    with pytest.raises(ValueError):
        broadcast_compat_data(var, ('x',), ('w',))

    with pytest.raises(ValueError):
        broadcast_compat_data(var, (), ())


def test_broadcast_compat_data_2d():
    data = np.arange(12).reshape(3, 4)
    var = xr.Variable(['x', 'y'], data)

    assert_identical(data, broadcast_compat_data(var, ('x', 'y'), ()))
    assert_identical(data, broadcast_compat_data(var, ('x',), ('y',)))
    assert_identical(data, broadcast_compat_data(var, (), ('x', 'y')))
    assert_identical(data.T, broadcast_compat_data(var, ('y', 'x'), ()))
    assert_identical(data.T, broadcast_compat_data(var, ('y',), ('x',)))
    assert_identical(data, broadcast_compat_data(var, ('w', 'x'), ('y',)))
    assert_identical(data, broadcast_compat_data(var, ('w',), ('x', 'y')))
    assert_identical(data.T, broadcast_compat_data(var, ('w',), ('y', 'x')))
    assert_identical(data[:, :, None],
                     broadcast_compat_data(var, ('w', 'x', 'y', 'z'), ()))
    assert_identical(data[None, :, :].T,
                     broadcast_compat_data(var, ('w', 'y', 'x', 'z'), ()))


def test_keep_attrs():

    def add(a, b, keep_attrs):
        if keep_attrs:
            return apply_ufunc(operator.add, a, b, keep_attrs=keep_attrs)
        else:
            return apply_ufunc(operator.add, a, b)

    a = xr.DataArray([0, 1], [('x', [0, 1])])
    a.attrs['attr'] = 'da'
    a['x'].attrs['attr'] = 'da_coord'
    b = xr.DataArray([1, 2], [('x', [0, 1])])

    actual = add(a, b, keep_attrs=False)
    assert not actual.attrs
    actual = add(a, b, keep_attrs=True)
    assert_identical(actual.attrs, a.attrs)
    assert_identical(actual['x'].attrs, a['x'].attrs)

    actual = add(a.variable, b.variable, keep_attrs=False)
    assert not actual.attrs
    actual = add(a.variable, b.variable, keep_attrs=True)
    assert_identical(actual.attrs, a.attrs)

    a = xr.Dataset({'x': [0, 1]})
    a.attrs['attr'] = 'ds'
    a.x.attrs['attr'] = 'da'
    b = xr.Dataset({'x': [0, 1]})

    actual = add(a, b, keep_attrs=False)
    assert not actual.attrs
    actual = add(a, b, keep_attrs=True)
    assert_identical(actual.attrs, a.attrs)
    assert_identical(actual.x.attrs, a.x.attrs)


def test_dataset_join():
    ds0 = xr.Dataset({'a': ('x', [1, 2]), 'x': [0, 1]})
    ds1 = xr.Dataset({'a': ('x', [99, 3]), 'x': [1, 2]})

    # by default, cannot have different labels
    with raises_regex(ValueError, 'indexes .* are not equal'):
        apply_ufunc(operator.add, ds0, ds1)
    with raises_regex(TypeError, 'must supply'):
        apply_ufunc(operator.add, ds0, ds1, dataset_join='outer')

    def add(a, b, join, dataset_join):
        return apply_ufunc(operator.add, a, b, join=join,
                           dataset_join=dataset_join,
                           dataset_fill_value=np.nan)

    actual = add(ds0, ds1, 'outer', 'inner')
    expected = xr.Dataset({'a': ('x', [np.nan, 101, np.nan]),
                           'x': [0, 1, 2]})
    assert_identical(actual, expected)

    actual = add(ds0, ds1, 'outer', 'outer')
    assert_identical(actual, expected)

    with raises_regex(ValueError, 'data variable names'):
        apply_ufunc(operator.add, ds0, xr.Dataset({'b': 1}))

    ds2 = xr.Dataset({'b': ('x', [99, 3]), 'x': [1, 2]})
    actual = add(ds0, ds2, 'outer', 'inner')
    expected = xr.Dataset({'x': [0, 1, 2]})
    assert_identical(actual, expected)

    # we used np.nan as the fill_value in add() above
    actual = add(ds0, ds2, 'outer', 'outer')
    expected = xr.Dataset({'a': ('x', [np.nan, np.nan, np.nan]),
                           'b': ('x', [np.nan, np.nan, np.nan]),
                           'x': [0, 1, 2]})
    assert_identical(actual, expected)


@requires_dask
def test_apply_dask():
    import dask.array as da

    array = da.ones((2,), chunks=2)
    variable = xr.Variable('x', array)
    coords = xr.DataArray(variable).coords.variables
    data_array = xr.DataArray(variable, dims=['x'], coords=coords)
    dataset = xr.Dataset({'y': variable})

    # encountered dask array, but did not set dask='allowed'
    with pytest.raises(ValueError):
        apply_ufunc(identity, array)
    with pytest.raises(ValueError):
        apply_ufunc(identity, variable)
    with pytest.raises(ValueError):
        apply_ufunc(identity, data_array)
    with pytest.raises(ValueError):
        apply_ufunc(identity, dataset)

    # unknown setting for dask array handling
    with pytest.raises(ValueError):
        apply_ufunc(identity, array, dask='unknown')

    def dask_safe_identity(x):
        return apply_ufunc(identity, x, dask='allowed')

    assert array is dask_safe_identity(array)

    actual = dask_safe_identity(variable)
    assert isinstance(actual.data, da.Array)
    assert_identical(variable, actual)

    actual = dask_safe_identity(data_array)
    assert isinstance(actual.data, da.Array)
    assert_identical(data_array, actual)

    actual = dask_safe_identity(dataset)
    assert isinstance(actual['y'].data, da.Array)
    assert_identical(dataset, actual)


@requires_dask
def test_apply_dask_parallelized_one_arg():
    import dask.array as da

    array = da.ones((2, 2), chunks=(1, 1))
    data_array = xr.DataArray(array, dims=('x', 'y'))

    def parallel_identity(x):
        return apply_ufunc(identity, x, dask='parallelized',
                           output_dtypes=[x.dtype])

    actual = parallel_identity(data_array)
    assert isinstance(actual.data, da.Array)
    assert actual.data.chunks == array.chunks
    assert_identical(data_array, actual)

    computed = data_array.compute()
    actual = parallel_identity(computed)
    assert_identical(computed, actual)


@requires_dask
def test_apply_dask_parallelized_two_args():
    import dask.array as da

    array = da.ones((2, 2), chunks=(1, 1), dtype=np.int64)
    data_array = xr.DataArray(array, dims=('x', 'y'))
    data_array.name = None

    def parallel_add(x, y):
        return apply_ufunc(operator.add, x, y,
                           dask='parallelized',
                           output_dtypes=[np.int64])

    def check(x, y):
        actual = parallel_add(x, y)
        assert isinstance(actual.data, da.Array)
        assert actual.data.chunks == array.chunks
        assert_identical(data_array, actual)

    check(data_array, 0),
    check(0, data_array)
    check(data_array, xr.DataArray(0))
    check(data_array, 0 * data_array)
    check(data_array, 0 * data_array[0])
    check(data_array[:, 0], 0 * data_array[0])
    check(data_array, 0 * data_array.compute())


@requires_dask
def test_apply_dask_parallelized_errors():
    import dask.array as da

    array = da.ones((2, 2), chunks=(1, 1))
    data_array = xr.DataArray(array, dims=('x', 'y'))

    with pytest.raises(NotImplementedError):
        apply_ufunc(identity, data_array, output_core_dims=[['z'], ['z']],
                    dask='parallelized')
    with raises_regex(ValueError, 'dtypes'):
        apply_ufunc(identity, data_array, dask='parallelized')
    with raises_regex(TypeError, 'list'):
        apply_ufunc(identity, data_array, dask='parallelized',
                    output_dtypes=float)
    with raises_regex(ValueError, 'must have the same length'):
        apply_ufunc(identity, data_array, dask='parallelized',
                    output_dtypes=[float, float])
    with raises_regex(ValueError, 'output_sizes'):
        apply_ufunc(identity, data_array, output_core_dims=[['z']],
                    output_dtypes=[float], dask='parallelized')
    with raises_regex(ValueError, 'at least one input is an xarray object'):
        apply_ufunc(identity, array, dask='parallelized')

    with raises_regex(ValueError, 'consists of multiple chunks'):
        apply_ufunc(identity, data_array, dask='parallelized',
                    output_dtypes=[float],
                    input_core_dims=[('y',)],
                    output_core_dims=[('y',)])


@requires_dask
def test_apply_dask_multiple_inputs():
    import dask.array as da

    def covariance(x, y):
        return ((x - x.mean(axis=-1, keepdims=True)) *
                (y - y.mean(axis=-1, keepdims=True))).mean(axis=-1)

    rs = np.random.RandomState(42)
    array1 = da.from_array(rs.randn(4, 4), chunks=(2, 4))
    array2 = da.from_array(rs.randn(4, 4), chunks=(2, 4))
    data_array_1 = xr.DataArray(array1, dims=('x', 'z'))
    data_array_2 = xr.DataArray(array2, dims=('y', 'z'))

    expected = apply_ufunc(
        covariance, data_array_1.compute(), data_array_2.compute(),
        input_core_dims=[['z'], ['z']])
    allowed = apply_ufunc(
        covariance, data_array_1, data_array_2, input_core_dims=[['z'], ['z']],
        dask='allowed')
    assert isinstance(allowed.data, da.Array)
    xr.testing.assert_allclose(expected, allowed.compute())

    parallelized = apply_ufunc(
        covariance, data_array_1, data_array_2, input_core_dims=[['z'], ['z']],
        dask='parallelized', output_dtypes=[float])
    assert isinstance(parallelized.data, da.Array)
    xr.testing.assert_allclose(expected, parallelized.compute())


@requires_dask
def test_apply_dask_new_output_dimension():
    import dask.array as da

    array = da.ones((2, 2), chunks=(1, 1))
    data_array = xr.DataArray(array, dims=('x', 'y'))

    def stack_negative(obj):
        def func(x):
            return np.stack([x, -x], axis=-1)
        return apply_ufunc(func, obj, output_core_dims=[['sign']],
                           dask='parallelized', output_dtypes=[obj.dtype],
                           output_sizes={'sign': 2})

    expected = stack_negative(data_array.compute())

    actual = stack_negative(data_array)
    assert actual.dims == ('x', 'y', 'sign')
    assert actual.shape == (2, 2, 2)
    assert isinstance(actual.data, da.Array)
    assert_identical(expected, actual)


def pandas_median(x):
    return pd.Series(x).median()


def test_vectorize():
    data_array = xr.DataArray([[0, 1, 2], [1, 2, 3]], dims=('x', 'y'))
    expected = xr.DataArray([1, 2], dims=['x'])
    actual = apply_ufunc(pandas_median, data_array,
                         input_core_dims=[['y']],
                         vectorize=True)
    assert_identical(expected, actual)


@requires_dask
def test_vectorize_dask():
    data_array = xr.DataArray([[0, 1, 2], [1, 2, 3]], dims=('x', 'y'))
    expected = xr.DataArray([1, 2], dims=['x'])
    actual = apply_ufunc(pandas_median, data_array.chunk({'x': 1}),
                         input_core_dims=[['y']],
                         vectorize=True,
                         dask='parallelized',
                         output_dtypes=[float])
    assert_identical(expected, actual)


def test_output_wrong_number():
    variable = xr.Variable('x', np.arange(10))

    def identity(x):
        return x

    def tuple3x(x):
        return (x, x, x)

    with raises_regex(ValueError, 'number of outputs'):
        apply_ufunc(identity, variable, output_core_dims=[(), ()])

    with raises_regex(ValueError, 'number of outputs'):
        apply_ufunc(tuple3x, variable, output_core_dims=[(), ()])


def test_output_wrong_dims():
    variable = xr.Variable('x', np.arange(10))

    def add_dim(x):
        return x[..., np.newaxis]

    def remove_dim(x):
        return x[..., 0]

    with raises_regex(ValueError, 'unexpected number of dimensions'):
        apply_ufunc(add_dim, variable, output_core_dims=[('y', 'z')])

    with raises_regex(ValueError, 'unexpected number of dimensions'):
        apply_ufunc(add_dim, variable)

    with raises_regex(ValueError, 'unexpected number of dimensions'):
        apply_ufunc(remove_dim, variable)


def test_output_wrong_dim_size():
    array = np.arange(10)
    variable = xr.Variable('x', array)
    data_array = xr.DataArray(variable, [('x', -array)])
    dataset = xr.Dataset({'y': variable}, {'x': -array})

    def truncate(array):
        return array[:5]

    def apply_truncate_broadcast_invalid(obj):
        return apply_ufunc(truncate, obj)

    with raises_regex(ValueError, 'size of dimension'):
        apply_truncate_broadcast_invalid(variable)
    with raises_regex(ValueError, 'size of dimension'):
        apply_truncate_broadcast_invalid(data_array)
    with raises_regex(ValueError, 'size of dimension'):
        apply_truncate_broadcast_invalid(dataset)

    def apply_truncate_x_x_invalid(obj):
        return apply_ufunc(truncate, obj, input_core_dims=[['x']],
                           output_core_dims=[['x']])

    with raises_regex(ValueError, 'size of dimension'):
        apply_truncate_x_x_invalid(variable)
    with raises_regex(ValueError, 'size of dimension'):
        apply_truncate_x_x_invalid(data_array)
    with raises_regex(ValueError, 'size of dimension'):
        apply_truncate_x_x_invalid(dataset)

    def apply_truncate_x_z(obj):
        return apply_ufunc(truncate, obj, input_core_dims=[['x']],
                           output_core_dims=[['z']])

    assert_identical(xr.Variable('z', array[:5]),
                     apply_truncate_x_z(variable))
    assert_identical(xr.DataArray(array[:5], dims=['z']),
                     apply_truncate_x_z(data_array))
    assert_identical(xr.Dataset({'y': ('z', array[:5])}),
                     apply_truncate_x_z(dataset))

    def apply_truncate_x_x_valid(obj):
        return apply_ufunc(truncate, obj, input_core_dims=[['x']],
                           output_core_dims=[['x']], exclude_dims={'x'})

    assert_identical(xr.Variable('x', array[:5]),
                     apply_truncate_x_x_valid(variable))
    assert_identical(xr.DataArray(array[:5], dims=['x']),
                     apply_truncate_x_x_valid(data_array))
    assert_identical(xr.Dataset({'y': ('x', array[:5])}),
                     apply_truncate_x_x_valid(dataset))


@pytest.mark.parametrize('use_dask', [True, False])
def test_dot(use_dask):
    if use_dask:
        if not has_dask:
            pytest.skip('test for dask.')

    a = np.arange(30 * 4).reshape(30, 4)
    b = np.arange(30 * 4 * 5).reshape(30, 4, 5)
    c = np.arange(5 * 60).reshape(5, 60)
    da_a = xr.DataArray(a, dims=['a', 'b'],
                        coords={'a': np.linspace(0, 1, 30)})
    da_b = xr.DataArray(b, dims=['a', 'b', 'c'],
                        coords={'a': np.linspace(0, 1, 30)})
    da_c = xr.DataArray(c, dims=['c', 'e'])
    if use_dask:
        da_a = da_a.chunk({'a': 3})
        da_b = da_b.chunk({'a': 3})
        da_c = da_c.chunk({'c': 3})

    actual = xr.dot(da_a, da_b, dims=['a', 'b'])
    assert actual.dims == ('c', )
    assert (actual.data == np.einsum('ij,ijk->k', a, b)).all()
    assert isinstance(actual.variable.data, type(da_a.variable.data))

    actual = xr.dot(da_a, da_b)
    assert actual.dims == ('c', )
    assert (actual.data == np.einsum('ij,ijk->k', a, b)).all()
    assert isinstance(actual.variable.data, type(da_a.variable.data))

    if use_dask:
        import dask
        if LooseVersion(dask.__version__) < LooseVersion('0.17.3'):
            pytest.skip("needs dask.array.einsum")

    # for only a single array is passed without dims argument, just return
    # as is
    actual = xr.dot(da_a)
    assert da_a.identical(actual)

    # test for variable
    actual = xr.dot(da_a.variable, da_b.variable)
    assert actual.dims == ('c', )
    assert (actual.data == np.einsum('ij,ijk->k', a, b)).all()
    assert isinstance(actual.data, type(da_a.variable.data))

    if use_dask:
        da_a = da_a.chunk({'a': 3})
        da_b = da_b.chunk({'a': 3})
        actual = xr.dot(da_a, da_b, dims=['b'])
        assert actual.dims == ('a', 'c')
        assert (actual.data == np.einsum('ij,ijk->ik', a, b)).all()
        assert isinstance(actual.variable.data, type(da_a.variable.data))

    actual = xr.dot(da_a, da_b, dims=['b'])
    assert actual.dims == ('a', 'c')
    assert (actual.data == np.einsum('ij,ijk->ik', a, b)).all()

    actual = xr.dot(da_a, da_b, dims='b')
    assert actual.dims == ('a', 'c')
    assert (actual.data == np.einsum('ij,ijk->ik', a, b)).all()

    actual = xr.dot(da_a, da_b, dims='a')
    assert actual.dims == ('b', 'c')
    assert (actual.data == np.einsum('ij,ijk->jk', a, b)).all()

    actual = xr.dot(da_a, da_b, dims='c')
    assert actual.dims == ('a', 'b')
    assert (actual.data == np.einsum('ij,ijk->ij', a, b)).all()

    actual = xr.dot(da_a, da_b, da_c, dims=['a', 'b'])
    assert actual.dims == ('c', 'e')
    assert (actual.data == np.einsum('ij,ijk,kl->kl ', a, b, c)).all()

    # should work with tuple
    actual = xr.dot(da_a, da_b, dims=('c', ))
    assert actual.dims == ('a', 'b')
    assert (actual.data == np.einsum('ij,ijk->ij', a, b)).all()

    # default dims
    actual = xr.dot(da_a, da_b, da_c)
    assert actual.dims == ('e', )
    assert (actual.data == np.einsum('ij,ijk,kl->l ', a, b, c)).all()

    # 1 array summation
    actual = xr.dot(da_a, dims='a')
    assert actual.dims == ('b', )
    assert (actual.data == np.einsum('ij->j ', a)).all()

    # empty dim
    actual = xr.dot(da_a.sel(a=[]), da_a.sel(a=[]), dims='a')
    assert actual.dims == ('b', )
    assert (actual.data == np.zeros(actual.shape)).all()

    # Invalid cases
    if not use_dask or LooseVersion(dask.__version__) > LooseVersion('0.17.4'):
        with pytest.raises(TypeError):
            xr.dot(da_a, dims='a', invalid=None)
    with pytest.raises(TypeError):
        xr.dot(da_a.to_dataset(name='da'), dims='a')
    with pytest.raises(TypeError):
        xr.dot(dims='a')

    # einsum parameters
    actual = xr.dot(da_a, da_b, dims=['b'], order='C')
    assert (actual.data == np.einsum('ij,ijk->ik', a, b)).all()
    assert actual.values.flags['C_CONTIGUOUS']
    assert not actual.values.flags['F_CONTIGUOUS']
    actual = xr.dot(da_a, da_b, dims=['b'], order='F')
    assert (actual.data == np.einsum('ij,ijk->ik', a, b)).all()
    # dask converts Fortran arrays to C order when merging the final array
    if not use_dask:
        assert not actual.values.flags['C_CONTIGUOUS']
        assert actual.values.flags['F_CONTIGUOUS']

    # einsum has a constant string as of the first parameter, which makes
    # it hard to pass to xarray.apply_ufunc.
    # make sure dot() uses functools.partial(einsum, subscripts), which
    # can be pickled, and not a lambda, which can't.
    pickle.loads(pickle.dumps(xr.dot(da_a)))


def test_where():
    cond = xr.DataArray([True, False], dims='x')
    actual = xr.where(cond, 1, 0)
    expected = xr.DataArray([1, 0], dims='x')
    assert_identical(expected, actual)
