import gc
from itertools import product

import numpy as np
from numpy.polynomial import polynomial as poly
from numpy.polynomial import polyutils as pu

from numba import jit, njit
from numba.tests.support import (TestCase, needs_lapack,
                                 EnableNRTStatsMixin, MemoryLeakMixin)
from numba.core.errors import TypingError


def roots_fn(p):
    return np.roots(p)


def polyadd(c1,c2):
    return poly.polyadd(c1,c2)


def polysub(c1,c2):
    return poly.polysub(c1,c2)


def polymul(c1,c2):
    return poly.polymul(c1,c2)


def trimseq(seq):
    return pu.trimseq(seq)


def polyasseries1(a):
    res = pu.as_series(a)
    return res


def polyasseries2(a, trim):
    res = pu.as_series(a, trim)
    return res


def polydiv(c1, c2):
    res = poly.polydiv(c1, c2)
    return res


def polyval2(x, c):
    res = poly.polyval(x, c)
    return res


def polyval3T(x, c):
    res = poly.polyval(x, c, True)
    return res


def polyval3F(x, c):
    res = poly.polyval(x, c, False)
    return res


def polyint(c, m=1):
    res = poly.polyint(c, m)
    return res


class TestPolynomialBase(EnableNRTStatsMixin, TestCase):
    """
    Provides setUp and common data/error modes for testing polynomial functions.
    """

    # supported dtypes
    dtypes = (np.float64, np.float32, np.complex128, np.complex64)

    def setUp(self):
        # Collect leftovers from previous test cases before checking for leaks
        gc.collect()
        super(TestPolynomialBase, self).setUp()

    def assert_error(self, cfunc, args, msg, err=ValueError):
        with self.assertRaises(err) as raises:
            cfunc(*args)
        self.assertIn(msg, str(raises.exception))

    def assert_1d_input(self, cfunc, args):
        msg = "Input must be a 1d array."
        self.assert_error(cfunc, args, msg)


class TestPoly1D(TestPolynomialBase):

    def assert_no_domain_change(self, name, cfunc, args):
        msg = name + "() argument must not cause a domain change."
        self.assert_error(cfunc, args, msg)

    @needs_lapack
    def test_roots(self):

        cfunc = jit(nopython=True)(roots_fn)

        default_resolution = np.finfo(np.float64).resolution

        def check(a, **kwargs):
            expected = roots_fn(a, **kwargs)
            got = cfunc(a, **kwargs)

            # eigen decomposition used so type specific impl
            # will be used in numba whereas a wide type impl
            # will be used in numpy, so compare using a more
            # fuzzy comparator

            if a.dtype in self.dtypes:
                resolution = np.finfo(a.dtype).resolution
            else:
                # this is for integer types when roots() will cast to float64
                resolution = default_resolution

            np.testing.assert_allclose(
                expected,
                got,
                rtol=10 * resolution,
                atol=100 * resolution  # zeros tend to be fuzzy
            )

            # Ensure proper resource management
            with self.assertNoNRTLeak():
                cfunc(a, **kwargs)

        # test vectors in real space
        # contrived examples to trip branches
        r_vectors = (
            np.array([1]),
            np.array([1, 3, 2]),
            np.array([0, 0, 0]),
            np.array([1, 6, 11, 6]),
            np.array([0, 0, 0, 1, 3, 2]),
            np.array([1, 1, 0, 0, 0]),
            np.array([0, 0, 1, 0, 0, 0])
        )

        # test loop real space
        for v, dtype in \
                product(r_vectors, [np.int32, np.int64] + list(self.dtypes)):
            a = v.astype(dtype)
            check(a)

        c_vectors = (
            np.array([1 + 1j]),
            np.array([1, 3 + 1j, 2]),
            np.array([0, 0 + 0j, 0]),
            np.array([1, 6 + 1j, 11, 6]),
            np.array([0, 0, 0, 1 + 1j, 3, 2]),
            np.array([1 + 1j, 1, 0, 0, 0]),
            np.array([0, 0, 1 + 1j, 0, 0, 0])
        )

        # test loop complex space
        for v, dtype in product(c_vectors, self.dtypes[2:]):
            a = v.astype(dtype)
            check(a)

        # check input with dimension > 1 raises
        self.assert_1d_input(cfunc, (np.arange(4.).reshape(2, 2),))

        # check real input with complex roots raises
        x = np.array([7., 2., 0., 1.])
        self.assert_no_domain_change("eigvals", cfunc, (x,))
        # but works fine if type conv to complex first
        cfunc(x.astype(np.complex128))


