"""
Tests for sub-components of parfors.
These tests are aimed to produce a good-enough coverage of parfor passes
so that refactoring on these passes are easier with faster testing turnaround.
"""
import unittest
from functools import reduce

import numpy as np

from numba import njit, typeof, prange, pndindex
import numba.parfors.parfor
from numba.core import (
    typing,
    rewrites,
    typed_passes,
    untyped_passes,
    inline_closurecall,
    compiler,
    cpu,
    errors
)
from numba.core.registry import cpu_target
from numba.tests.support import (TestCase, is_parfors_unsupported)


class MyPipeline(object):
    def __init__(self, typingctx, targetctx, args, test_ir):
        self.state = compiler.StateDict()
        self.state.typingctx = typingctx
        self.state.targetctx = targetctx
        self.state.args = args
        self.state.func_ir = test_ir
        self.state.typemap = None
        self.state.return_type = None
        self.state.calltypes = None
        self.state.metadata = {}


class BaseTest(TestCase):
    @classmethod
    def _run_parfor(cls, test_func, args, swap_map=None):
        # TODO: refactor this with get_optimized_numba_ir() where this is
        #       copied from
        typingctx = typing.Context()
        targetctx = cpu.CPUContext(typingctx)
        test_ir = compiler.run_frontend(test_func)
        options = cpu.ParallelOptions(True)

        tp = MyPipeline(typingctx, targetctx, args, test_ir)

        with cpu_target.nested_context(typingctx, targetctx):
            typingctx.refresh()
            targetctx.refresh()

            inline_pass = inline_closurecall.InlineClosureCallPass(
                tp.state.func_ir, options, typed=True
            )
            inline_pass.run()

            rewrites.rewrite_registry.apply("before-inference", tp.state)

            untyped_passes.ReconstructSSA().run_pass(tp.state)

            (
                tp.state.typemap,
                tp.state.return_type,
                tp.state.calltypes,
                _
            ) = typed_passes.type_inference_stage(
                tp.state.typingctx, tp.state.targetctx, tp.state.func_ir,
                tp.state.args, None
            )

            typed_passes.PreLowerStripPhis().run_pass(tp.state)

            diagnostics = numba.parfors.parfor.ParforDiagnostics()

            preparfor_pass = numba.parfors.parfor.PreParforPass(
                tp.state.func_ir,
                tp.state.typemap,
                tp.state.calltypes,
                tp.state.typingctx,
                tp.state.targetctx,
                options,
                swapped=diagnostics.replaced_fns,
                replace_functions_map=swap_map,
            )
            preparfor_pass.run()

            rewrites.rewrite_registry.apply("after-inference", tp.state)
            return tp, options, diagnostics, preparfor_pass

    @classmethod
    def run_parfor_sub_pass(cls, test_func, args):
        tp, options, diagnostics, _ = cls._run_parfor(test_func, args)

        flags = compiler.Flags()
        parfor_pass = numba.parfors.parfor.ParforPass(
            tp.state.func_ir,
            tp.state.typemap,
            tp.state.calltypes,
            tp.state.return_type,
            tp.state.typingctx,
            tp.state.targetctx,
            options,
            flags,
            tp.state.metadata,
            diagnostics=diagnostics,
        )
        parfor_pass._pre_run()
        # Run subpass
        sub_pass = cls.sub_pass_class(parfor_pass)
        sub_pass.run(parfor_pass.func_ir.blocks)

        return sub_pass

    @classmethod
    def run_parfor_pre_pass(cls, test_func, args, swap_map=None):
        tp, options, diagnostics, preparfor_pass = cls._run_parfor(
            test_func, args, swap_map
        )
        return preparfor_pass

    def _run_parallel(self, func, *args, **kwargs):
        cfunc = njit(parallel=True)(func)
        expect = func(*args, **kwargs)
        got = cfunc(*args, **kwargs)
        return expect, got

    def run_parallel(self, func, *args, **kwargs):
        if is_parfors_unsupported:
            # Skip
            return
        expect, got = self._run_parallel(func, *args, **kwargs)
        self.assertPreciseEqual(expect, got)

    def run_parallel_check_output_array(self, func, *args, **kwargs):
        if is_parfors_unsupported:
            # Skip
            return
        expect, got = self._run_parallel(func, *args, **kwargs)
        # Don't match the value, just the return type. must return array
        self.assertIsInstance(expect, np.ndarray)
        self.assertIsInstance(got, np.ndarray)
        self.assertEqual(expect.shape, got.shape)

    def check_records(self, records):
        for rec in records:
            self.assertIsInstance(rec["new"], numba.parfors.parfor.Parfor)


