import unittest

import numpy
import pytest

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


class TestDims(unittest.TestCase):
    def check_atleast(self, func, xp):
        a = testing.shaped_arange((), xp)
        b = testing.shaped_arange((2,), xp)
        c = testing.shaped_arange((2, 2), xp)
        d = testing.shaped_arange((4, 3, 2), xp)
        e = 1
        f = numpy.float32(1)
        return func(a, b, c, d, e, f)

    @pytest.mark.skip(reason="Scalar input is not supported")
    @testing.numpy_cupy_array_equal()
    def test_atleast_1d1(self, xp):
        return self.check_atleast(xp.atleast_1d, xp)

    @testing.numpy_cupy_array_equal()
    def test_atleast_1d2(self, xp):
        a = testing.shaped_arange((1, 3, 2), xp)
        return xp.atleast_1d(a)

    @pytest.mark.skip(reason="Scalar input is not supported")
    @testing.numpy_cupy_array_equal()
    def test_atleast_2d1(self, xp):
        return self.check_atleast(xp.atleast_2d, xp)

    @testing.numpy_cupy_array_equal()
    def test_atleast_2d2(self, xp):
        a = testing.shaped_arange((1, 3, 2), xp)
        return xp.atleast_2d(a)

    @pytest.mark.skip(reason="Scalar input is not supported")
    @testing.numpy_cupy_array_equal()
    def test_atleast_3d1(self, xp):
        return self.check_atleast(xp.atleast_3d, xp)

    @testing.numpy_cupy_array_equal()
    def test_atleast_3d2(self, xp):
        a = testing.shaped_arange((1, 3, 2), xp)
        return xp.atleast_3d(a)

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_broadcast_to(self, xp, dtype):
        # Note that broadcast_to is only supported on numpy>=1.10
        a = testing.shaped_arange((3, 1, 4), xp, dtype)
        b = xp.broadcast_to(a, (2, 3, 3, 4))
        return b

    @testing.numpy_cupy_array_equal()
    def test_broadcast_to_int(self, xp):
        # Broadcast 0-dim array to 1-dim.
        a = xp.array(10, dtype=xp.float32)
        b = xp.broadcast_to(a, 10)
        return b

    @testing.for_all_dtypes()
    def test_broadcast_to_fail(self, dtype):
        for xp in (numpy, cupy):
            # Note that broadcast_to is only supported on numpy>=1.10
            a = testing.shaped_arange((3, 1, 4), xp, dtype)
            with pytest.raises(ValueError):
                xp.broadcast_to(a, (1, 3, 4))

    @testing.for_all_dtypes()
    def test_broadcast_to_short_shape(self, dtype):
        for xp in (numpy, cupy):
            # Note that broadcast_to is only supported on numpy>=1.10
            a = testing.shaped_arange((1, 3, 4), xp, dtype)
            with pytest.raises(ValueError):
                xp.broadcast_to(a, (3, 4))

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_broadcast_to_numpy19(self, xp, dtype):
        # Note that broadcast_to is only supported on numpy>=1.10
        a = testing.shaped_arange((3, 1, 4), xp, dtype)
        if xp is cupy:
            b = xp.broadcast_to(a, (2, 3, 3, 4))
        else:
            dummy = xp.empty((2, 3, 3, 4))
            b, _ = xp.broadcast_arrays(a, dummy)
        return b

    @testing.for_all_dtypes()
    def test_broadcast_to_fail_numpy19(self, dtype):
        for xp in (numpy, cupy):
            # Note that broadcast_to is only supported on numpy>=1.10
            a = testing.shaped_arange((3, 1, 4), xp, dtype)
            with pytest.raises(ValueError):
                xp.broadcast_to(a, (1, 3, 4))

    @testing.for_all_dtypes()
    def test_broadcast_to_short_shape_numpy19(self, dtype):
        for xp in (numpy, cupy):
            # Note that broadcast_to is only supported on numpy>=1.10
            a = testing.shaped_arange((1, 3, 4), xp, dtype)
            with pytest.raises(ValueError):
                xp.broadcast_to(a, (3, 4))

    @testing.numpy_cupy_array_equal()
    def test_expand_dims0(self, xp):
        a = testing.shaped_arange((2, 3), xp)
        return xp.expand_dims(a, 0)

    @testing.numpy_cupy_array_equal()
    def test_expand_dims1(self, xp):
        a = testing.shaped_arange((2, 3), xp)
        return xp.expand_dims(a, 1)

    @testing.numpy_cupy_array_equal()
    def test_expand_dims2(self, xp):
        a = testing.shaped_arange((2, 3), xp)
        return xp.expand_dims(a, 2)

    @testing.numpy_cupy_array_equal()
    def test_expand_dims_negative1(self, xp):
        a = testing.shaped_arange((2, 3), xp)
        return xp.expand_dims(a, -2)

    @testing.with_requires("numpy>=1.18")
    def test_expand_dims_negative2(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((2, 3), xp)
            with pytest.raises(numpy.AxisError):
                xp.expand_dims(a, -4)

    @testing.with_requires("numpy>=1.18")
    @testing.numpy_cupy_array_equal()
    def test_expand_dims_tuple_axis(self, xp):
        a = testing.shaped_arange((2, 2, 2), xp)
        return [
            xp.expand_dims(a, axis)
            for axis in [
                (0, 1, 2),
                (0, -1, -2),
                (0, 3, 5),
                (0, -3, -5),
                (),
                (1,),
            ]
        ]

    @testing.with_requires("numpy>=1.18")
    def test_expand_dims_out_of_range(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((2, 2, 2), xp)
            for axis in [(1, -6), (1, 5)]:
                with pytest.raises(numpy.AxisError):
                    xp.expand_dims(a, axis)

    @testing.with_requires("numpy>=1.18")
    def test_expand_dims_repeated_axis(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((2, 2, 2), xp)
            with pytest.raises(ValueError):
                xp.expand_dims(a, (1, 1))

    @testing.numpy_cupy_array_equal()
    def test_squeeze1(self, xp):
        a = testing.shaped_arange((1, 2, 1, 3, 1, 1, 4, 1), xp)
        return a.squeeze()

    @testing.numpy_cupy_array_equal()
    def test_squeeze2(self, xp):
        a = testing.shaped_arange((2, 3, 4), xp)
        return a.squeeze()

    @testing.numpy_cupy_array_equal()
    def test_squeeze_int_axis1(self, xp):
        a = testing.shaped_arange((1, 2, 1, 3, 1, 1, 4, 1), xp)
        return a.squeeze(axis=2)

    @testing.numpy_cupy_array_equal()
    def test_squeeze_int_axis2(self, xp):
        a = testing.shaped_arange((1, 2, 1, 3, 1, 1, 4, 1), xp)
        return a.squeeze(axis=-3)

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

    def test_squeeze_int_axis_failure2(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((1, 2, 1, 3, 1, 1, 4, 1), xp)
            with pytest.raises(numpy.AxisError):
                a.squeeze(axis=-9)

    @testing.numpy_cupy_array_equal()
    def test_squeeze_tuple_axis1(self, xp):
        a = testing.shaped_arange((1, 2, 1, 3, 1, 1, 4, 1), xp)
        return a.squeeze(axis=(2, 4))

    @testing.numpy_cupy_array_equal()
    def test_squeeze_tuple_axis2(self, xp):
        a = testing.shaped_arange((1, 2, 1, 3, 1, 1, 4, 1), xp)
        return a.squeeze(axis=(-4, -3))

    @testing.numpy_cupy_array_equal()
    def test_squeeze_tuple_axis3(self, xp):
        a = testing.shaped_arange((1, 2, 1, 3, 1, 1, 4, 1), xp)
        return a.squeeze(axis=(4, 2))

    @testing.numpy_cupy_array_equal()
    def test_squeeze_tuple_axis4(self, xp):
        a = testing.shaped_arange((1, 2, 1, 3, 1, 1, 4, 1), xp)
        return a.squeeze(axis=())

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

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

    def test_squeeze_tuple_axis_failure3(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((1, 2, 1, 3, 1, 1, 4, 1), xp)
            with pytest.raises(numpy.AxisError):
                a.squeeze(axis=(-9,))

    @testing.numpy_cupy_array_equal()
    def test_squeeze_scalar1(self, xp):
        a = testing.shaped_arange((), xp)
        return a.squeeze(axis=0)

    @testing.numpy_cupy_array_equal()
    def test_squeeze_scalar2(self, xp):
        a = testing.shaped_arange((), xp)
        return a.squeeze(axis=-1)

    def test_squeeze_scalar_failure1(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((), xp)
            with pytest.raises(ValueError):
                a.squeeze(axis=-2)

    def test_squeeze_scalar_failure2(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((), xp)
            with pytest.raises(ValueError):
                a.squeeze(axis=1)

    def test_squeeze_scalar_failure3(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((), xp)
            with pytest.raises(numpy.AxisError):
                a.squeeze(axis=-2)

    def test_squeeze_scalar_failure4(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((), cupy)
            with pytest.raises(numpy.AxisError):
                a.squeeze(axis=1)

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

    @testing.numpy_cupy_array_equal()
    def test_external_squeeze(self, xp):
        a = testing.shaped_arange((1, 2, 1, 3, 1, 1, 4, 1), xp)
        return xp.squeeze(a)


@testing.parameterize(
    {"shapes": [(), ()]},
    {"shapes": [(0,), (0,)]},
    {"shapes": [(1,), (1,)]},
    {"shapes": [(2,), (2,)]},
    {"shapes": [(0,), (1,)]},
    {"shapes": [(2, 3), (1, 3)]},
    {"shapes": [(2, 1, 3, 4), (3, 1, 4)]},
    {"shapes": [(4, 3, 2, 3), (2, 3)]},
    {"shapes": [(2, 0, 1, 1, 3), (2, 1, 0, 0, 3)]},
    {"shapes": [(0, 1, 1, 3), (2, 1, 0, 0, 3)]},
    {"shapes": [(0, 1, 1, 0, 3), (5, 2, 0, 1, 0, 0, 3), (2, 1, 0, 0, 0, 3)]},
)
class TestBroadcast(unittest.TestCase):
    def _broadcast(self, xp, dtype, shapes):
        arrays = [testing.shaped_arange(s, xp, dtype) for s in shapes]
        return xp.broadcast(*arrays)

    @testing.for_all_dtypes()
    def test_broadcast(self, dtype):
        broadcast_np = self._broadcast(numpy, dtype, self.shapes)
        broadcast_cp = self._broadcast(cupy, dtype, self.shapes)
        assert broadcast_np.shape == broadcast_cp.shape
        assert broadcast_np.size == broadcast_cp.size
        assert broadcast_np.nd == broadcast_cp.nd

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_broadcast_arrays(self, xp, dtype):
        arrays = [testing.shaped_arange(s, xp, dtype) for s in self.shapes]
        return xp.broadcast_arrays(*arrays)


@testing.parameterize(
    {"shapes": [(3,), (2,)]},
    {
        "shapes": [
            (3, 2),
            (
                2,
                3,
            ),
        ]
    },
    {
        "shapes": [
            (3, 2),
            (
                3,
                4,
            ),
        ]
    },
    {"shapes": [(0,), (2,)]},
)
class TestInvalidBroadcast(unittest.TestCase):
    @testing.for_all_dtypes()
    def test_invalid_broadcast(self, dtype):
        for xp in (numpy, cupy):
            arrays = [testing.shaped_arange(s, xp, dtype) for s in self.shapes]
            with pytest.raises(ValueError):
                xp.broadcast(*arrays)

    @testing.for_all_dtypes()
    def test_invalid_broadcast_arrays(self, dtype):
        for xp in (numpy, cupy):
            arrays = [testing.shaped_arange(s, xp, dtype) for s in self.shapes]
            with pytest.raises(ValueError):
                xp.broadcast_arrays(*arrays)
