import pytest

import dpnp as inp

import numpy
from numpy.testing import (
    assert_array_equal
)


@pytest.mark.parametrize("lhs", [[[-7, -6, -5, -4, -3, -2, -1], [0, 1, 2, 3, 4, 5, 6]], [-3, -2, -1, 0, 1, 2, 3], 0])
@pytest.mark.parametrize("rhs", [[[0, 1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12, 13]], [0, 1, 2, 3, 4, 5, 6], 3])
@pytest.mark.parametrize("dtype", [inp.bool, inp.int32, inp.int64])
class TestBitwise:

    @staticmethod
    def array_or_scalar(xp, data, dtype=None):
        if numpy.isscalar(data):
            if dtype == inp.bool:
                return numpy.dtype(dtype).type(data)
            return data

        return xp.array(data, dtype=dtype)

    def _test_unary_int(self, name, data, dtype):
        dp_a = self.array_or_scalar(inp, data, dtype=dtype)
        result = getattr(inp, name)(dp_a)

        np_a = self.array_or_scalar(numpy, data, dtype=dtype)
        expected = getattr(numpy, name)(np_a)

        assert_array_equal(result, expected)
        return (dp_a, np_a)

    def _test_binary_int(self, name, lhs, rhs, dtype):
        if name in ('left_shift', 'right_shift') and dtype == inp.bool:
            pytest.skip("A shift operation isn't implemented for bool type")
        elif numpy.isscalar(lhs) and numpy.isscalar(rhs):
            pytest.skip("Both inputs can't be scalars")

        dp_a = self.array_or_scalar(inp, lhs, dtype=dtype)
        dp_b = self.array_or_scalar(inp, rhs, dtype=dtype)
        result = getattr(inp, name)(dp_a, dp_b)

        np_a = self.array_or_scalar(numpy, lhs, dtype=dtype)
        np_b = self.array_or_scalar(numpy, rhs, dtype=dtype)
        expected = getattr(numpy, name)(np_a, np_b)

        assert_array_equal(result, expected)
        return (dp_a, dp_b, np_a, np_b)

    def test_bitwise_and(self, lhs, rhs, dtype):
        dp_a, dp_b, np_a, np_b = self._test_binary_int('bitwise_and', lhs, rhs, dtype)
        assert_array_equal(dp_a & dp_b, np_a & np_b)

        if not (inp.isscalar(dp_a) or inp.isscalar(dp_b)) and dp_a.shape == dp_b.shape:
            dp_a &= dp_b
            np_a &= np_b
            assert_array_equal(dp_a, np_a)

    def test_bitwise_or(self, lhs, rhs, dtype):
        dp_a, dp_b, np_a, np_b = self._test_binary_int('bitwise_or', lhs, rhs, dtype)
        assert_array_equal(dp_a | dp_b, np_a | np_b)

        if not (inp.isscalar(dp_a) or inp.isscalar(dp_b)) and dp_a.shape == dp_b.shape:
            dp_a |= dp_b
            np_a |= np_b
            assert_array_equal(dp_a, np_a)

    def test_bitwise_xor(self, lhs, rhs, dtype):
        dp_a, dp_b, np_a, np_b = self._test_binary_int('bitwise_xor', lhs, rhs, dtype)
        assert_array_equal(dp_a ^ dp_b, np_a ^ np_b)

        if not (inp.isscalar(dp_a) or inp.isscalar(dp_b)) and dp_a.shape == dp_b.shape:
            dp_a ^= dp_b
            np_a ^= np_b
            assert_array_equal(dp_a, np_a)

    def test_invert(self, lhs, rhs, dtype):
        dp_a, np_a = self._test_unary_int('invert', lhs, dtype)
        assert_array_equal(~dp_a, ~np_a)

    def test_left_shift(self, lhs, rhs, dtype):
        dp_a, dp_b, np_a, np_b = self._test_binary_int('left_shift', lhs, rhs, dtype)
        assert_array_equal(dp_a << dp_b, np_a << np_b)

        if not (inp.isscalar(dp_a) or inp.isscalar(dp_b)) and dp_a.shape == dp_b.shape:
            dp_a <<= dp_b
            np_a <<= np_b
            assert_array_equal(dp_a, np_a)

    def test_right_shift(self, lhs, rhs, dtype):
        dp_a, dp_b, np_a, np_b = self._test_binary_int('right_shift', lhs, rhs, dtype)
        assert_array_equal(dp_a >> dp_b, np_a >> np_b)

        if not (inp.isscalar(dp_a) or inp.isscalar(dp_b)) and dp_a.shape == dp_b.shape:
            dp_a >>= dp_b
            np_a >>= np_b
            assert_array_equal(dp_a, np_a)