class TestConvertSetItemPass(BaseTest):
    sub_pass_class = numba.parfors.parfor.ConvertSetItemPass

    def test_setitem_full_slice(self):
        def test_impl():
            n = 10
            a = np.ones(n)
            a[:] = 7
            return a

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "slice")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl)

    def test_setitem_slice_stop_bound(self):
        def test_impl():
            n = 10
            a = np.ones(n)
            a[:5] = 7
            return a

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "slice")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl)

    def test_setitem_slice_start_bound(self):
        def test_impl():
            n = 10
            a = np.ones(n)
            a[4:] = 7
            return a

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "slice")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl)

    def test_setitem_gather_if_scalar(self):
        def test_impl():
            n = 10
            a = np.ones(n)
            b = np.ones_like(a, dtype=np.bool_)
            a[b] = 7
            return a

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "masked_assign_broadcast_scalar")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl)

    def test_setitem_gather_if_array(self):
        def test_impl():
            n = 10
            a = np.ones(n)
            b = np.ones_like(a, dtype=np.bool_)
            c = np.ones_like(a)
            a[b] = c[b]
            return a

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "masked_assign_array")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl)


class TestConvertNumpyPass(BaseTest):
    sub_pass_class = numba.parfors.parfor.ConvertNumpyPass

    def check_numpy_allocators(self, fn):
        def test_impl():
            n = 10
            a = fn(n)
            return a

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "numpy_allocator")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl)

    def check_numpy_random(self, fn):
        def test_impl():
            n = 10
            a = fn(n)
            return a

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "numpy_allocator")
        self.check_records(sub_pass.rewritten)

        self.run_parallel_check_output_array(test_impl)

    def test_numpy_allocators(self):
        fns = [np.ones, np.zeros]
        for fn in fns:
            with self.subTest(fn.__name__):
                self.check_numpy_allocators(fn)

    def test_numpy_random(self):
        fns = [np.random.random]
        for fn in fns:
            with self.subTest(fn.__name__):
                self.check_numpy_random(fn)

    def test_numpy_arrayexpr(self):
        def test_impl(a, b):
            return a + b

        a = b = np.ones(10)

        args = (a, b)
        argtypes = [typeof(x) for x in args]

        sub_pass = self.run_parfor_sub_pass(test_impl, argtypes)
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "arrayexpr")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl, *args)

    def test_numpy_arrayexpr_ufunc(self):
        def test_impl(a, b):
            return np.sin(-a) + np.float64(1) / np.sqrt(b)

        a = b = np.ones(10)

        args = (a, b)
        argtypes = [typeof(x) for x in args]

        sub_pass = self.run_parfor_sub_pass(test_impl, argtypes)
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "arrayexpr")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl, *args)

    def test_numpy_arrayexpr_boardcast(self):
        def test_impl(a, b):
            return a + b + np.array(1)

        a = np.ones(10)
        b = np.ones((3, 10))

        args = (a, b)
        argtypes = [typeof(x) for x in args]

        sub_pass = self.run_parfor_sub_pass(test_impl, argtypes)
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "arrayexpr")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl, *args)

    def test_numpy_arrayexpr_reshaped(self):
        def test_impl(a, b):
            a = a.reshape(1, a.size)  # shape[0] is now constant
            return a + b

        a = np.ones(10)
        b = np.ones(10)

        args = (a, b)
        argtypes = [typeof(x) for x in args]

        sub_pass = self.run_parfor_sub_pass(test_impl, argtypes)
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "arrayexpr")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl, *args)


