import unittest
from numba.tests.support import TestCase

import sys
import operator
import numpy as np
import numpy

from numba import jit, njit, typed
from numba.core import types, utils
from numba.core.errors import TypingError, LoweringError
from numba.core.types.functions import _header_lead
from numba.np.numpy_support import numpy_version
from numba.tests.support import tag, _32bit, captured_stdout


# deliberately imported twice for different use cases


PARALLEL_SUPPORTED = not _32bit

def comp_list(n):
    l = [i for i in range(n)]
    s = 0
    for i in l:
        s += i
    return s


class TestListComprehension(TestCase):

    def test_comp_list(self):
        pyfunc = comp_list
        cfunc = njit((types.intp,))(pyfunc)
        self.assertEqual(cfunc(5), pyfunc(5))
        self.assertEqual(cfunc(0), pyfunc(0))
        self.assertEqual(cfunc(-1), pyfunc(-1))

    def test_bulk_use_cases(self):
        """ Tests the large number of use cases defined below """

        # jitted function used in some tests
        @jit(nopython=True)
        def fib3(n):
            if n < 2:
                return n
            return fib3(n - 1) + fib3(n - 2)

        def list1(x):
            """ Test basic list comprehension """
            return [i for i in range(1, len(x) - 1)]

        def list2(x):
            """ Test conditional list comprehension """
            return [y for y in x if y < 2]

        def list3(x):
            """ Test ternary list comprehension """
            return [y if y < 2 else -1 for y in x]

        def list4(x):
            """ Test list comprehension to np.array ctor """
            return np.array([1, 2, 3])

        # expected fail, unsupported type in sequence
        def list5(x):
            """ Test nested list comprehension to np.array ctor """
            return np.array([np.array([z for z in x]) for y in x])

        def list6(x):
            """ Test use of inner function in list comprehension """
            def inner(x):
                return x + 1
            return [inner(z) for z in x]

        def list7(x):
            """ Test use of closure in list comprehension """
            y = 3

            def inner(x):
                return x + y
            return [inner(z) for z in x]

        def list8(x):
            """ Test use of list comprehension as arg to inner function """
            l = [z + 1 for z in x]

            def inner(x):
                return x[0] + 1
            q = inner(l)
            return q

        def list9(x):
            """ Test use of list comprehension access in closure """
            l = [z + 1 for z in x]

            def inner(x):
                return x[0] + l[1]
            return inner(x)

        def list10(x):
            """ Test use of list comprehension access in closure and as arg """
            l = [z + 1 for z in x]

            def inner(x):
                return [y + l[0] for y in x]
            return inner(l)

        def list11(x):
            """ Test scalar array construction in list comprehension """
            l = [np.array(z) for z in x]
            return l

        def list12(x):
            """ Test scalar type conversion construction in list comprehension """
            l = [np.float64(z) for z in x]
            return l

        def list13(x):
            """ Test use of explicit numpy scalar ctor reference in list comprehension """
            l = [numpy.float64(z) for z in x]
            return l

        def list14(x):
            """ Test use of python scalar ctor reference in list comprehension """
            l = [float(z) for z in x]
            return l

        def list15(x):
            """ Test use of python scalar ctor reference in list comprehension followed by np array construction from the list"""
            l = [float(z) for z in x]
            return np.array(l)

        def list16(x):
            """ Test type unification from np array ctors consuming list comprehension """
            l1 = [float(z) for z in x]
            l2 = [z for z in x]
            ze = np.array(l1)
            oe = np.array(l2)
            return ze + oe

        def list17(x):
            """ Test complex list comprehension including math calls """
            return [(a, b, c)
                    for a in x for b in x for c in x if np.sqrt(a**2 + b**2) == c]

        _OUTER_SCOPE_VAR = 9

        def list18(x):
            """ Test loop list with outer scope var as conditional"""
            z = []
            for i in x:
                if i < _OUTER_SCOPE_VAR:
                    z.append(i)
            return z

        _OUTER_SCOPE_VAR = 9

        def list19(x):
            """ Test list comprehension with outer scope as conditional"""
            return [i for i in x if i < _OUTER_SCOPE_VAR]

        def list20(x):
            """ Test return empty list """
            return [i for i in x if i == -1000]

        def list21(x):
            """ Test call a jitted function in a list comprehension """
            return [fib3(i) for i in x]

        def list22(x):
            """ Test create two lists comprehensions and a third walking the first two """
            a = [y - 1 for y in x]
            b = [y + 1 for y in x]
            return [x for x in a for y in b if x == y]

        def list23(x):
            """ Test operation on comprehension generated list """
            z = [y for y in x]
            z.append(1)
            return z

        def list24(x):
            """ Test type promotion """
            z = [float(y) if y > 3 else y for y in x]
            return z

        def list25(x):
            # See issue #6260. Old style inline_closure_call uses get_ir_of_code
            # for the closure->IR transform, without SSA there's multiply
            # defined labels, the unary negation is self referent and DCE runs
            # eliminating the duplicated labels.
            included = np.array([1, 2, 6, 8])
            not_included = [i for i in range(10) if i not in list(included)]
            return not_included

        # functions to test that are expected to pass
        f = [list1, list2, list3, list4,
             list6, list7, list8, list9, list10, list11,
             list12, list13, list14, list15,
             list16, list17, list18, list19, list20,
             list21, list22, list23, list24, list25]

        var = [1, 2, 3, 4, 5]
        for ref in f:
            try:
                cfunc = jit(nopython=True)(ref)
                self.assertEqual(cfunc(var), ref(var))
            except ValueError:  # likely np array returned
                try:
                    np.testing.assert_allclose(cfunc(var), ref(var))
                except Exception:
                    raise

        # test functions that are expected to fail
        with self.assertRaises(TypingError) as raises:
            cfunc = jit(nopython=True)(list5)
            cfunc(var)
        # TODO: we can't really assert the error message for the above
        # Also, test_nested_array is a similar case (but without list) that works.

        if sys.maxsize > 2 ** 32:
            bits = 64
        else:
            bits = 32

    def test_objmode_inlining(self):
        def objmode_func(y):
            z = object()
            inlined = [x for x in y]
            return inlined

        cfunc = jit(forceobj=True)(objmode_func)
        t = [1, 2, 3]
        expected = objmode_func(t)
        got = cfunc(t)
        self.assertPreciseEqual(expected, got)


