"""
A collection of functions to find the weights and abscissas for
Gaussian Quadrature.

These calculations are done by finding the eigenvalues of a
tridiagonal matrix whose entries are dependent on the coefficients
in the recursion formula for the orthogonal polynomials with the
corresponding weighting function over the interval.

Many recursion relations for orthogonal polynomials are given:

.. math::

    a1n f_{n+1} (x) = (a2n + a3n x ) f_n (x) - a4n f_{n-1} (x)

The recursion relation of interest is

.. math::

    P_{n+1} (x) = (x - A_n) P_n (x) - B_n P_{n-1} (x)

where :math:`P` has a different normalization than :math:`f`.

The coefficients can be found as:

.. math::

    A_n = -a2n / a3n
    \\qquad
    B_n = ( a4n / a3n \\sqrt{h_n-1 / h_n})^2

where

.. math::

    h_n = \\int_a^b w(x) f_n(x)^2

assume:

.. math::

    P_0 (x) = 1
    \\qquad
    P_{-1} (x) == 0

For the mathematical background, see [golub.welsch-1969-mathcomp]_ and
[abramowitz.stegun-1965]_.

References
----------
.. [golub.welsch-1969-mathcomp]
   Golub, Gene H, and John H Welsch. 1969. Calculation of Gauss
   Quadrature Rules. *Mathematics of Computation* 23, 221-230+s1--s10.

.. [abramowitz.stegun-1965]
   Abramowitz, Milton, and Irene A Stegun. (1965) *Handbook of
   Mathematical Functions: with Formulas, Graphs, and Mathematical
   Tables*. Gaithersburg, MD: National Bureau of Standards.
   http://www.math.sfu.ca/~cbm/aands/

.. [townsend.trogdon.olver-2014]
   Townsend, A. and Trogdon, T. and Olver, S. (2014)
   *Fast computation of Gauss quadrature nodes and
   weights on the whole real line*. :arXiv:`1410.5286`.

.. [townsend.trogdon.olver-2015]
   Townsend, A. and Trogdon, T. and Olver, S. (2015)
   *Fast computation of Gauss quadrature nodes and
   weights on the whole real line*.
   IMA Journal of Numerical Analysis
   :doi:`10.1093/imanum/drv002`.
"""
#
# Author:  Travis Oliphant 2000
# Updated Sep. 2003 (fixed bugs --- tested to be accurate)

from __future__ import division, print_function, absolute_import

# SciPy imports.
import numpy as np
from numpy import (exp, inf, pi, sqrt, floor, sin, cos, around, int,
                   hstack, arccos, arange)
from scipy import linalg
from scipy.special import airy

# Local imports.
from . import _ufuncs
from . import _ufuncs as cephes
_gam = cephes.gamma
from . import specfun

_polyfuns = ['legendre', 'chebyt', 'chebyu', 'chebyc', 'chebys',
             'jacobi', 'laguerre', 'genlaguerre', 'hermite',
             'hermitenorm', 'gegenbauer', 'sh_legendre', 'sh_chebyt',
             'sh_chebyu', 'sh_jacobi']

# Correspondence between new and old names of root functions
_rootfuns_map = {'roots_legendre': 'p_roots',
               'roots_chebyt': 't_roots',
               'roots_chebyu': 'u_roots',
               'roots_chebyc': 'c_roots',
               'roots_chebys': 's_roots',
               'roots_jacobi': 'j_roots',
               'roots_laguerre': 'l_roots',
               'roots_genlaguerre': 'la_roots',
               'roots_hermite': 'h_roots',
               'roots_hermitenorm': 'he_roots',
               'roots_gegenbauer': 'cg_roots',
               'roots_sh_legendre': 'ps_roots',
               'roots_sh_chebyt': 'ts_roots',
               'roots_sh_chebyu': 'us_roots',
               'roots_sh_jacobi': 'js_roots'}

_evalfuns = ['eval_legendre', 'eval_chebyt', 'eval_chebyu',
             'eval_chebyc', 'eval_chebys', 'eval_jacobi',
             'eval_laguerre', 'eval_genlaguerre', 'eval_hermite',
             'eval_hermitenorm', 'eval_gegenbauer',
             'eval_sh_legendre', 'eval_sh_chebyt', 'eval_sh_chebyu',
             'eval_sh_jacobi']