class TestConvertReducePass(BaseTest):
    sub_pass_class = numba.parfors.parfor.ConvertReducePass

    def test_reduce_max_basic(self):
        def test_impl(arr):
            return reduce(lambda x, y: max(x, y), arr, 0.0)

        x = np.ones(10)
        args = (x,)
        argtypes = [typeof(x) for x in args]

        sub_pass = self.run_parfor_sub_pass(test_impl, argtypes)
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "reduce")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl, *args)

    def test_reduce_max_masked(self):
        def test_impl(arr):
            return reduce(lambda x, y: max(x, y), arr[arr > 5], 0.0)

        x = np.ones(10)
        args = (x,)
        argtypes = [typeof(x) for x in args]

        sub_pass = self.run_parfor_sub_pass(test_impl, argtypes)
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "reduce")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl, *args)


class TestConvertLoopPass(BaseTest):
    sub_pass_class = numba.parfors.parfor.ConvertLoopPass

    def test_prange_reduce_simple(self):
        def test_impl():
            n = 20
            c = 0
            for i in prange(n):
                c += i
            return c

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "loop")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl)

    def test_prange_map_simple(self):
        def test_impl():
            n = 20
            arr = np.ones(n)
            for i in prange(n):
                arr[i] += i
            return arr

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "loop")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl)

    def test_prange_two_args(self):
        def test_impl():
            n = 20
            arr = np.ones(n)
            for i in prange(3, n):
                arr[i] += i
            return arr

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "loop")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl)

    def test_prange_three_args(self):
        def test_impl():
            n = 20
            arr = np.ones(n)
            for i in prange(3, n, 2):
                arr[i] += i
            return arr

        with self.assertRaises(errors.UnsupportedRewriteError) as raises:
            self.run_parfor_sub_pass(test_impl, ())
        self.assertIn(
            "Only constant step size of 1 is supported for prange",
            str(raises.exception),
        )

    def test_prange_map_inner_loop(self):
        def test_impl():
            n = 20
            arr = np.ones((n, n))
            for i in prange(n):
                for j in range(i):
                    arr[i, j] += i + j * n
            return arr

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "loop")
        self.check_records(sub_pass.rewritten)

        self.run_parallel(test_impl)

    def test_prange_map_nested_prange(self):
        def test_impl():
            n = 20
            arr = np.ones((n, n))
            for i in prange(n):
                for j in prange(i):
                    arr[i, j] += i + j * n
            return arr

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 2)
        self.check_records(sub_pass.rewritten)
        for record in sub_pass.rewritten:
            self.assertEqual(record["reason"], "loop")

        self.run_parallel(test_impl)

    def test_prange_map_none_index(self):
        def test_impl():
            n = 20
            arr = np.ones(n)
            for i in prange(n):
                inner = arr[i : i + 1]
                inner[()] += 1
            return arr

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        self.check_records(sub_pass.rewritten)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "loop")

        self.run_parallel(test_impl)

    def test_prange_map_overwrite_index(self):
        def test_impl():
            n = 20
            arr = np.ones(n)
            for i in prange(n):
                i += 1
                arr[i - 1] = i
            return arr

        with self.assertRaises(errors.UnsupportedRewriteError) as raises:
            self.run_parfor_sub_pass(test_impl, ())
        self.assertIn(
            "Overwrite of parallel loop index",
            str(raises.exception),
        )

    def test_init_prange(self):
        def test_impl():
            n = 20
            arr = np.ones(n)
            numba.parfors.parfor.init_prange()
            val = 0
            for i in numba.parfors.parfor.internal_prange(len(arr)):
                val += arr[i]
            return val

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        self.check_records(sub_pass.rewritten)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "loop")

        self.run_parallel(test_impl)

    def test_pndindex(self):
        def test_impl():
            n = 20
            arr = np.ones((n, n))
            val = 0
            for idx in pndindex(arr.shape):
                val += idx[0] * idx[1]
            return val

        sub_pass = self.run_parfor_sub_pass(test_impl, ())
        self.assertEqual(len(sub_pass.rewritten), 1)
        self.check_records(sub_pass.rewritten)

        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "loop")

        self.run_parallel(test_impl)

    def test_numpy_sum(self):
        def test_impl(arr):
            return np.sum(arr)

        shape = 11, 13
        arr = np.arange(np.prod(shape)).reshape(shape)
        args = (arr,)
        argtypes = [typeof(x) for x in args]

        sub_pass = self.run_parfor_sub_pass(test_impl, argtypes)
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "loop")
        self.check_records(sub_pass.rewritten)
        self.run_parallel(test_impl, *args)

    def test_numpy_sum_bool_array_masked(self):
        def test_impl(arr):
            sliced = arr[:, 0]
            return np.sum(arr[sliced >= 3, 1:2])

        shape = 11, 13
        arr = np.arange(np.prod(shape)).reshape(shape)
        args = (arr,)
        argtypes = [typeof(x) for x in args]

        sub_pass = self.run_parfor_sub_pass(test_impl, argtypes)
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "loop")
        self.check_records(sub_pass.rewritten)
        self.run_parallel(test_impl, *args)

    def test_numpy_sum_int_array_masked(self):
        def test_impl(arr):
            sel = np.arange(arr.shape[1])
            return np.sum(arr[:, sel])

        shape = 11, 13
        arr = np.arange(np.prod(shape)).reshape(shape)
        args = (arr,)
        argtypes = [typeof(x) for x in args]

        sub_pass = self.run_parfor_sub_pass(test_impl, argtypes)
        # 1 for arange; 1 for sum
        self.assertEqual(len(sub_pass.rewritten), 2)
        for record in sub_pass.rewritten:
            self.assertEqual(record["reason"], "loop")
        self.check_records(sub_pass.rewritten)
        self.run_parallel(test_impl, *args)

    def test_numpy_fill_method(self):
        def test_impl(arr):
            arr.fill(3)
            return arr

        shape = 11, 13
        arr = np.arange(np.prod(shape)).reshape(shape)
        args = (arr,)
        argtypes = [typeof(x) for x in args]

        sub_pass = self.run_parfor_sub_pass(test_impl, argtypes)
        # 1 for arange; 1 for sum
        self.assertEqual(len(sub_pass.rewritten), 1)
        [record] = sub_pass.rewritten
        self.assertEqual(record["reason"], "loop")
        self.check_records(sub_pass.rewritten)
        self.run_parallel(test_impl, *args)