class TestArrayComprehension(unittest.TestCase):

    _numba_parallel_test_ = False

    def check(self, pyfunc, *args, **kwargs):
        """A generic check function that run both pyfunc, and jitted pyfunc,
        and compare results."""
        run_parallel = kwargs.get('run_parallel', False)
        assert_allocate_list = kwargs.get('assert_allocate_list', False)
        assert_dtype = kwargs.get('assert_dtype', False)
        cfunc = jit(nopython=True,parallel=run_parallel)(pyfunc)
        pyres = pyfunc(*args)
        cres = cfunc(*args)
        np.testing.assert_array_equal(pyres, cres)
        if assert_dtype:
            self.assertEqual(cres[1].dtype, assert_dtype)
        if assert_allocate_list:
            self.assertIn('allocate list', cfunc.inspect_llvm(cfunc.signatures[0]))
        else:
            self.assertNotIn('allocate list', cfunc.inspect_llvm(cfunc.signatures[0]))
        if run_parallel:
            self.assertIn('@do_scheduling', cfunc.inspect_llvm(cfunc.signatures[0]))

    def test_comp_with_array_1(self):
        def comp_with_array_1(n):
            m = n * 2
            l = np.array([i + m for i in range(n)])
            return l

        self.check(comp_with_array_1, 5)
        if PARALLEL_SUPPORTED:
            self.check(comp_with_array_1, 5, run_parallel=True)

    def test_comp_with_array_2(self):
        def comp_with_array_2(n, threshold):
            A = np.arange(-n, n)
            return np.array([ x * x if x < threshold else x * 2 for x in A ])

        self.check(comp_with_array_2, 5, 0)

    def test_comp_with_array_noinline(self):
        def comp_with_array_noinline(n):
            m = n * 2
            l = np.array([i + m for i in range(n)])
            return l

        import numba.core.inline_closurecall as ic
        try:
            ic.enable_inline_arraycall = False
            self.check(comp_with_array_noinline, 5, assert_allocate_list=True)
        finally:
            ic.enable_inline_arraycall = True

    def test_comp_with_array_noinline_issue_6053(self):
        def comp_with_array_noinline(n):
            lst = [0]
            for i in range(n):
                lst.append(i)
            l = np.array(lst)
            return l

        self.check(comp_with_array_noinline, 5, assert_allocate_list=True)

    def test_comp_nest_with_array(self):
        def comp_nest_with_array(n):
            l = np.array([[i * j for j in range(n)] for i in range(n)])
            return l

        self.check(comp_nest_with_array, 5)
        if PARALLEL_SUPPORTED:
            self.check(comp_nest_with_array, 5, run_parallel=True)

    def test_comp_nest_with_array_3(self):
        def comp_nest_with_array_3(n):
            l = np.array([[[i * j * k for k in range(n)] for j in range(n)] for i in range(n)])
            return l

        self.check(comp_nest_with_array_3, 5)
        if PARALLEL_SUPPORTED:
            self.check(comp_nest_with_array_3, 5, run_parallel=True)

    def test_comp_nest_with_array_noinline(self):
        def comp_nest_with_array_noinline(n):
            l = np.array([[i * j for j in range(n)] for i in range(n)])
            return l

        import numba.core.inline_closurecall as ic
        try:
            ic.enable_inline_arraycall = False
            self.check(comp_nest_with_array_noinline, 5,
                       assert_allocate_list=True)
        finally:
            ic.enable_inline_arraycall = True

    def test_comp_with_array_range(self):
        def comp_with_array_range(m, n):
            l = np.array([i for i in range(m, n)])
            return l

        self.check(comp_with_array_range, 5, 10)

    def test_comp_with_array_range_and_step(self):
        def comp_with_array_range_and_step(m, n):
            l = np.array([i for i in range(m, n, 2)])
            return l

        self.check(comp_with_array_range_and_step, 5, 10)

    def test_comp_with_array_conditional(self):
        def comp_with_array_conditional(n):
            l = np.array([i for i in range(n) if i % 2 == 1])
            return l
        # arraycall inline would not happen when conditional is present
        self.check(comp_with_array_conditional, 10, assert_allocate_list=True)

    def test_comp_nest_with_array_conditional(self):
        def comp_nest_with_array_conditional(n):
            l = np.array([[i * j for j in range(n)] for i in range(n) if i % 2 == 1])
            return l
        self.check(comp_nest_with_array_conditional, 5,
                   assert_allocate_list=True)

    @unittest.skipUnless(numpy_version < (1, 24),
                         'Setting an array element with a sequence is removed '
                         'in NumPy 1.24')
    def test_comp_nest_with_dependency(self):
        def comp_nest_with_dependency(n):
            l = np.array([[i * j for j in range(i+1)] for i in range(n)])
            return l
        # test is expected to fail
        with self.assertRaises(TypingError) as raises:
            self.check(comp_nest_with_dependency, 5)
        self.assertIn(_header_lead, str(raises.exception))
        self.assertIn('array(undefined,', str(raises.exception))

    def test_comp_unsupported_iter(self):
        def comp_unsupported_iter():
            val = zip([1, 2, 3], [4, 5, 6])
            return np.array([a for a, b in val])
        with self.assertRaises(TypingError) as raises:
            self.check(comp_unsupported_iter)
        self.assertIn(_header_lead, str(raises.exception))
        self.assertIn('Unsupported iterator found in array comprehension',
                      str(raises.exception))

    def test_no_array_comp(self):
        def no_array_comp1(n):
            l = [1,2,3,4]
            a = np.array(l)
            return a
        # const 1D array is actually inlined
        self.check(no_array_comp1, 10, assert_allocate_list=False)
        def no_array_comp2(n):
            l = [1,2,3,4]
            a = np.array(l)
            l.append(5)
            return a
        self.check(no_array_comp2, 10, assert_allocate_list=True)

    def test_nested_array(self):
        def nested_array(n):
            l = np.array([ np.array([x for x in range(n)]) for y in range(n)])
            return l

        self.check(nested_array, 10)

    def test_nested_array_with_const(self):
        def nested_array(n):
            l = np.array([ np.array([x for x in range(3)]) for y in range(4)])
            return l

        self.check(nested_array, 0)

    def test_array_comp_with_iter(self):
        def array_comp(a):
            l = np.array([ x * x for x in a ])
            return l
        # with list iterator
        l = [1,2,3,4,5]
        self.check(array_comp, l)
        # with array iterator
        self.check(array_comp, np.array(l))
        # with tuple iterator (issue #7394)
        self.check(array_comp, tuple(l))
        # with typed.List iterator (issue #6550)
        self.check(array_comp, typed.List(l))

    def test_array_comp_with_dtype(self):
        def array_comp(n):
            l = np.array([i for i in range(n)], dtype=np.complex64)
            return l

        self.check(array_comp, 10, assert_dtype=np.complex64)

    def test_array_comp_inferred_dtype(self):
        def array_comp(n):
            l = np.array([i * 1j for i in range(n)])
            return l

        self.check(array_comp, 10)

    def test_array_comp_inferred_dtype_nested(self):
        def array_comp(n):
            l = np.array([[i * j for j in range(n)] for i in range(n)])
            return l

        self.check(array_comp, 10)

    def test_array_comp_inferred_dtype_nested_sum(self):
        def array_comp(n):
            l = np.array([[i * j for j in range(n)] for i in range(n)])
            # checks that operations on the inferred array
            return l

        self.check(array_comp, 10)

    def test_array_comp_inferred_dtype_outside_setitem(self):
        def array_comp(n, v):
            arr = np.array([i for i in range(n)])
            # the following should not change the dtype
            arr[0] = v
            return arr

        # float to int cast is valid
        v = 1.2
        self.check(array_comp, 10, v, assert_dtype=np.intp)
        # complex to int cast is invalid
        with self.assertRaises(TypingError) as raises:
            cfunc = jit(nopython=True)(array_comp)
            cfunc(10, 2.3j)
        self.assertIn(
            _header_lead + " Function({})".format(operator.setitem),
            str(raises.exception),
        )
        self.assertIn(
            "(array({}, 1d, C), Literal[int](0), complex128)".format(types.intp),
            str(raises.exception),
        )

    def test_array_comp_shuffle_sideeffect(self):
        nelem = 100

        @jit(nopython=True)
        def foo():
            numbers = np.array([i for i in range(nelem)])
            np.random.shuffle(numbers)
            print(numbers)

        with captured_stdout() as gotbuf:
            foo()
        got = gotbuf.getvalue().strip()

        with captured_stdout() as expectbuf:
            print(np.array([i for i in range(nelem)]))
        expect = expectbuf.getvalue().strip()

        # For a large enough array, the chances of shuffle to not move any
        # element is tiny enough.
        self.assertNotEqual(got, expect)
        self.assertRegex(got, r'\[(\s*\d+)+\]')

    def test_empty_list_not_removed(self):
        # see issue #3724
        def f(x):
            t = []
            myList = np.array([1])
            a = np.random.choice(myList, 1)
            t.append(x + a)
            return a
        self.check(f, 5, assert_allocate_list=True)

    def test_reuse_of_array_var(self):
        """ Test issue 3742 """
        # redefinition of z breaks array comp as there's multiple defn
        def foo(n):
            # doesn't matter where this is in the code, it's just to ensure a
            # `make_function` opcode exists
            [i for i in range(1)]
            z = np.empty(n)
            for i in range(n):
                z = np.zeros(n)
                z[i] = i # write is required to trip the bug

            return z

        self.check(foo, 10, assert_allocate_list=True)


if __name__ == '__main__':
    unittest.main()