import unittest
import warnings

import numpy
import pytest

import dpnp as cupy
from tests.third_party.cupy import testing


def _dec_shape(shape, dec):
    # Test smaller shape
    return tuple(1 if s == 1 else max(0, s - dec) for s in shape)


class TestEinSumError(unittest.TestCase):

    def test_irregular_ellipsis1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('..', xp.zeros((2, 2, 2)))

    def test_irregular_ellipsis2(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('...i...', xp.zeros((2, 2, 2)))

    def test_irregular_ellipsis3(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i...->...i...', xp.zeros((2, 2, 2)))

    def test_irregular_ellipsis4(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('...->', xp.zeros((2, 2, 2)))

    def test_no_arguments(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum()

    def test_one_argument(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('')

    def test_not_string_subject(self):
        for xp in (numpy, cupy):
            with pytest.raises(TypeError):
                xp.einsum(0, 0)

    def test_bad_argument(self):
        for xp in (numpy, cupy):
            with pytest.raises(TypeError):
                xp.einsum('', 0, bad_arg=0)

    def test_too_many_operands1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('', 0, 0)

    def test_too_many_operands2(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(
                    'i,j',
                    xp.array([0, 0]),
                    xp.array([0, 0]),
                    xp.array([0, 0]))

    def test_too_few_operands1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(',', 0)

    def test_too_many_dimension1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i', 0)

    def test_too_many_dimension2(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('ij', xp.array([0, 0]))

    def test_too_many_dimension3(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('ijk...->...', xp.arange(6).reshape(2, 3))

    def test_too_few_dimension(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i->i', xp.arange(6).reshape(2, 3))

    def test_invalid_char1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i%', xp.array([0, 0]))

    def test_invalid_char2(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('j$', xp.array([0, 0]))

    def test_invalid_char3(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i->&', xp.array([0, 0]))

    # output subscripts must appear in inumpy.t
    def test_invalid_output_subscripts1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i->ij', xp.array([0, 0]))

    # output subscripts may only be specified once
    def test_invalid_output_subscripts2(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('ij->jij', xp.array([[0, 0], [0, 0]]))

    # output subscripts must not incrudes comma
    def test_invalid_output_subscripts3(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('ij->i,j', xp.array([[0, 0], [0, 0]]))

    # dimensions much match when being collapsed
    def test_invalid_diagonal1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('ii', xp.arange(6).reshape(2, 3))

    def test_invalid_diagonal2(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('ii->', xp.arange(6).reshape(2, 3))

    def test_invalid_diagonal3(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('ii', xp.arange(3).reshape(1, 3))

    def test_dim_mismatch_char1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i,i', xp.arange(2), xp.arange(3))

    def test_dim_mismatch_ellipsis1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('...,...', xp.arange(2), xp.arange(3))

    def test_dim_mismatch_ellipsis2(self):
        for xp in (numpy, cupy):
            a = xp.arange(12).reshape(2, 3, 2)
            with pytest.raises(ValueError):
                xp.einsum('i...,...i', a, a)

    def test_dim_mismatch_ellipsis3(self):
        for xp in (numpy, cupy):
            a = xp.arange(12).reshape(2, 3, 2)
            with pytest.raises(ValueError):
                xp.einsum('...,...', a, a[:, :2])

    # invalid -> operator
    def test_invalid_arrow1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i-i', xp.array([0, 0]))

    def test_invalid_arrow2(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i>i', xp.array([0, 0]))

    def test_invalid_arrow3(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i->->i', xp.array([0, 0]))

    def test_invalid_arrow4(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum('i-', xp.array([0, 0]))


class TestListArgEinSumError(unittest.TestCase):

    def test_invalid_sub1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(xp.arange(2), [None])

    def test_invalid_sub2(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(xp.arange(2), [0], [1])

    def test_invalid_sub3(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(xp.arange(2), [Ellipsis, 0, Ellipsis])

    def test_dim_mismatch1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(xp.arange(2), [0], xp.arange(3), [0])

    def test_dim_mismatch2(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(xp.arange(2), [0], xp.arange(3), [0], [0])

    def test_dim_mismatch3(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(xp.arange(6).reshape(2, 3), [0, 0])

    def test_too_many_dims1(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(3, [0])

    def test_too_many_dims2(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(xp.arange(2), [0, 1])

    def test_too_many_dims3(self):
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                xp.einsum(xp.arange(6).reshape(2, 3), [Ellipsis, 0, 1, 2])


class TestEinSumUnaryOperationWithScalar(unittest.TestCase):
    @testing.for_all_dtypes()
    @testing.numpy_cupy_allclose()
    def test_scalar_int(self, xp, dtype):
        return xp.asarray(xp.einsum('->', 2, dtype=dtype))

    @testing.for_all_dtypes()
    @testing.numpy_cupy_allclose()
    def test_scalar_float(self, xp, dtype):
        return xp.asarray(xp.einsum('', 2.0, dtype=dtype))


class TestEinSumBinaryOperationWithScalar(unittest.TestCase):
    @testing.for_all_dtypes()
    @testing.numpy_cupy_allclose(contiguous_check=False)
    def test_scalar_1(self, xp, dtype):
        shape_a = (2,)
        a = testing.shaped_arange(shape_a, xp, dtype)
        return xp.asarray(xp.einsum(',i->', 3, a))

    @testing.for_all_dtypes()
    @testing.numpy_cupy_allclose(contiguous_check=False)
    def test_scalar_2(self, xp, dtype):
        shape_a = (2,)
        a = testing.shaped_arange(shape_a, xp, dtype)
        return xp.asarray(xp.einsum('i,->', a, 4))


@testing.parameterize(*([
    # memory constraint
    {'subscript': 'a,b,c->abc', 'opt': ('greedy', 0)},
    {'subscript': 'acdf,jbje,gihb,hfac', 'opt': ('greedy', 0)},
] + testing.product({'subscript': [
    # long paths
    'acdf,jbje,gihb,hfac,gfac,gifabc,hfac',
    'chd,bde,agbc,hiad,bdi,cgh,agdb',
    # edge cases
    'eb,cb,fb->cef',
    'dd,fb,be,cdb->cef',
    'bca,cdb,dbf,afc->',
    'dcc,fce,ea,dbf->ab',
    'a,ac,ab,ad,cd,bd,bc->',
], 'opt': ['greedy', 'optimal'],
})))
class TestEinSumLarge(unittest.TestCase):

    def setUp(self):
        chars = 'abcdefghij'
        sizes = numpy.array([2, 3, 4, 5, 4, 3, 2, 6, 5, 4, 3])
        size_dict = {}
        for size, char in zip(sizes, chars):
            size_dict[char] = size

        # Builds views based off initial operands
        string = self.subscript
        operands = [string]
        terms = string.split('->')[0].split(',')
        for term in terms:
            dims = [size_dict[x] for x in term]
            operands.append(numpy.ones(*dims))

        self.operands = operands

    @testing.numpy_cupy_allclose(contiguous_check=False)
    def test_einsum(self, xp):
        # TODO(kataoka): support memory efficient cupy.einsum
        with warnings.catch_warnings(record=True) as ws:
            # I hope there's no problem with np.einsum for these cases...
            out = xp.einsum(*self.operands, optimize=self.opt)
            if xp is not numpy and \
                    isinstance(self.opt, tuple):  # with memory limit
                for w in ws:
                    self.assertIn('memory', str(w.message))
            else:
                self.assertEqual(len(ws), 0)
        return out