class TestPreParforPass(BaseTest):
    class sub_pass_class:
        def __init__(self, pass_states):
            pass

        def run(self, blocks):
            pass

    def test_dtype_conversion(self):
        # array.dtype are converted to np.dtype(array) in the PreParforPass
        def test_impl(a):
            b = np.ones(20, dtype=a.dtype)
            return b

        arr = np.arange(10)
        args = (arr,)
        argtypes = [typeof(x) for x in args]

        pre_pass = self.run_parfor_pre_pass(test_impl, argtypes)
        self.assertEqual(pre_pass.stats["replaced_func"], 0)
        self.assertEqual(pre_pass.stats["replaced_dtype"], 1)
        self.run_parallel(test_impl, *args)

    def test_sum_replacement(self):
        def test_impl(a):
            return np.sum(a)

        arr = np.arange(10)
        args = (arr,)
        argtypes = [typeof(x) for x in args]

        pre_pass = self.run_parfor_pre_pass(test_impl, argtypes)
        self.assertEqual(pre_pass.stats["replaced_func"], 1)
        self.assertEqual(pre_pass.stats["replaced_dtype"], 0)
        self.run_parallel(test_impl, *args)

    def test_replacement_map(self):
        def test_impl(a):
            return np.sum(a)

        arr = np.arange(10)
        args = (arr,)
        argtypes = [typeof(x) for x in args]

        swap_map = numba.parfors.parfor.swap_functions_map.copy()
        swap_map.pop(("sum", "numpy"))
        pre_pass = self.run_parfor_pre_pass(test_impl, argtypes, swap_map)
        self.assertEqual(pre_pass.stats["replaced_func"], 0)
        self.assertEqual(pre_pass.stats["replaced_dtype"], 0)
        self.run_parallel(test_impl, *args)


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