import numpy as np
from numpy.testing import assert_, assert_allclose
import pytest

from scipy.special import _ufuncs, gammasgn
import scipy.special._orthogonal as orth
from scipy.special._testutils import FuncData


def test_eval_chebyt():
    n = np.arange(0, 10000, 7, dtype=np.dtype("long"))
    x = 2*np.random.rand() - 1
    v1 = np.cos(n*np.arccos(x))
    v2 = _ufuncs.eval_chebyt(n, x)
    assert_(np.allclose(v1, v2, rtol=1e-15))


def test_eval_chebyt_gh20129():
    # https://github.com/scipy/scipy/issues/20129
    assert _ufuncs.eval_chebyt(7, 2 + 0j) == 5042.0


def test_eval_genlaguerre_restriction():
    # check it returns nan for alpha <= -1
    assert_(np.isnan(_ufuncs.eval_genlaguerre(0, -1, 0)))
    assert_(np.isnan(_ufuncs.eval_genlaguerre(0.1, -1, 0)))


def test_warnings():
    # ticket 1334
    with np.errstate(all='raise'):
        # these should raise no fp warnings
        _ufuncs.eval_legendre(1, 0)
        _ufuncs.eval_laguerre(1, 1)
        _ufuncs.eval_gegenbauer(1, 1, 0)


class TestPolys:
    """
    Check that the eval_* functions agree with the constructed polynomials

    """

    def check_poly(self, func, cls, param_ranges=(), x_range=(), nn=10,
                   nparam=10, nx=10, rtol=1e-8):
        rng = np.random.RandomState(1234)

        dataset = []
        for n in np.arange(nn):
            params = [a + (b-a)*rng.rand(nparam) for a,b in param_ranges]
            params = np.asarray(params).T
            if not param_ranges:
                params = [0]
            for p in params:
                if param_ranges:
                    p = (n,) + tuple(p)
                else:
                    p = (n,)
                x = x_range[0] + (x_range[1] - x_range[0])*rng.rand(nx)
                x[0] = x_range[0]  # always include domain start point
                x[1] = x_range[1]  # always include domain end point
                poly = np.poly1d(cls(*p).coef)
                z = np.c_[np.tile(p, (nx,1)), x, poly(x)]
                dataset.append(z)

        dataset = np.concatenate(dataset, axis=0)

        def polyfunc(*p):
            p = (p[0].astype(np.dtype("long")),) + p[1:]
            return func(*p)

        with np.errstate(all='raise'):
            ds = FuncData(polyfunc, dataset, list(range(len(param_ranges)+2)), -1,
                          rtol=rtol)
            ds.check()

    def test_jacobi(self):
        self.check_poly(_ufuncs.eval_jacobi, orth.jacobi,
                        param_ranges=[(-0.99, 10), (-0.99, 10)],
                        x_range=[-1, 1], rtol=1e-5)

    def test_sh_jacobi(self):
        self.check_poly(_ufuncs.eval_sh_jacobi, orth.sh_jacobi,
                        param_ranges=[(1, 10), (0, 1)], x_range=[0, 1],
                        rtol=1e-5)

    def test_gegenbauer(self):
        self.check_poly(_ufuncs.eval_gegenbauer, orth.gegenbauer,
                        param_ranges=[(-0.499, 10)], x_range=[-1, 1],
                        rtol=1e-7)

    def test_chebyt(self):
        self.check_poly(_ufuncs.eval_chebyt, orth.chebyt,
                        param_ranges=[], x_range=[-1, 1])

    def test_chebyu(self):
        self.check_poly(_ufuncs.eval_chebyu, orth.chebyu,
                        param_ranges=[], x_range=[-1, 1])

    def test_chebys(self):
        self.check_poly(_ufuncs.eval_chebys, orth.chebys,
                        param_ranges=[], x_range=[-2, 2])

    def test_chebyc(self):
        self.check_poly(_ufuncs.eval_chebyc, orth.chebyc,
                        param_ranges=[], x_range=[-2, 2])

    def test_sh_chebyt(self):
        with np.errstate(all='ignore'):
            self.check_poly(_ufuncs.eval_sh_chebyt, orth.sh_chebyt,
                            param_ranges=[], x_range=[0, 1])

    def test_sh_chebyu(self):
        self.check_poly(_ufuncs.eval_sh_chebyu, orth.sh_chebyu,
                        param_ranges=[], x_range=[0, 1])

    def test_legendre(self):
        self.check_poly(_ufuncs.eval_legendre, orth.legendre,
                        param_ranges=[], x_range=[-1, 1])

    def test_sh_legendre(self):
        with np.errstate(all='ignore'):
            self.check_poly(_ufuncs.eval_sh_legendre, orth.sh_legendre,
                            param_ranges=[], x_range=[0, 1])

    def test_genlaguerre(self):
        self.check_poly(_ufuncs.eval_genlaguerre, orth.genlaguerre,
                        param_ranges=[(-0.99, 10)], x_range=[0, 100])

    def test_laguerre(self):
        self.check_poly(_ufuncs.eval_laguerre, orth.laguerre,
                        param_ranges=[], x_range=[0, 100])

    def test_hermite(self):
        self.check_poly(_ufuncs.eval_hermite, orth.hermite,
                        param_ranges=[], x_range=[-100, 100])

    def test_hermitenorm(self):
        self.check_poly(_ufuncs.eval_hermitenorm, orth.hermitenorm,
                        param_ranges=[], x_range=[-100, 100])