__all__ = _polyfuns + list(_rootfuns_map.keys()) + _evalfuns + ['poch', 'binom']


class orthopoly1d(np.poly1d):

    def __init__(self, roots, weights=None, hn=1.0, kn=1.0, wfunc=None,
                 limits=None, monic=False, eval_func=None):
        equiv_weights = [weights[k] / wfunc(roots[k]) for
                         k in range(len(roots))]
        mu = sqrt(hn)
        if monic:
            evf = eval_func
            if evf:
                knn = kn
                eval_func = lambda x: evf(x) / knn
            mu = mu / abs(kn)
            kn = 1.0

        # compute coefficients from roots, then scale
        poly = np.poly1d(roots, r=True)
        np.poly1d.__init__(self, poly.coeffs * float(kn))

        # TODO: In numpy 1.13, there is no need to use __dict__ to access attributes
        self.__dict__['weights'] = np.array(list(zip(roots,
                                                     weights, equiv_weights)))
        self.__dict__['weight_func'] = wfunc
        self.__dict__['limits'] = limits
        self.__dict__['normcoef'] = mu

        # Note: eval_func will be discarded on arithmetic
        self.__dict__['_eval_func'] = eval_func

    def __call__(self, v):
        if self._eval_func and not isinstance(v, np.poly1d):
            return self._eval_func(v)
        else:
            return np.poly1d.__call__(self, v)

    def _scale(self, p):
        if p == 1.0:
            return
        try:
            self._coeffs
        except AttributeError:
            self.__dict__['coeffs'] *= p
        else:
            # the coeffs attr is be made private in future versions of numpy
            self._coeffs *= p

        evf = self._eval_func
        if evf:
            self.__dict__['_eval_func'] = lambda x: evf(x) * p
        self.__dict__['normcoef'] *= p


def _gen_roots_and_weights(n, mu0, an_func, bn_func, f, df, symmetrize, mu):
    """[x,w] = gen_roots_and_weights(n,an_func,sqrt_bn_func,mu)

    Returns the roots (x) of an nth order orthogonal polynomial,
    and weights (w) to use in appropriate Gaussian quadrature with that
    orthogonal polynomial.

    The polynomials have the recurrence relation
          P_n+1(x) = (x - A_n) P_n(x) - B_n P_n-1(x)

    an_func(n)          should return A_n
    sqrt_bn_func(n)     should return sqrt(B_n)
    mu ( = h_0 )        is the integral of the weight over the orthogonal
                        interval
    """
    k = np.arange(n, dtype='d')
    c = np.zeros((2, n))
    c[0,1:] = bn_func(k[1:])
    c[1,:] = an_func(k)
    x = linalg.eigvals_banded(c, overwrite_a_band=True)

    # improve roots by one application of Newton's method
    y = f(n, x)
    dy = df(n, x)
    x -= y/dy

    fm = f(n-1, x)
    fm /= np.abs(fm).max()
    dy /= np.abs(dy).max()
    w = 1.0 / (fm * dy)

    if symmetrize:
        w = (w + w[::-1]) / 2
        x = (x - x[::-1]) / 2

    w *= mu0 / w.sum()

    if mu:
        return x, w, mu0
    else:
        return x, w

# Jacobi Polynomials 1               P^(alpha,beta)_n(x)


