from sage.all import *

def legendre_p(n, x, a = -1, b = 1):
    r"""
    Returns the ``n``th Legendre polynomial of the first kind over the
    interval [a, b] with respect to ``x``.

    When [a,b] is not [-1,1], we scale the standard Legendre
    polynomial (which is defined over [-1,1]) via an affine map. The
    resulting polynomials are still orthogonal and possess the
    property that `P(a) = P(b) = 1`.

    INPUT:

      * ``n`` -- The index of the polynomial.

      * ``x`` -- Either the variable to use as the independent
        variable in the polynomial, or a point at which to evaluate
        the polynomial.

      * ``a`` -- The "left" endpoint of the interval. Must be a real number.

      * ``b`` -- The "right" endpoint of the interval. Must be a real number.

    OUTPUT:

    If ``x`` is a variable, a polynomial (symbolic expression) will be
    returned. Otherwise, the value of the ``n``th polynomial at ``x``
    will be returned.

    SETUP::

        sage: from mjo.orthogonal_polynomials import legendre_p

    EXAMPLES:

    Create the standard Legendre polynomials in `x`::

        sage: legendre_p(0,x)
        1
        sage: legendre_p(1,x)
        x

    Reuse the variable from a polynomial ring::
        sage: P.<t> = QQ[]
        sage: legendre_p(2,t)
        3/2*t^2 - 1/2

    If ``x`` is a real number, the result should be as well::

        sage: legendre_p(3, 1.1)
        1.67750000000000

    Similarly for complex numbers::

        sage: legendre_p(3, 1 + I)
        7/2*I - 13/2

    Even matrices work::

        sage: legendre_p(3, MatrixSpace(ZZ, 2)([1, 2, -4, 7]))
        [-179  242]
        [-484  547]

    And finite field elements::

        sage: legendre_p(3, GF(11)(5))
        8

    Solve a simple least squares problem over `[-\pi, \pi]`::

        sage: a = -pi
        sage: b = pi
        sage: def inner_product(v1, v2):
        ....:     return integrate(v1*v2, x, a, b)
        sage: def norm(v):
        ....:     return sqrt(inner_product(v,v))
        sage: def project(basis, v):
        ....:     return sum( inner_product(v, b)*b/norm(b)**2
        ....:                 for b in basis)
        sage: f = sin(x)
        sage: legendre_basis = [ legendre_p(k, x, a, b) for k in range(4) ]
        sage: proj = project(legendre_basis, f)
        sage: proj.simplify_trig()
        5/2*(7*(pi^2 - 15)*x^3 - 3*(pi^4 - 21*pi^2)*x)/pi^6

    TESTS:

    We should agree with Maxima for all `n`::

        sage: eq = lambda k: bool(legendre_p(k,x) == legendre_P(k,x))
        sage: all( eq(k) for k in range(20) ) # long time
        True

    We can evaluate the result of the zeroth polynomial::

        sage: f = legendre_p(0,x)
        sage: f(x=10)
        1

    We should have |P(a)| = |P(b)| = 1 for all a,b::

        sage: a = RR.random_element()
        sage: b = RR.random_element()
        sage: k = ZZ.random_element(20)
        sage: P = legendre_p(k, x, a, b)
        sage: abs(P(x=a)) # abs tol 1e-12
        1
        sage: abs(P(x=b)) # abs tol 1e-12
        1

    Two different polynomials should be orthogonal with respect to the
    inner product over `[a,b]`. Note that this test can fail if QQ is
    replaced with RR, since integrate() can return NaN::

        sage: a = QQ.random_element()
        sage: b = QQ.random_element()
        sage: j = ZZ.random_element(20)
        sage: k = j + 1
        sage: Pj = legendre_p(j, x, a, b)
        sage: Pk = legendre_p(k, x, a, b)
        sage: integrate(Pj*Pk, x, a, b) # abs tol 1e-12
        0

    The first few polynomials shifted to [0,1] are known to be::

        sage: p0 = 1
        sage: p1 = 2*x - 1
        sage: p2 = 6*x^2 - 6*x + 1
        sage: p3 = 20*x^3 - 30*x^2 + 12*x - 1
        sage: bool(legendre_p(0, x, 0, 1) == p0)
        True
        sage: bool(legendre_p(1, x, 0, 1) == p1)
        True
        sage: bool(legendre_p(2, x, 0, 1) == p2)
        True
        sage: bool(legendre_p(3, x, 0, 1) == p3)
        True

    The zeroth polynomial is an element of the ring that we're working
    in::

        sage: legendre_p(0, MatrixSpace(ZZ, 2)([1, 2, -4, 7]))
        [1 0]
        [0 1]

    """
    if not n in ZZ:
        raise TypeError('n must be a natural number')

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

    if not (a in RR and b in RR):
        raise TypeError('both `a` and `b` must be real numbers')

    if n == 0:
        # Easy base case, save time. Attempt to return a value in the
        # same field/ring as `x`.
        return x.parent()(1)

    # Even though we know a,b are real we use a different ring. We
    # prefer ZZ so that we can support division of finite field
    # elements by (a-b). Eventually this should be supported for QQ as
    # well, although it does not work at the moment. The preference of
    # SR over RR is to return something attractive when e.g. a=pi.
    if a in ZZ:
        a = ZZ(a)
    else:
        a = SR(a)

    if b in ZZ:
        b = ZZ(b)
    else:
        b = SR(b)

    # Ensure that (2**n) is an element of ZZ. This is used later --
    # we can divide finite field elements by integers but we can't
    # multiply them by rationals at the moment.
    n = ZZ(n)

    def phi(t):
        # This is an affine map from [a,b] into [-1,1] and so
        # preserves orthogonality.
        return (a + b - 2*t)/(a - b)

    def c(m):
        return binomial(n,m)*binomial(n, n-m)

    def g(m):
        # As given in A&S, but with `x` replaced by `phi(x)`.
        return ( ((phi(x) - 1)**(n-m)) * (phi(x) + 1)**m )

    # From Abramowitz & Stegun, (22.3.2) with alpha = beta = 0.
    # Also massaged to support finite field elements.
    P = sum( c(m)*g(m) for m in range(n+1) )/(2**n)

    return P