class TestRecurrence:
    """
    Check that the eval_* functions sig='ld->d' and 'dd->d' agree.

    """

    def check_poly(self, func, param_ranges=(), x_range=(), nn=10,
                   nparam=10, nx=10, rtol=1e-8):
        rng = np.random.default_rng(1234)

        dataset = []
        for n in np.arange(nn):
            params = [a + (b-a)*rng.random(nparam) for a,b in param_ranges]
            params = np.asarray(params).T
            if not param_ranges:
                params = [0]
            for p in params:
                if param_ranges:
                    p = (n,) + tuple(p)
                else:
                    p = (n,)
                x = x_range[0] + (x_range[1] - x_range[0])*rng.random(nx)
                x[0] = x_range[0]  # always include domain start point
                x[1] = x_range[1]  # always include domain end point
                kw = dict(sig=(len(p)+1)*'d'+'->d')
                z = np.c_[np.tile(p, (nx,1)), x, func(*(p + (x,)), **kw)]
                dataset.append(z)

        dataset = np.concatenate(dataset, axis=0)

        def polyfunc(*p):
            p0 = p[0].astype(np.intp)
            p = (p0,) + p[1:]
            p0_type_char = p0.dtype.char
            kw = dict(sig=p0_type_char + (len(p)-1)*'d' + '->d')
            return func(*p, **kw)

        with np.errstate(all='raise'):
            ds = FuncData(polyfunc, dataset, list(range(len(param_ranges)+2)), -1,
                          rtol=rtol)
            ds.check()

    def test_jacobi(self):
        self.check_poly(_ufuncs.eval_jacobi,
                        param_ranges=[(-0.99, 10), (-0.99, 10)],
                        x_range=[-1, 1])

    def test_sh_jacobi(self):
        self.check_poly(_ufuncs.eval_sh_jacobi,
                        param_ranges=[(1, 10), (0, 1)], x_range=[0, 1])

    def test_gegenbauer(self):
        self.check_poly(_ufuncs.eval_gegenbauer,
                        param_ranges=[(-0.499, 10)], x_range=[-1, 1])

    def test_chebyt(self):
        self.check_poly(_ufuncs.eval_chebyt,
                        param_ranges=[], x_range=[-1, 1])

    def test_chebyu(self):
        self.check_poly(_ufuncs.eval_chebyu,
                        param_ranges=[], x_range=[-1, 1])

    def test_chebys(self):
        self.check_poly(_ufuncs.eval_chebys,
                        param_ranges=[], x_range=[-2, 2])

    def test_chebyc(self):
        self.check_poly(_ufuncs.eval_chebyc,
                        param_ranges=[], x_range=[-2, 2])

    def test_sh_chebyt(self):
        self.check_poly(_ufuncs.eval_sh_chebyt,
                        param_ranges=[], x_range=[0, 1])

    def test_sh_chebyu(self):
        self.check_poly(_ufuncs.eval_sh_chebyu,
                        param_ranges=[], x_range=[0, 1])

    def test_legendre(self):
        self.check_poly(_ufuncs.eval_legendre,
                        param_ranges=[], x_range=[-1, 1])

    def test_sh_legendre(self):
        self.check_poly(_ufuncs.eval_sh_legendre,
                        param_ranges=[], x_range=[0, 1])

    def test_genlaguerre(self):
        self.check_poly(_ufuncs.eval_genlaguerre,
                        param_ranges=[(-0.99, 10)], x_range=[0, 100])

    def test_laguerre(self):
        self.check_poly(_ufuncs.eval_laguerre,
                        param_ranges=[], x_range=[0, 100])

    def test_hermite(self):
        v = _ufuncs.eval_hermite(70, 1.0)
        a = -1.457076485701412e60
        assert_allclose(v, a)


def test_hermite_domain():
    # Regression test for gh-11091.
    assert np.isnan(_ufuncs.eval_hermite(-1, 1.0))
    assert np.isnan(_ufuncs.eval_hermitenorm(-1, 1.0))


@pytest.mark.parametrize("n", [0, 1, 2])
@pytest.mark.parametrize("x", [0, 1, np.nan])
def test_hermite_nan(n, x):
    # Regression test for gh-11369.
    assert np.isnan(_ufuncs.eval_hermite(n, x)) == np.any(np.isnan([n, x]))
    assert np.isnan(_ufuncs.eval_hermitenorm(n, x)) == np.any(np.isnan([n, x]))


