import numpy as np
from numpy import array, nan
from xray.core import ops
from xray.core.ops import (
    first, last, count, mean
)
from xray.core.nputils import (
    interleaved_concat as interleaved_concat_numpy, _calc_concat_shape,
)

from . import TestCase


class TestOps(TestCase):
    def setUp(self):
        self.x = array([[[ nan,  nan,   2.,  nan],
                         [ nan,   5.,   6.,  nan],
                         [  8.,   9.,  10.,  nan]],

                        [[ nan,  13.,  14.,  15.],
                         [ nan,  17.,  18.,  nan],
                         [ nan,  21.,  nan,  nan]]])

    def test_first(self):
        expected_results = [array([[nan, 13, 2, 15],
                                   [nan, 5, 6, nan],
                                   [8, 9, 10, nan]]),
                            array([[8, 5, 2, nan],
                                   [nan, 13, 14, 15]]),
                            array([[2, 5, 8],
                                   [13, 17, 21]])]
        for axis, expected in zip([0, 1, 2, -3, -2, -1],
                                  2 * expected_results):
            actual = first(self.x, axis)
            self.assertArrayEqual(expected, actual)

        expected = self.x[0]
        actual = first(self.x, axis=0, skipna=False)
        self.assertArrayEqual(expected, actual)

        expected = self.x[..., 0]
        actual = first(self.x, axis=-1, skipna=False)
        self.assertArrayEqual(expected, actual)

        with self.assertRaisesRegexp(IndexError, 'out of bounds'):
            first(self.x, 3)

    def test_last(self):
        expected_results = [array([[nan, 13, 14, 15],
                                   [nan, 17, 18, nan],
                                   [8, 21, 10, nan]]),
                            array([[8, 9, 10, nan],
                                   [nan, 21, 18, 15]]),
                            array([[2, 6, 10],
                                   [15, 18, 21]])]
        for axis, expected in zip([0, 1, 2, -3, -2, -1],
                                  2 * expected_results):
            actual = last(self.x, axis)
            self.assertArrayEqual(expected, actual)

        expected = self.x[-1]
        actual = last(self.x, axis=0, skipna=False)
        self.assertArrayEqual(expected, actual)

        expected = self.x[..., -1]
        actual = last(self.x, axis=-1, skipna=False)
        self.assertArrayEqual(expected, actual)

        with self.assertRaisesRegexp(IndexError, 'out of bounds'):
            last(self.x, 3)

    def test_count(self):
        self.assertEqual(12, count(self.x))

        expected = array([[1, 2, 3], [3, 2, 1]])
        self.assertArrayEqual(expected, count(self.x, axis=-1))

    def test_all_nan_arrays(self):
        assert np.isnan(mean([np.nan, np.nan]))

    def test_interleaved_concat(self):
        for interleaved_concat in [interleaved_concat_numpy,
                                   ops._interleaved_concat_slow,
                                   ops.interleaved_concat]:
            x = np.arange(5)
            self.assertArrayEqual(x, interleaved_concat([x], [x]))

            arrays = np.arange(10).reshape(2, -1)
            indices = np.arange(10).reshape(2, -1, order='F')
            actual = interleaved_concat(arrays, indices)
            expected = np.array([0, 5, 1, 6, 2, 7, 3, 8, 4, 9])
            self.assertArrayEqual(expected, actual)

            arrays2 = arrays.reshape(2, 5, 1)

            actual = interleaved_concat(arrays2, indices, axis=0)
            self.assertArrayEqual(expected.reshape(10, 1), actual)

            actual = interleaved_concat(arrays2, [[0], [1]], axis=1)
            self.assertArrayEqual(arrays.T, actual)

            actual = interleaved_concat(arrays2, [slice(1), slice(1, 2)], axis=-1)
            self.assertArrayEqual(arrays.T, actual)

            with self.assertRaises(IndexError):
                interleaved_concat(arrays, indices, axis=1)
            with self.assertRaises(IndexError):
                interleaved_concat(arrays, indices, axis=-2)
            with self.assertRaises(IndexError):
                interleaved_concat(arrays2, [0, 1], axis=2)

    def test_interleaved_concat_dtypes(self):
        for interleaved_concat in [interleaved_concat_numpy,
                                   ops._interleaved_concat_slow,
                                   ops.interleaved_concat]:
            a = np.array(['a'])
            b = np.array(['bc'])
            actual = interleaved_concat([a, b], [[0], [1]])
            expected = np.array(['a', 'bc'])
            self.assertArrayEqual(expected, actual)

            c = np.array([np.nan], dtype=object)
            actual = interleaved_concat([a, b, c], [[0], [1], [2]])
            expected = np.array(['a', 'bc', np.nan], dtype=object)
            self.assertArrayEqual(expected, actual)

    def test_interleaved_indices_required(self):
        self.assertFalse(ops._interleaved_indices_required([[0]]))
        self.assertFalse(ops._interleaved_indices_required([[0, 1], [2, 3, 4]]))
        self.assertFalse(ops._interleaved_indices_required([slice(3), slice(3, 4)]))
        self.assertFalse(ops._interleaved_indices_required([slice(0, 2, 1)]))
        self.assertTrue(ops._interleaved_indices_required([[0], [2]]))
        self.assertTrue(ops._interleaved_indices_required([[1], [2, 3]]))
        self.assertTrue(ops._interleaved_indices_required([[0, 1], [2, 4]]))
        self.assertTrue(ops._interleaved_indices_required([[0, 1], [3.5, 4]]))
        self.assertTrue(ops._interleaved_indices_required([slice(None, None, 2)]))
