import unittest

import numpy
import pytest

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


@testing.parameterize(
    *testing.product(
        {
            "shape": [(2, 3), (), (4,)],
        }
    )
)
class TestShape(unittest.TestCase):
    def test_shape(self):
        shape = self.shape
        for xp in (numpy, cupy):
            a = testing.shaped_arange(shape, xp)
            assert cupy.shape(a) == shape

    def test_shape_list(self):
        shape = self.shape
        a = testing.shaped_arange(shape, numpy)
        a = a.tolist()
        assert cupy.shape(a) == shape


class TestReshape(unittest.TestCase):
    # order = 'A' is out of support currently
    _supported_orders = "CF"

    def test_reshape_shapes(self):
        def func(xp):
            a = testing.shaped_arange((1, 1, 1, 2, 2), xp)
            return a.shape

        assert func(numpy) == func(cupy)

    def test_reshape2(self):
        def func(xp):
            a = xp.zeros((8,), dtype=xp.float32)
            return a.reshape((1, 1, 1, 4, 1, 2)).shape

        assert func(numpy) == func(cupy)

    @testing.for_orders(_supported_orders)
    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_nocopy_reshape(self, xp, dtype, order):
        a = xp.zeros((2, 3, 4), dtype=dtype)
        b = a.reshape(4, 3, 2, order=order)
        b[1] = 1
        return a

    @testing.for_orders(_supported_orders)
    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_nocopy_reshape_with_order(self, xp, dtype, order):
        a = xp.zeros((2, 3, 4), dtype=dtype)
        b = a.reshape(4, 3, 2, order=order)
        b[1] = 1
        return a

    @testing.for_orders(_supported_orders)
    @testing.numpy_cupy_array_equal()
    def test_transposed_reshape2(self, xp, order):
        a = testing.shaped_arange((2, 3, 4), xp).transpose(2, 0, 1)
        return a.reshape(2, 3, 4, order=order)

    @testing.for_orders(_supported_orders)
    @testing.numpy_cupy_array_equal()
    def test_reshape_with_unknown_dimension(self, xp, order):
        a = testing.shaped_arange((2, 3, 4), xp)
        return a.reshape(3, -1, order=order)

    def test_reshape_with_multiple_unknown_dimensions(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((2, 3, 4), xp)
            with pytest.raises(ValueError):
                a.reshape(3, -1, -1)

    def test_reshape_with_changed_arraysize(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((2, 3, 4), xp)
            with pytest.raises(ValueError):
                a.reshape(2, 4, 4)

    def test_reshape_invalid_order(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((2, 3, 4), xp)
            with pytest.raises(ValueError):
                a.reshape(2, 4, 4, order="K")

    def test_reshape_zerosize_invalid(self):
        for xp in (numpy, cupy):
            a = xp.zeros((0,))
            with pytest.raises(ValueError):
                a.reshape(())

    def test_reshape_zerosize_invalid_unknown(self):
        for xp in (numpy, cupy):
            a = xp.zeros((0,))
            with pytest.raises(ValueError):
                a.reshape((-1, 0))

    @pytest.mark.skip("array.base is not implemented")
    @testing.numpy_cupy_array_equal()
    def test_reshape_zerosize(self, xp):
        a = xp.zeros((0,))
        b = a.reshape((0,))
        assert b.base is a
        return b

    @pytest.mark.skip("array.base is not implemented")
    @testing.for_orders(_supported_orders)
    @testing.numpy_cupy_array_equal(strides_check=True)
    def test_reshape_zerosize2(self, xp, order):
        a = xp.zeros((2, 0, 3))
        b = a.reshape((5, 0, 4), order=order)
        assert b.base is a
        return b

    @testing.for_orders(_supported_orders)
    @testing.numpy_cupy_array_equal()
    def test_external_reshape(self, xp, order):
        a = xp.zeros((8,), dtype=xp.float32)
        return xp.reshape(a, (1, 1, 1, 4, 1, 2), order=order)

    def _test_ndim_limit(self, xp, ndim, dtype, order):
        idx = [1] * ndim
        idx[-1] = ndim
        a = xp.ones(ndim, dtype=dtype)
        a = a.reshape(idx, order=order)
        assert a.ndim == ndim
        return a

    @testing.for_orders(_supported_orders)
    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_ndim_limit1(self, xp, dtype, order):
        # from cupy/cupy#4193
        a = self._test_ndim_limit(xp, 32, dtype, order)
        return a

    @pytest.mark.skip("no max ndim limit for reshape in dpctl")
    @testing.for_orders(_supported_orders)
    @testing.for_all_dtypes()
    def test_ndim_limit2(self, dtype, order):
        # from cupy/cupy#4193
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                self._test_ndim_limit(xp, 33, dtype, order)


class TestRavel(unittest.TestCase):
    @testing.for_orders("CF")
    # order = 'A' is out of support currently
    @testing.numpy_cupy_array_equal()
    def test_ravel(self, xp, order):
        a = testing.shaped_arange((2, 3, 4), xp)
        a = a.transpose(2, 0, 1)
        return a.ravel(order)

    @testing.for_orders("CF")
    # order = 'A' is out of support currently
    @testing.numpy_cupy_array_equal()
    def test_ravel2(self, xp, order):
        a = testing.shaped_arange((2, 3, 4), xp)
        return a.ravel(order)

    @testing.for_orders("CF")
    # order = 'A' is out of support currently
    @testing.numpy_cupy_array_equal()
    def test_ravel3(self, xp, order):
        a = testing.shaped_arange((2, 3, 4), xp)
        a = xp.asfortranarray(a)
        return a.ravel(order)

    @testing.numpy_cupy_array_equal()
    def test_external_ravel(self, xp):
        a = testing.shaped_arange((2, 3, 4), xp)
        a = a.transpose(2, 0, 1)
        return xp.ravel(a)


@testing.parameterize(
    *testing.product(
        {
            "order_init": ["C", "F"],
            # order = 'A' is out of support currently
            # 'order_reshape': ['C', 'F', 'A', 'c', 'f', 'a'],
            "order_reshape": ["C", "F", "c", "f"],
            "shape_in_out": [
                ((2, 3), (1, 6, 1)),  # (shape_init, shape_final)
                ((6,), (2, 3)),
                ((3, 3, 3), (9, 3)),
            ],
        }
    )
)
class TestReshapeOrder(unittest.TestCase):
    def test_reshape_contiguity(self):
        shape_init, shape_final = self.shape_in_out

        a_cupy = testing.shaped_arange(shape_init, xp=cupy)
        a_cupy = cupy.asarray(a_cupy, order=self.order_init)
        b_cupy = a_cupy.reshape(shape_final, order=self.order_reshape)

        a_numpy = testing.shaped_arange(shape_init, xp=numpy)
        a_numpy = numpy.asarray(a_numpy, order=self.order_init)
        b_numpy = a_numpy.reshape(shape_final, order=self.order_reshape)

        assert b_cupy.flags.f_contiguous == b_numpy.flags.f_contiguous
        assert b_cupy.flags.c_contiguous == b_numpy.flags.c_contiguous

        testing.assert_array_equal(b_cupy, b_numpy)