@pytest.mark.parametrize('n', [0, 1, 2, 3.2])
@pytest.mark.parametrize('alpha', [1, np.nan])
@pytest.mark.parametrize('x', [2, np.nan])
def test_genlaguerre_nan(n, alpha, x):
    # Regression test for gh-11361.
    nan_laguerre = np.isnan(_ufuncs.eval_genlaguerre(n, alpha, x))
    nan_arg = np.any(np.isnan([n, alpha, x]))
    assert nan_laguerre == nan_arg


@pytest.mark.parametrize('n', [0, 1, 2, 3.2])
@pytest.mark.parametrize('alpha', [0.0, 1, np.nan])
@pytest.mark.parametrize('x', [1e-6, 2, np.nan])
def test_gegenbauer_nan(n, alpha, x):
    # Regression test for gh-11370.
    nan_gegenbauer = np.isnan(_ufuncs.eval_gegenbauer(n, alpha, x))
    nan_arg = np.any(np.isnan([n, alpha, x]))
    assert nan_gegenbauer == nan_arg

@pytest.mark.parametrize("n", np.arange(10))
@pytest.mark.parametrize("alpha", [-0.45, -0.25, -0.125, 0.0, 0.25, 1.0, 2.0])
@pytest.mark.parametrize("x", [-np.inf, np.inf])
def test_gegenbauer_infinity(n, alpha, x):
    # gh-11713 - check correct handling of x = +inf and x = -inf
    if alpha == 0.0:
        expected = 0.0
    elif n == 0.0:
        expected = 1.0
    else:
        # sign of leading coefficient: 2^n * Gamma(n+alpha) / (n! Gamma(alpha))
        lead_sign = gammasgn(n + alpha) * gammasgn(alpha)
        expected = lead_sign * (np.sign(x) ** n) * np.inf
    assert_allclose(_ufuncs.eval_gegenbauer(int(n), alpha, x), expected, rtol=1e-10)
    assert_allclose(_ufuncs.eval_gegenbauer(float(n), alpha, x), expected, rtol=1e-10)

@pytest.mark.parametrize(
    "n, expected",
    [
        (0, [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]),
        (1, [-2.0, -1.8, -1.6, -1.4, -1.2, -1.0, -0.8, -0.6, -0.4, -0.2, 0.0]),
        (2, [3.0, 2.16, 1.44, 0.84, 0.36, 0.0, -0.24, -0.36, -0.36, -0.24, 0.0]),
        (3, [-4.0, -1.98, -0.64, 0.14, 0.48, 0.5, 0.32, 0.06, -0.16, -0.22, 0.0]),
        (
            4,
            [5.0, 1.332, -0.288, -0.658, -0.408, 0.0, 0.272, 0.282, 0.072, -0.148, 0.0],
        ),
        (
            5,
            [
                -6.0,
                -0.43308,
                0.79104,
                0.36876,
                -0.21312,
                -0.375,
                -0.14208,
                0.15804,
                0.19776,
                -0.04812,
                0.0,
            ],
        ),
    ],
)
def test_jacobi_alpha_minus_one_beta_plus_one(n, expected):
    # gh-7001 - expected values were computed with mathematica
    x = np.linspace(-1.0, 1.0, 11)
    a, b = -1, 1  # alpha, beta
    assert_allclose(_ufuncs.eval_jacobi(n, a, b, x), expected, rtol=1e-10, atol=1e-14)
    assert_allclose(
        _ufuncs.eval_jacobi(float(n), a, b, x), expected, rtol=1e-10, atol=1e-14
    )


@pytest.mark.parametrize(
    "n, expected",
    [
        (0, [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]),
        (1, [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
        (2, [0.0, -0.09, -0.16, -0.21, -0.24, -0.25, -0.24, -0.21, -0.16, -0.09, 0.0]),
        (
            3,
            [0.0, 0.144, 0.192, 0.168, 0.096, 0.0, -0.096, -0.168, -0.192, -0.144, 0.0],
        ),
        (
            4,
            [
                0.0,
                -0.1485,
                -0.096,
                0.0315,
                0.144,
                0.1875,
                0.144,
                0.0315,
                -0.096,
                -0.1485,
                0.0,
            ],
        ),
        (
            5,
            [
                0.0,
                0.10656,
                -0.04608,
                -0.15792,
                -0.13056,
                0.0,
                0.13056,
                0.15792,
                0.04608,
                -0.10656,
                0.0,
            ],
        ),
    ],
)
def test_jacobi_alpha_minus_one_beta_minus_one(n, expected):
    # gh-7001 - expected values were computed with mathematica
    x = np.linspace(-1.0, 1.0, 11)
    a, b = -1, -1  # alpha, beta
    assert_allclose(_ufuncs.eval_jacobi(n, a, b, x), expected, rtol=1e-10, atol=1e-14)
    assert_allclose(
        _ufuncs.eval_jacobi(float(n), a, b, x), expected, rtol=1e-10, atol=1e-14
    )