def roots_jacobi(n, alpha, beta, mu=False):
    r"""Gauss-Jacobi quadrature.

    Computes the sample points and weights for Gauss-Jacobi quadrature. The
    sample points are the roots of the n-th degree Jacobi polynomial,
    :math:`P^{\alpha, \beta}_n(x)`.  These sample points and weights
    correctly integrate polynomials of degree :math:`2n - 1` or less over the
    interval :math:`[-1, 1]` with weight function
    :math:`f(x) = (1 - x)^{\alpha} (1 + x)^{\beta}`.

    Parameters
    ----------
    n : int
        quadrature order
    alpha : float
        alpha must be > -1
    beta : float
        beta must be > -1
    mu : bool, optional
        If True, return the sum of the weights, optional.

    Returns
    -------
    x : ndarray
        Sample points
    w : ndarray
        Weights
    mu : float
        Sum of the weights

    See Also
    --------
    scipy.integrate.quadrature
    scipy.integrate.fixed_quad
    """
    m = int(n)
    if n < 1 or n != m:
        raise ValueError("n must be a positive integer.")
    if alpha <= -1 or beta <= -1:
        raise ValueError("alpha and beta must be greater than -1.")

    if alpha == 0.0 and beta == 0.0:
        return roots_legendre(m, mu)
    if alpha == beta:
        return roots_gegenbauer(m, alpha+0.5, mu)

    mu0 = 2.0**(alpha+beta+1)*cephes.beta(alpha+1, beta+1)
    a = alpha
    b = beta
    if a + b == 0.0:
        an_func = lambda k: np.where(k == 0, (b-a)/(2+a+b), 0.0)
    else:
        an_func = lambda k: np.where(k == 0, (b-a)/(2+a+b),
                  (b*b - a*a) / ((2.0*k+a+b)*(2.0*k+a+b+2)))

    bn_func = lambda k: 2.0 / (2.0*k+a+b)*np.sqrt((k+a)*(k+b) / (2*k+a+b+1)) \
              * np.where(k == 1, 1.0, np.sqrt(k*(k+a+b) / (2.0*k+a+b-1)))

    f = lambda n, x: cephes.eval_jacobi(n, a, b, x)
    df = lambda n, x: 0.5 * (n + a + b + 1) \
                      * cephes.eval_jacobi(n-1, a+1, b+1, x)
    return _gen_roots_and_weights(m, mu0, an_func, bn_func, f, df, False, mu)


def jacobi(n, alpha, beta, monic=False):
    r"""Jacobi polynomial.

    Defined to be the solution of

    .. math::
        (1 - x^2)\frac{d^2}{dx^2}P_n^{(\alpha, \beta)}
          + (\beta - \alpha - (\alpha + \beta + 2)x)
            \frac{d}{dx}P_n^{(\alpha, \beta)}
          + n(n + \alpha + \beta + 1)P_n^{(\alpha, \beta)} = 0

    for :math:`\alpha, \beta > -1`; :math:`P_n^{(\alpha, \beta)}` is a
    polynomial of degree :math:`n`.

    Parameters
    ----------
    n : int
        Degree of the polynomial.
    alpha : float
        Parameter, must be greater than -1.
    beta : float
        Parameter, must be greater than -1.
    monic : bool, optional
        If `True`, scale the leading coefficient to be 1. Default is
        `False`.

    Returns
    -------
    P : orthopoly1d
        Jacobi polynomial.

    Notes
    -----
    For fixed :math:`\alpha, \beta`, the polynomials
    :math:`P_n^{(\alpha, \beta)}` are orthogonal over :math:`[-1, 1]`
    with weight function :math:`(1 - x)^\alpha(1 + x)^\beta`.

    """
    if n < 0:
        raise ValueError("n must be nonnegative.")

    wfunc = lambda x: (1 - x)**alpha * (1 + x)**beta
    if n == 0:
        return orthopoly1d([], [], 1.0, 1.0, wfunc, (-1, 1), monic,
                           eval_func=np.ones_like)
    x, w, mu = roots_jacobi(n, alpha, beta, mu=True)
    ab1 = alpha + beta + 1.0
    hn = 2**ab1 / (2 * n + ab1) * _gam(n + alpha + 1)
    hn *= _gam(n + beta + 1.0) / _gam(n + 1) / _gam(n + ab1)
    kn = _gam(2 * n + ab1) / 2.0**n / _gam(n + 1) / _gam(n + ab1)
    # here kn = coefficient on x^n term
    p = orthopoly1d(x, w, hn, kn, wfunc, (-1, 1), monic,
                    lambda x: eval_jacobi(n, alpha, beta, x))
    return p

# Jacobi Polynomials shifted         G_n(p,q,x)


def roots_sh_jacobi(n, p1, q1, mu=False):
    """Gauss-Jacobi (shifted) quadrature.