import os
import subprocess
import sys
import warnings
import numpy as np

import unittest
from numba import jit
from numba.core.errors import NumbaWarning, deprecated, NumbaDeprecationWarning
from numba.core import errors


class TestBuiltins(unittest.TestCase):

    def check_objmode_deprecation_warning(self, w):
        # Object mode fall-back is slated for deprecation, check the warning
        msg = ("Fall-back from the nopython compilation path to the object "
               "mode compilation path has been detected")
        self.assertEqual(w.category, NumbaDeprecationWarning)
        self.assertIn(msg, str(w.message))

    def test_type_infer_warning(self):
        def add(x, y):
            a = {} # noqa dead
            return x + y

        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always', NumbaWarning)

            cfunc = jit(add)
            cfunc(1, 2)

            self.assertEqual(len(w), 3)
            # Type inference failure
            self.assertEqual(w[0].category, NumbaWarning)
            self.assertIn('type inference', str(w[0].message))

            # Object mode
            self.assertEqual(w[1].category, NumbaWarning)
            self.assertIn('object mode', str(w[1].message))

            # check objmode deprecation warning
            self.check_objmode_deprecation_warning(w[2])

    def test_return_type_warning(self):
        y = np.ones(4, dtype=np.float32)

        def return_external_array():
            return y

        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always', NumbaWarning)

            cfunc = jit(_nrt=False)(return_external_array)
            cfunc()

            self.assertEqual(len(w), 3)

            # Legal return value failure
            self.assertEqual(w[0].category, NumbaWarning)
            self.assertIn('return type', str(w[0].message))

            # Object mode fall-back
            self.assertEqual(w[1].category, NumbaWarning)
            self.assertIn('object mode without forceobj=True',
                          str(w[1].message))

            # check objmode deprecation warning
            self.check_objmode_deprecation_warning(w[2])

    def test_return_type_warning_with_nrt(self):
        """
        Rerun test_return_type_warning with nrt
        """
        y = np.ones(4, dtype=np.float32)

        def return_external_array():
            return y

        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always', NumbaWarning)

            cfunc = jit(return_external_array)
            cfunc()
            # No more warning
            self.assertEqual(len(w), 0)

    def test_no_warning_with_forceobj(self):
        def add(x, y):
            a = [] # noqa dead
            return x + y

        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always', NumbaWarning)

            cfunc = jit(add, forceobj=True)
            cfunc(1, 2)

            self.assertEqual(len(w), 0)

    def test_loop_lift_warn(self):
        def do_loop(x):
            a = {} # noqa dead
            for i in range(x.shape[0]):
                x[i] *= 2

        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always', NumbaWarning)

            x = np.ones(4, dtype=np.float32)
            cfunc = jit(do_loop)
            cfunc(x)

            self.assertEqual(len(w), 4)

            # Type inference failure (1st pass, in npm, fall-back to objmode
            # with looplift)
            self.assertEqual(w[0].category, NumbaWarning)
            self.assertIn('type inference', str(w[0].message))
            self.assertIn('WITH looplifting', str(w[0].message))

            # Type inference failure (2nd pass, objmode with lifted loops,
            # loop found but still failed, fall back to objmode no looplift)
            self.assertEqual(w[1].category, NumbaWarning)
            self.assertIn('type inference', str(w[1].message))
            self.assertIn('WITHOUT looplifting', str(w[1].message))

            # States compilation outcome
            self.assertEqual(w[2].category, NumbaWarning)
            self.assertIn('compiled in object mode without forceobj=True',
                          str(w[2].message))
            self.assertIn('but has lifted loops', str(w[2].message))

            # check objmode deprecation warning
            self.check_objmode_deprecation_warning(w[3])

    def test_deprecated(self):
        @deprecated('foo')
        def bar():
            pass

        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always')
            bar()

            self.assertEqual(len(w), 1)
            self.assertEqual(w[0].category, DeprecationWarning)
            self.assertIn('bar', str(w[0].message))
            self.assertIn('foo', str(w[0].message))

    def test_warnings_fixer(self):
        # For some context, see #4083

        wfix = errors.WarningsFixer(errors.NumbaWarning)
        with wfix.catch_warnings('foo', 10):
            warnings.warn(errors.NumbaWarning('same'))
            warnings.warn(errors.NumbaDeprecationWarning('same'))

        with warnings.catch_warnings(record=True) as w:
            warnings.simplefilter('always')
            wfix.flush()

            self.assertEqual(len(w), 2)
            # the order of these will be backwards to the above, the
            # WarningsFixer flush method sorts with a key based on str
            # comparison
            self.assertEqual(w[0].category, NumbaDeprecationWarning)
            self.assertEqual(w[1].category, NumbaWarning)
            self.assertIn('same', str(w[0].message))
            self.assertIn('same', str(w[1].message))

    def test_disable_performance_warnings(self):

        not_found_ret_code = 55
        found_ret_code = 99
        expected = "'parallel=True' was specified but no transformation"

        # NOTE: the error_usecases is needed as the NumbaPerformanceWarning's
        # for parallel=True failing to parallelise do not appear for functions
        # defined by string eval/exec etc.
        parallel_code = """if 1:
            import warnings
            from numba.tests.error_usecases import foo
            import numba
            with warnings.catch_warnings(record=True) as w:
                warnings.simplefilter('always')
                foo()
            for x in w:
                if x.category == numba.errors.NumbaPerformanceWarning:
                    if "%s" in str(x.message):
                        exit(%s)
            exit(%s)
        """ % (expected, found_ret_code, not_found_ret_code)

        # run in the standard env, warning should raise
        popen = subprocess.Popen([sys.executable, "-c", parallel_code])
        out, err = popen.communicate()
        self.assertEqual(popen.returncode, found_ret_code)

        # run in an env with performance warnings disabled, should not warn
        env = dict(os.environ)
        env['NUMBA_DISABLE_PERFORMANCE_WARNINGS'] = "1"
        popen = subprocess.Popen([sys.executable, "-c", parallel_code], env=env)
        out, err = popen.communicate()
        self.assertEqual(popen.returncode, not_found_ret_code)


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