class TestPolynomial(MemoryLeakMixin, TestCase):

    #
    # tests for Polyutils functions
    #

    def test_trimseq_basic(self):
        pyfunc = trimseq
        cfunc = njit(trimseq)

        def inputs():
            for i in range(5):
                yield np.array([1] + [0] * i)

        for coefs in inputs():
            self.assertPreciseEqual(pyfunc(coefs), cfunc(coefs))

    def test_trimseq_exception(self):
        cfunc = njit(trimseq)

        self.disable_leak_check()

        with self.assertRaises(TypingError) as raises:
            cfunc("abc")
        self.assertIn('The argument "seq" must be array-like',
                      str(raises.exception))

        with self.assertRaises(TypingError) as e:
            cfunc(np.arange(10).reshape(5, 2))
        self.assertIn('Coefficient array is not 1-d',
                      str(e.exception))

        with self.assertRaises(TypingError) as e:
            cfunc((1, 2, 3, 0))
        self.assertIn('Unsupported type UniTuple(int64, 4) for argument "seq"',
                      str(e.exception))

    def test_pu_as_series_basic(self):
        pyfunc1 = polyasseries1
        cfunc1 = njit(polyasseries1)
        pyfunc2 = polyasseries2
        cfunc2 = njit(polyasseries2)

        def inputs():
            yield np.arange(4)
            yield np.arange(6).reshape((2,3))
            yield (1, np.arange(3), np.arange(2, dtype=np.float32))
            yield ([1, 2, 3, 4, 0], [1, 2, 3])
            yield ((0, 0, 1e-3, 0, 1e-5, 0, 0), (1, 2, 3, 4, 5, 6, 7))
            yield ((0, 0, 1e-3, 0, 1e-5, 0, 0), (1j, 2, 3j, 4j, 5, 6j, 7))
            yield (2, [1.1, 0.])
            yield ([1, 2, 3, 0], )
            yield ((1, 2, 3, 0), )
            yield (np.array([1, 2, 3, 0]), )
            yield [np.array([1, 2, 3, 0]), np.array([1, 2, 3, 0])]
            yield [np.array([1,2,3]), ]

        for input in inputs():
            self.assertPreciseEqual(pyfunc1(input), cfunc1(input))
            self.assertPreciseEqual(pyfunc2(input, False), cfunc2(input, False))
            self.assertPreciseEqual(pyfunc2(input, True), cfunc2(input, True))

    def test_pu_as_series_exception(self):
        cfunc1 = njit(polyasseries1)
        cfunc2 = njit(polyasseries2)

        self.disable_leak_check()

        with self.assertRaises(TypingError) as raises:
            cfunc1("abc")
        self.assertIn('The argument "alist" must be array-like',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc2("abc", True)
        self.assertIn('The argument "alist" must be array-like',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc2(np.arange(4), "abc")
        self.assertIn('The argument "trim" must be boolean',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc1(([1, 2, 3], np.arange(16).reshape(4,4)))
        self.assertIn('Coefficient array is not 1-d',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc1(np.arange(8).reshape((2, 2, 2)))
        self.assertIn('Coefficient array is not 1-d',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc1([np.array([[1,2,3],[1,2,3]]), ])
        self.assertIn('Coefficient array is not 1-d',
                      str(raises.exception))

        with self.assertRaises(ValueError) as raises:
            cfunc1(np.array([[]], dtype=np.float64))
        self.assertIn('Coefficient array is empty',
                      str(raises.exception))

        with self.assertRaises(ValueError) as raises:
            cfunc1(([1, 2, 3], np.array([], dtype=np.float64),
                    np.array([1,2,1])))
        self.assertIn('Coefficient array is empty',
                      str(raises.exception))

    def _test_polyarithm_basic(self, pyfunc, ignore_sign_on_zero=False):
        # test suite containing tests for polyadd, polysub, polymul, polydiv
        cfunc = njit(pyfunc)

        def inputs():
            # basic, taken from https://github.com/numpy/numpy/blob/48a8277855849be094a5979c48d9f5f1778ee4de/numpy/polynomial/tests/test_polynomial.py#L58-L123 # noqa: E501
            for i in range(5):
                for j in range(5):
                    p1 = np.array([0] * i + [1])
                    p2 = np.array([0] * j + [1])
                    yield p1, p2
            # test lists, tuples, scalars
            yield [1, 2, 3], [1, 2, 3]
            yield [1, 2, 3], (1, 2, 3)
            yield (1, 2, 3), [1, 2, 3]
            yield [1, 2, 3], 3
            yield 3, (1, 2, 3)
            # test different dtypes
            yield np.array([1, 2, 3]), np.array([1.0, 2.0, 3.0])
            yield np.array([1j, 2j, 3j]), np.array([1.0, 2.0, 3.0])
            yield np.array([1, 2, 3]), np.array([1j, 2j, 3j])
            yield (1, 2, 3), 3.0
            yield (1, 2, 3), 3j
            yield (1, 1e-3, 3), (1, 2, 3)

        for p1, p2 in inputs():
            self.assertPreciseEqual(pyfunc(p1,p2), cfunc(p1,p2),
                                    ignore_sign_on_zero=ignore_sign_on_zero)

    def _test_polyarithm_exception(self, pyfunc):
        # test suite containing tests for polyadd, polysub, polymul, polydiv
        cfunc = njit(pyfunc)

        self.disable_leak_check()

        with self.assertRaises(TypingError) as raises:
            cfunc("abc", np.array([1,2,3]))
        self.assertIn('The argument "c1" must be array-like',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc(np.array([1,2,3]), "abc")
        self.assertIn('The argument "c2" must be array-like',
                      str(raises.exception))

        with self.assertRaises(TypingError) as e:
            cfunc(np.arange(10).reshape(5, 2), np.array([1, 2, 3]))
        self.assertIn('Coefficient array is not 1-d',
                      str(e.exception))

        with self.assertRaises(TypingError) as e:
            cfunc(np.array([1, 2, 3]), np.arange(10).reshape(5, 2))
        self.assertIn('Coefficient array is not 1-d',
                      str(e.exception))

    def test_polyadd_basic(self):
        self._test_polyarithm_basic(polyadd)

    def test_polyadd_exception(self):
        self._test_polyarithm_exception(polyadd)

    def test_polysub_basic(self):
        self._test_polyarithm_basic(polysub, ignore_sign_on_zero=True)

    def test_polysub_exception(self):
        self._test_polyarithm_exception(polysub)

    def test_polymul_basic(self):
        self._test_polyarithm_basic(polymul)

    def test_polymul_exception(self):
        self._test_polyarithm_exception(polymul)

    def test_poly_polydiv_basic(self):
        pyfunc = polydiv
        cfunc = njit(polydiv)
        self._test_polyarithm_basic(polydiv)

        def inputs():
            # Based on https://github.com/numpy/numpy/blob/160c16f055d4d2fce072004e286d8075b31955cd/numpy/polynomial/tests/test_polynomial.py#L99-L114 # noqa: E501
            # check scalar division
            yield [2], [2]
            yield [2, 2], [2]
            # check rest.
            for i in range(5):
                for j in range(5):
                    ci = [0] * i + [1, 2]
                    cj = [0] * j + [1, 2]
                    tgt = poly.polyadd(ci, cj)
                    yield tgt, ci
            yield np.array([1,0,0,0,0,0,-1]), np.array([1,0,0,-1])

        for c1, c2 in inputs():
            self.assertPreciseEqual(pyfunc(c1, c2), cfunc(c1, c2))

    def test_poly_polydiv_exception(self):
        self._test_polyarithm_exception(polydiv)
        cfunc = njit(polydiv)
        # Based on https://github.com/numpy/numpy/blob/160c16f055d4d2fce072004e286d8075b31955cd/numpy/polynomial/tests/test_polynomial.py#L97 # noqa: E501
        # check zero division
        with self.assertRaises(ZeroDivisionError) as _:
            cfunc([1], [0])

    def test_poly_polyval_basic(self):
        pyfunc2 = polyval2
        cfunc2 = njit(polyval2)
        pyfunc3T = polyval3T
        cfunc3T = njit(polyval3T)
        pyfunc3F = polyval3F
        cfunc3F = njit(polyval3F)

        def inputs():
            # Based on https://github.com/numpy/numpy/blob/160c16f055d4d2fce072004e286d8075b31955cd/numpy/polynomial/tests/test_polynomial.py#L137-L157 # noqa: E501
            # check empty input
            yield np.array([], dtype=np.float64), [1]
            yield 1, [1,2,3]
            yield np.arange(4).reshape(2,2), [1,2,3]
            # check normal input
            for i in range(5):
                yield np.linspace(-1, 1), [0] * i + [1]
            yield np.linspace(-1, 1), [0, -1, 0, 1]
            # check that shape is preserved
            for i in range(3):
                dims = [2] * i
                x = np.zeros(dims)
                yield x, [1]
                yield x, [1, 0]
                yield x, [1, 0, 0]
            # Check that behaviour corresponds to tensor = False
            yield np.array([1, 2]), np.arange(4).reshape(2,2)
            yield [1, 2], np.arange(4).reshape(2,2)

        for x, c in inputs():
            self.assertPreciseEqual(pyfunc2(x, c), cfunc2(x, c))
            # test tensor argument
            self.assertPreciseEqual(pyfunc3T(x, c), cfunc3T(x, c))
            self.assertPreciseEqual(pyfunc3F(x, c), cfunc3F(x, c))

    def test_poly_polyval_exception(self):
        cfunc2 = njit(polyval2)
        cfunc3T = njit(polyval3T)
        cfunc3F = njit(polyval3F)

        self.disable_leak_check()

        with self.assertRaises(TypingError) as raises:
            cfunc2(3, "abc")
        self.assertIn('The argument "c" must be array-like',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc2("abc", 3)
        self.assertIn('The argument "x" must be array-like',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc2("abc", "def")
        self.assertIn('The argument "x" must be array-like',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc3T(3, "abc")
        self.assertIn('The argument "c" must be array-like',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc3T("abc", 3)
        self.assertIn('The argument "x" must be array-like',
                      str(raises.exception))

        @njit
        def polyval3(x, c, tensor):
            res = poly.polyval(x, c, tensor)
            return res
        with self.assertRaises(TypingError) as raises:
            polyval3(3, 3, "abc")
        self.assertIn('The argument "tensor" must be boolean',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc3F("abc", "def")
        self.assertIn('The argument "x" must be array-like',
                      str(raises.exception))

    def test_poly_polyint_basic(self):
        pyfunc = polyint
        cfunc = njit(polyint)
        # basic
        self.assertPreciseEqual(pyfunc([1,2,3]), cfunc([1,2,3]))
        # Based on https://github.com/numpy/numpy/blob/160c16f055d4d2fce072004e286d8075b31955cd/numpy/polynomial/tests/test_polynomial.py#L314-L381 # noqa: E501
        # test integration of zero polynomial
        for i in range(2, 5):
            self.assertPreciseEqual(pyfunc([0], m=i), cfunc([0], m=i))

        # check single integration with integration constant
        for i in range(5):
            pol = [0] * i + [1]
            self.assertPreciseEqual(pyfunc(pol, m=1), pyfunc(pol, m=1))

        # check multiple integrations with default k
        for i in range(5):
            for j in range(2, 5):
                pol = [0] * i + [1]
                self.assertPreciseEqual(pyfunc(pol, m=j), cfunc(pol, m=j))

        # test multidimensional arrays
        c2 = np.array([[0,1], [0,2]])
        self.assertPreciseEqual(pyfunc(c2), cfunc(c2))
        c3 = np.arange(8).reshape((2,2,2))
        self.assertPreciseEqual(pyfunc(c3), cfunc(c3))

    def test_poly_polyint_exception(self):
        cfunc = njit(polyint)

        self.disable_leak_check()

        with self.assertRaises(TypingError) as raises:
            cfunc("abc")
        self.assertIn('The argument "c" must be array-like',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc(np.array([1,2,3]), "abc")
        self.assertIn('The argument "m" must be an integer',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc(['a', 'b', 'c'], 1)
        self.assertIn('Input dtype must be scalar.',
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc(('a', 'b', 'c'), 1)
        self.assertIn('Input dtype must be scalar.',
                      str(raises.exception))

    #
    # tests for Polynomial class
    #

    def test_Polynomial_constructor(self):

        def pyfunc3(c, dom, win):
            p = poly.Polynomial(c, dom, win)
            return p
        cfunc3 = njit(pyfunc3)

        def pyfunc1(c):
            p = poly.Polynomial(c)
            return p
        cfunc1 = njit(pyfunc1)
        list1 = (np.array([0, 1]), np.array([0., 1.]))
        list2 = (np.array([0, 1]), np.array([0., 1.]))
        list3 = (np.array([0, 1]), np.array([0., 1.]))
        for c in list1:
            for dom in list2:
                for win in list3:
                    p1 = pyfunc3(c, dom, win)
                    p2 = cfunc3(c, dom, win)
                    q1 = pyfunc1(c)
                    q2 = cfunc1(c)
                    self.assertPreciseEqual(p1, p2)
                    self.assertPreciseEqual(p1.coef, p2.coef)
                    self.assertPreciseEqual(p1.domain, p2.domain)
                    self.assertPreciseEqual(p1.window, p2.window)
                    self.assertPreciseEqual(q1.coef, q2.coef)
                    self.assertPreciseEqual(q1.domain, q2.domain)
                    self.assertPreciseEqual(q1.window, q2.window)

    def test_Polynomial_exeption(self):
        def pyfunc3(c, dom, win):
            p = poly.Polynomial(c, dom, win)
            return p
        cfunc3 = njit(pyfunc3)

        self.disable_leak_check()

        input2 = np.array([1, 2])
        input3 = np.array([1, 2, 3])
        input2D = np.arange(4).reshape((2, 2))

        with self.assertRaises(ValueError) as raises:
            cfunc3(input2, input3, input2)
        self.assertIn("Domain has wrong number of elements.",
                      str(raises.exception))

        with self.assertRaises(ValueError) as raises:
            cfunc3(input2, input2, input3)
        self.assertIn("Window has wrong number of elements.",
                      str(raises.exception))

        with self.assertRaises(TypingError) as raises:
            cfunc3(input2D, input2, input2)
        self.assertIn("Coefficient array is not 1-d",
                      str(raises.exception))