import numpy
import pytest

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


def astype_without_warning(x, dtype, *args, **kwargs):
    dtype = numpy.dtype(dtype)
    if x.dtype.kind == "c" and dtype.kind not in ["b", "c"]:
        with testing.assert_warns(numpy.ComplexWarning):
            return x.astype(dtype, *args, **kwargs)
    else:
        return x.astype(dtype, *args, **kwargs)


def get_strides(xp, a):
    if xp is numpy:
        return tuple(el // a.itemsize for el in a.strides)
    return a.strides


@pytest.mark.skip("'dpnp_array' object has no attribute 'view' yet")
class TestView:
    @testing.numpy_cupy_array_equal()
    def test_view(self, xp):
        a = testing.shaped_arange((4,), xp, dtype=numpy.float32)
        b = a.view(dtype=numpy.int32)
        b[:] = 0
        return a

    @testing.for_dtypes([numpy.int16, numpy.int64])
    @testing.numpy_cupy_array_equal()
    def test_view_itemsize(self, xp, dtype):
        a = testing.shaped_arange((4,), xp, dtype=numpy.int32)
        b = a.view(dtype=dtype)
        return b

    @testing.numpy_cupy_array_equal()
    def test_view_0d(self, xp):
        a = xp.array(1.5, dtype=numpy.float32)
        return a.view(dtype=numpy.int32)

    @testing.for_dtypes([numpy.int16, numpy.int64])
    def test_view_0d_raise(self, dtype):
        for xp in (numpy, cupy):
            a = xp.array(3, dtype=numpy.int32)
            with pytest.raises(ValueError):
                a.view(dtype=dtype)

    @testing.for_dtypes([numpy.int16, numpy.int64])
    def test_view_non_contiguous_raise(self, dtype):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((2, 2, 2), xp, dtype=numpy.int32)
            a = a.transpose(0, 2, 1)
            with pytest.raises(ValueError):
                a.view(dtype=dtype)

    @testing.for_dtypes([numpy.int16, numpy.int64])
    @testing.with_requires("numpy>=1.23")
    def test_view_f_contiguous(self, dtype):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((2, 2, 2), xp, dtype=numpy.float32)
            a = a.T
            with pytest.raises(ValueError):
                a.view(dtype=dtype)

    def test_view_assert_divisible(self):
        for xp in (numpy, cupy):
            a = testing.shaped_arange((3,), xp, dtype=numpy.int32)
            with pytest.raises(ValueError):
                a.view(dtype=numpy.int64)

    @testing.for_dtypes([numpy.float32, numpy.float64])
    @testing.numpy_cupy_array_equal(strides_check=True)
    def test_view_relaxed_contiguous(self, xp, dtype):
        a = testing.shaped_arange((1, 3, 5), xp, dtype=dtype)
        a = xp.moveaxis(a, 0, 2)  # (3, 5, 1)
        b = a.view(dtype=numpy.int32)
        return b

    @pytest.mark.parametrize(
        ("order", "shape"),
        [
            ("C", (3,)),
            ("C", (3, 5)),
            ("C", (0,)),
            ("C", (1, 3)),
            ("C", (3, 1)),
        ],
        ids=str,
    )
    @testing.numpy_cupy_equal()
    def test_view_flags_smaller(self, xp, order, shape):
        a = xp.zeros(shape, numpy.int32, order)
        b = a.view(numpy.int16)
        return b.flags.c_contiguous, b.flags.f_contiguous, b.flags.owndata

    @pytest.mark.parametrize(
        ("order", "shape"),
        [
            ("F", (3, 5)),
        ],
        ids=str,
    )
    @testing.with_requires("numpy>=1.23")
    def test_view_flags_smaller_invalid(self, order, shape):
        for xp in (numpy, cupy):
            a = xp.zeros(shape, numpy.int32, order)
            with pytest.raises(ValueError):
                a.view(numpy.int16)

    @pytest.mark.parametrize(
        ("order", "shape"),
        [
            ("C", (6,)),
            ("C", (3, 10)),
            ("C", (0,)),
            ("C", (1, 6)),
            ("C", (3, 2)),
        ],
        ids=str,
    )
    @testing.numpy_cupy_equal()
    def test_view_flags_larger(self, xp, order, shape):
        a = xp.zeros(shape, numpy.int16, order)
        b = a.view(numpy.int32)
        return b.flags.c_contiguous, b.flags.f_contiguous, b.flags.owndata

    @pytest.mark.parametrize(
        ("order", "shape"),
        [
            ("F", (6, 5)),
            ("F", (2, 3)),
        ],
        ids=str,
    )
    @testing.with_requires("numpy>=1.23")
    def test_view_flags_larger_invalid(self, order, shape):
        for xp in (numpy, cupy):
            a = xp.zeros(shape, numpy.int16, order)
            with pytest.raises(ValueError):
                a.view(numpy.int32)

    @testing.with_requires("numpy>=1.23")
    @testing.numpy_cupy_array_equal()
    def test_view_smaller_dtype_multiple(self, xp):
        # x is non-contiguous
        x = xp.arange(10, dtype=xp.int32)[::2]
        with pytest.raises(ValueError):
            x.view(xp.int16)
        return x[:, xp.newaxis].view(xp.int16)

    @testing.with_requires("numpy>=1.23")
    @testing.numpy_cupy_array_equal()
    def test_view_smaller_dtype_multiple2(self, xp):
        # x is non-contiguous, and stride[-1] != 0
        x = xp.ones((3, 4), xp.int32)[:, :1:2]
        return x.view(xp.int16)

    @testing.with_requires("numpy>=1.23")
    @testing.numpy_cupy_array_equal()
    def test_view_larger_dtype_multiple(self, xp):
        # x is non-contiguous in the first dimension, contiguous in the last
        x = xp.arange(20, dtype=xp.int16).reshape(10, 2)[::2, :]
        return x.view(xp.int32)

    @testing.with_requires("numpy>=1.23")
    @testing.numpy_cupy_array_equal()
    def test_view_non_c_contiguous(self, xp):
        # x is contiguous in axis=-1, but not C-contiguous in other axes
        x = (
            xp.arange(2 * 3 * 4, dtype=xp.int8)
            .reshape(2, 3, 4)
            .transpose(1, 0, 2)
        )
        return x.view(xp.int16)

    @testing.numpy_cupy_array_equal()
    def test_view_larger_dtype_zero_sized(self, xp):
        x = xp.ones((3, 20), xp.int16)[:0, ::2]
        return x.view(xp.int32)


class TestArrayCopy:
    @testing.for_orders("CF")
    @testing.for_dtypes(
        [numpy.int16, numpy.int64, numpy.float16, numpy.float64]
    )
    @testing.numpy_cupy_array_equal()
    def test_isinstance_numpy_copy(self, xp, dtype, order):
        a = numpy.arange(100, dtype=dtype).reshape(10, 10, order=order)
        b = xp.empty(a.shape, dtype=dtype, order=order)
        b[:] = a
        return b

    @pytest.mark.skip("Doesn't raise ValueError in numpy")
    def test_isinstance_numpy_copy_wrong_dtype(self):
        a = numpy.arange(100, dtype=numpy.float32).reshape(10, 10)
        b = cupy.empty(a.shape, dtype=numpy.int32)
        with pytest.raises(ValueError):
            b[:] = a

    def test_isinstance_numpy_copy_wrong_shape(self):
        for xp in (numpy, cupy):
            a = numpy.arange(100, dtype=numpy.float32).reshape(10, 10)
            b = cupy.empty(100, dtype=a.dtype)
            with pytest.raises(ValueError):
                b[:] = a

    @testing.numpy_cupy_array_equal()
    def test_isinstance_numpy_copy_not_slice(self, xp):
        a = xp.arange(5, dtype=numpy.float32)
        a[a < 3] = 0
        return a

    @pytest.mark.skip("copy from host to device is allowed")
    def test_copy_host_to_device_view(self):
        dev = cupy.empty((10, 10), dtype=numpy.float32)[2:5, 1:8]
        host = numpy.arange(3 * 7, dtype=numpy.float32).reshape(3, 7)
        with pytest.raises(ValueError):
            dev[:] = host


class TestArrayFlatten:
    @testing.numpy_cupy_array_equal()
    def test_flatten(self, xp):
        a = testing.shaped_arange((2, 3, 4), xp)
        return a.flatten()

    @testing.numpy_cupy_array_equal()
    def test_flatten_copied(self, xp):
        a = testing.shaped_arange((4,), xp)
        b = a.flatten()
        a[:] = 1
        return b

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

    @testing.for_orders("CFAK")
    @testing.numpy_cupy_array_equal()
    def test_flatten_order(self, xp, order):
        a = testing.shaped_arange((2, 3, 4), xp)
        return a.flatten(order)

    @testing.for_orders("CFAK")
    @testing.numpy_cupy_array_equal()
    def test_flatten_order_copied(self, xp, order):
        a = testing.shaped_arange((4,), xp)
        b = a.flatten(order=order)
        a[:] = 1
        return b

    @testing.for_orders("CFAK")
    @testing.numpy_cupy_array_equal()
    def test_flatten_order_transposed(self, xp, order):
        a = testing.shaped_arange((2, 3, 4), xp).transpose(2, 0, 1)
        return a.flatten(order=order)


class TestArrayFill:
    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_fill(self, xp, dtype):
        a = testing.shaped_arange((2, 3, 4), xp, dtype)
        a.fill(1)
        return a

    @testing.with_requires("numpy>=1.24.0")
    @testing.for_all_dtypes_combination(("dtype1", "dtype2"))
    @testing.numpy_cupy_array_equal(accept_error=numpy.ComplexWarning)
    def test_fill_with_numpy_scalar_ndarray(self, xp, dtype1, dtype2):
        a = testing.shaped_arange((2, 3, 4), xp, dtype1)
        a.fill(numpy.ones((), dtype=dtype2))
        return a

    @testing.with_requires("numpy>=1.24.0")
    @testing.for_all_dtypes_combination(("dtype1", "dtype2"))
    @testing.numpy_cupy_array_equal(accept_error=numpy.ComplexWarning)
    def test_fill_with_cupy_scalar_ndarray(self, xp, dtype1, dtype2):
        a = testing.shaped_arange((2, 3, 4), xp, dtype1)
        b = xp.ones((), dtype=dtype2)
        a.fill(b)
        return a

    @pytest.mark.skip(
        "it's allowed to broadcast dpnp array while filling, no exception then"
    )
    @testing.for_all_dtypes()
    def test_fill_with_nonscalar_ndarray(self, dtype):
        a = testing.shaped_arange((2, 3, 4), cupy, dtype)
        for xp in (numpy, cupy):
            with pytest.raises(ValueError):
                a.fill(xp.ones((1,), dtype=dtype))

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_transposed_fill(self, xp, dtype):
        a = testing.shaped_arange((2, 3, 4), xp, dtype)
        b = a.transpose(2, 0, 1)
        b.fill(1)
        return b


class TestArrayAsType:
    @testing.for_orders(["C", "F", "A", "K", None])
    @testing.for_all_dtypes_combination(("src_dtype", "dst_dtype"))
    @testing.numpy_cupy_array_equal()
    def test_astype(self, xp, src_dtype, dst_dtype, order):
        a = testing.shaped_arange((2, 3, 4), xp, src_dtype)
        return astype_without_warning(a, dst_dtype, order=order)

    @testing.for_orders(["C", "F", "A", "K", None])
    @testing.for_all_dtypes_combination(("src_dtype", "dst_dtype"))
    @testing.numpy_cupy_array_equal()
    def test_astype_empty(self, xp, src_dtype, dst_dtype, order):
        a = testing.shaped_arange((2, 0, 4), xp, src_dtype)
        return astype_without_warning(a, dst_dtype, order=order)

    @testing.for_orders("CFAK")
    @testing.for_all_dtypes_combination(("src_dtype", "dst_dtype"))
    def test_astype_type(self, src_dtype, dst_dtype, order):
        a = testing.shaped_arange((2, 3, 4), cupy, src_dtype)
        b = astype_without_warning(a, dst_dtype, order=order)
        a_cpu = testing.shaped_arange((2, 3, 4), numpy, src_dtype)
        b_cpu = astype_without_warning(a_cpu, dst_dtype, order=order)
        assert b.dtype.type == b_cpu.dtype.type

    @testing.for_orders("CAK")
    @testing.for_all_dtypes()
    def test_astype_type_c_contiguous_no_copy(self, dtype, order):
        a = testing.shaped_arange((2, 3, 4), cupy, dtype)
        b = a.astype(dtype, order=order, copy=False)
        assert b is a

    @testing.for_orders("FAK")
    @testing.for_all_dtypes()
    def test_astype_type_f_contiguous_no_copy(self, dtype, order):
        a = testing.shaped_arange((2, 3, 4), cupy, dtype)
        a = cupy.asfortranarray(a)
        b = a.astype(dtype, order=order, copy=False)
        assert b is a

    @testing.for_all_dtypes_combination(("src_dtype", "dst_dtype"))
    @testing.numpy_cupy_equal()
    def test_astype_strides(self, xp, src_dtype, dst_dtype):
        src = testing.shaped_arange((1, 2, 3), xp, dtype=src_dtype)
        dst = astype_without_warning(src, dst_dtype, order="K")
        return get_strides(xp, dst)

    @testing.for_all_dtypes_combination(("src_dtype", "dst_dtype"))
    @testing.numpy_cupy_equal()
    def test_astype_strides_negative(self, xp, src_dtype, dst_dtype):
        src = testing.shaped_arange((2, 3), xp, dtype=src_dtype)
        src = src[::-1, :]
        dst = astype_without_warning(src, dst_dtype, order="K")
        return tuple(abs(x) for x in get_strides(xp, dst))

    @testing.for_all_dtypes_combination(("src_dtype", "dst_dtype"))
    @testing.numpy_cupy_equal()
    def test_astype_strides_swapped(self, xp, src_dtype, dst_dtype):
        src = testing.shaped_arange((2, 3, 4), xp, dtype=src_dtype)
        src = xp.swapaxes(src, 1, 0)
        dst = astype_without_warning(src, dst_dtype, order="K")
        return get_strides(xp, dst)

    @testing.for_all_dtypes_combination(("src_dtype", "dst_dtype"))
    @testing.numpy_cupy_equal()
    def test_astype_strides_broadcast(self, xp, src_dtype, dst_dtype):
        src1 = testing.shaped_arange((2, 3, 2), xp, dtype=src_dtype)
        src2 = testing.shaped_arange((2,), xp, dtype=src_dtype)
        src, _ = xp.broadcast_arrays(src1, src2)
        dst = astype_without_warning(src, dst_dtype, order="K")
        return get_strides(xp, dst)

    @pytest.mark.skip("'dpnp_array' object has no attribute 'view' yet")
    @testing.numpy_cupy_array_equal()
    def test_astype_boolean_view(self, xp):
        # See #4354
        a = xp.array([0, 1, 2], dtype=numpy.int8).view(dtype=numpy.bool_)
        return a.astype(numpy.int8)


class TestArrayDiagonal:
    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_diagonal1(self, xp, dtype):
        a = testing.shaped_arange((3, 4, 5), xp, dtype)
        return a.diagonal(1, 2, 0)

    @testing.for_all_dtypes()
    @testing.numpy_cupy_array_equal()
    def test_diagonal2(self, xp, dtype):
        a = testing.shaped_arange((3, 4, 5), xp, dtype)
        return a.diagonal(-1, 2, 0)


@testing.parameterize(
    {"src_order": "C"},
    {"src_order": "F"},
)
class TestNumPyArrayCopyView:
    @testing.for_orders("CF")
    @testing.for_dtypes(
        [numpy.int16, numpy.int64, numpy.float16, numpy.float64]
    )
    @testing.numpy_cupy_array_equal()
    def test_isinstance_numpy_view_copy_f(self, xp, dtype, order):
        a = numpy.arange(100, dtype=dtype).reshape(10, 10, order=self.src_order)
        a = a[2:5, 1:8]
        b = xp.empty(a.shape, dtype=dtype, order=order)
        b[:] = a
        return b


class C_cp(cupy.ndarray):
    def __new__(cls, *args, info=None, **kwargs):
        obj = super().__new__(cls, *args, **kwargs)
        obj.info = info
        return obj

    def __array_finalize__(self, obj):
        if obj is None:
            return
        self.info = getattr(obj, "info", None)


class C_np(numpy.ndarray):
    def __new__(cls, *args, info=None, **kwargs):
        obj = super().__new__(cls, *args, **kwargs)
        obj.info = info
        return obj

    def __array_finalize__(self, obj):
        if obj is None:
            return
        self.info = getattr(obj, "info", None)


@pytest.mark.skip("'dpnp_array' object has no attribute 'view' yet")
class TestSubclassArrayView:
    def test_view_casting(self):
        for xp, C in [(numpy, C_np), (cupy, C_cp)]:
            a = xp.arange(5, dtype="i").view("F")
            assert type(a) is xp.ndarray
            assert a.dtype == xp.float32

            a = xp.arange(5, dtype="i").view(dtype="F")
            assert type(a) is xp.ndarray
            assert a.dtype == xp.float32

            with pytest.raises(TypeError):
                xp.arange(5, dtype="i").view("F", dtype="F")

            a = xp.arange(5, dtype="i").view(C)
            assert type(a) is C
            assert a.dtype == xp.int32
            assert a.info is None

            a = xp.arange(5, dtype="i").view(type=C)
            assert type(a) is C
            assert a.dtype == xp.int32
            assert a.info is None

            # When an instance of ndarray's subclass is supplied to `dtype`,
            # view() interprets it as if it is supplied to `type`
            a = xp.arange(5, dtype="i").view(dtype=C)
            assert type(a) is C
            assert a.dtype == xp.int32
            assert a.info is None

            with pytest.raises(TypeError):
                xp.arange(5).view("F", C, type=C)

        with pytest.raises(ValueError):
            cupy.arange(5).view(type=numpy.ndarray)
