"""
This module implements some special functions that commonly appear in
combinatorial contexts (e.g. in power series); in particular,
sequences of rational numbers such as Bernoulli and Fibonacci numbers.

Factorials, binomial coefficients and related functions are located in
the separate 'factorials' module.
"""

from sympy import Function, S, Symbol, Rational, oo, Integer, C, Add, expand_mul

from sympy.mpmath import bernfrac
from sympy.mpmath.libmp import ifib as _ifib

def _product(a, b):
    p = 1
    for k in range(a, b+1):
        p *= k
    return p

from sympy.utilities.memoization import recurrence_memo


# Dummy symbol used for computing polynomial sequences
_sym = Symbol('x')
_symbols = Function('x')


#----------------------------------------------------------------------------#
#                                                                            #
#                           Fibonacci numbers                                #
#                                                                            #
#----------------------------------------------------------------------------#

class fibonacci(Function):
    """
    Fibonacci numbers / Fibonacci polynomials

    The Fibonacci numbers are the integer sequence defined by the
    initial terms F_0 = 0, F_1 = 1 and the two-term recurrence
    relation F_n = F_{n-1} + F_{n-2}.

    The Fibonacci polynomials are defined by F_1(x) = 1,
    F_2(x) = x, and F_n(x) = x*F_{n-1}(x) + F_{n-2}(x) for n > 2.
    For all positive integers n, F_n(1) = F_n.

    * fibonacci(n) gives the nth Fibonacci number, F_n
    * fibonacci(n, x) gives the nth Fibonacci polynomial in x, F_n(x)

    Examples
    ========

    >>> from sympy import fibonacci, Symbol

    >>> [fibonacci(x) for x in range(11)]
    [0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55]
    >>> fibonacci(5, Symbol('t'))
    t**4 + 3*t**2 + 1

    References
    ==========

    * http://en.wikipedia.org/wiki/Fibonacci_number
    * http://mathworld.wolfram.com/FibonacciNumber.html

    See Also
    ========

    lucas
    """

    @staticmethod
    def _fib(n):
        return _ifib(n)

    @staticmethod
    @recurrence_memo([None, S.One, _sym])
    def _fibpoly(n, prev):
        return (prev[-2] + _sym*prev[-1]).expand()

    @classmethod
    def eval(cls, n, sym=None):
        if n.is_Integer:
            n = int(n)
            if n < 0:
                return S.NegativeOne**(n+1) * fibonacci(-n)
            if sym is None:
                return Integer(cls._fib(n))
            else:
                if n < 1:
                    raise ValueError("Fibonacci polynomials are defined "
                       "only for positive integer indices.")
                return cls._fibpoly(n).subs(_sym, sym)

class lucas(Function):
    """
    Lucas numbers

    Lucas numbers satisfy a recurrence relation similar to that of
    the Fibonacci sequence, in which each term is the sum of the
    preceding two. They are generated by choosing the initial
    values L_0 = 2 and L_1 = 1.

    * lucas(n) gives the nth Lucas number

    Examples
    ========

    >>> from sympy import lucas

    >>> [lucas(x) for x in range(11)]
    [2, 1, 3, 4, 7, 11, 18, 29, 47, 76, 123]

    References
    ==========

    * http://en.wikipedia.org/wiki/Lucas_number

    See Also
    ========

    fibonacci
    """

    @classmethod
    def eval(cls, n):
        if n.is_Integer:
            return fibonacci(n+1) + fibonacci(n-1)


#----------------------------------------------------------------------------#
#                                                                            #
#                           Bernoulli numbers                                #
#                                                                            #
#----------------------------------------------------------------------------#

class bernoulli(Function):
    r"""
    Bernoulli numbers / Bernoulli polynomials

    The Bernoulli numbers are a sequence of rational numbers
    defined by B_0 = 1 and the recursive relation (n > 0)::

                n
               ___
              \      / n + 1 \
          0 =  )     |       | * B .
              /___   \   k   /    k
              k = 0

    They are also commonly defined by their exponential generating
    function, which is x/(exp(x) - 1). For odd indices > 1, the
    Bernoulli numbers are zero.

    The Bernoulli polynomials satisfy the analogous formula::

                    n
                   ___
                  \      / n \         n-k
          B (x) =  )     |   | * B  * x   .
           n      /___   \ k /    k
                  k = 0

    Bernoulli numbers and Bernoulli polynomials are related as
    B_n(0) = B_n.

    We compute Bernoulli numbers using Ramanujan's formula::

                                   / n + 3 \
          B   =  (A(n) - S(n))  /  |       |
           n                       \   n   /

    where A(n) = (n+3)/3 when n = 0 or 2 (mod 6), A(n) = -(n+3)/6
    when n = 4 (mod 6), and::

                 [n/6]
                  ___
                 \      /  n + 3  \
          S(n) =  )     |         | * B
                 /___   \ n - 6*k /    n-6*k
                 k = 1

    This formula is similar to the sum given in the definition, but
    cuts 2/3 of the terms. For Bernoulli polynomials, we use the
    formula in the definition.

    * bernoulli(n) gives the nth Bernoulli number, B_n
    * bernoulli(n, x) gives the nth Bernoulli polynomial in x, B_n(x)

    Examples
    ========

    >>> from sympy import bernoulli

    >>> [bernoulli(n) for n in range(11)]
    [1, -1/2, 1/6, 0, -1/30, 0, 1/42, 0, -1/30, 0, 5/66]
    >>> bernoulli(1000001)
    0

    References
    ==========

    * http://en.wikipedia.org/wiki/Bernoulli_number
    * http://en.wikipedia.org/wiki/Bernoulli_polynomial

    See Also
    ========

    euler, bell
    """

    # Calculates B_n for positive even n
    @staticmethod
    def _calc_bernoulli(n):
        s = 0
        a = int(C.binomial(n+3, n-6))
        for j in range(1, n//6+1):
            s += a * bernoulli(n - 6*j)
            # Avoid computing each binomial coefficient from scratch
            a *= _product(n-6 - 6*j + 1, n-6*j)
            a //= _product(6*j+4, 6*j+9)
        if n % 6 == 4:
            s = -Rational(n+3, 6) - s
        else:
            s = Rational(n+3, 3) - s
        return s / C.binomial(n+3, n)

    # We implement a specialized memoization scheme to handle each
    # case modulo 6 separately
    _cache = {0: S.One, 2:Rational(1,6), 4:Rational(-1,30)}
    _highest = {0:0, 2:2, 4:4}

    @classmethod
    def eval(cls, n, sym=None):
        if n.is_Number:
            if n.is_Integer and n.is_nonnegative:
                if n is S.Zero:
                    return S.One
                elif n is S.One:
                    if sym is None: return -S.Half
                    else:           return sym - S.Half
                # Bernoulli numbers
                elif sym is None:
                    if n.is_odd:
                        return S.Zero
                    n = int(n)
                    # Use mpmath for enormous Bernoulli numbers
                    if n > 500:
                        p, q = bernfrac(n)
                        return Rational(int(p), int(q))
                    case = n % 6
                    highest_cached = cls._highest[case]
                    if n <= highest_cached:
                        return cls._cache[n]
                    # To avoid excessive recursion when, say, bernoulli(1000) is
                    # requested, calculate and cache the entire sequence ... B_988,
                    # B_994, B_1000 in increasing order
                    for i in range(highest_cached+6, n+6, 6):
                        b = cls._calc_bernoulli(i)
                        cls._cache[i] = b
                        cls._highest[case] = i
                    return b
                # Bernoulli polynomials
                else:
                    n, result = int(n), []
                    for k in range(n + 1):
                        result.append(C.binomial(n, k)*cls(k)*sym**(n-k))
                    return Add(*result)
            else:
                raise ValueError("Bernoulli numbers are defined only"
                                 " for nonnegative integer indices.")


#----------------------------------------------------------------------------#
#                                                                            #
#                             Bell numbers                                   #
#                                                                            #
#----------------------------------------------------------------------------#

class bell(Function):
    r"""
    Bell numbers / Bell polynomials

    The Bell numbers satisfy `B_0 = 1` and

    .. math:: B_n = \sum_{k=0}^{n-1} \binom{n-1}{k} B_k.

    They are also given by:

    .. math:: B_n = \frac{1}{e} \sum_{k=0}^{\infty} \frac{k^n}{k!}.

    The Bell polynomials are given by `B_0(x) = 1` and

    .. math:: B_n(x) = x \sum_{k=1}^{n-1} \binom{n-1}{k-1} B_{k-1}(x).

    The second kind of Bell polynomials (are sometimes called "partial" Bell
    polynomials or incomplete Bell polynomials) are defined as

    .. math:: B_{n,k}(x_1, x_2,\dotsc x_{n-k+1}) =
            \sum_{j_1+j_2+j_2+\dotsb=k \atop j_1+2j_2+3j_2+\dotsb=n}
                \frac{n!}{j_1!j_2!\dotsb j_{n-k+1}!}
                \left(\frac{x_1}{1!} \right)^{j_1}
                \left(\frac{x_2}{2!} \right)^{j_2} \dotsb
                \left(\frac{x_{n-k+1}}{(n-k+1)!} \right) ^{j_{n-k+1}}.

    * bell(n) gives the `n^{th}` Bell number, `B_n`.
    * bell(n, x) gives the `n^{th}` Bell polynomial, `B_n(x)`.
    * bell(n, k, (x1, x2, ...)) gives Bell polynomials of the second kind,
      `B_{n,k}(x_1, x_2, \dotsc, x_{n-k+1})`.

    Notes
    =====

    Not to be confused with Bernoulli numbers and Bernoulli polynomials,
    which use the same notation.

    Examples
    ========

    >>> from sympy import bell, Symbol, symbols

    >>> [bell(n) for n in range(11)]
    [1, 1, 2, 5, 15, 52, 203, 877, 4140, 21147, 115975]
    >>> bell(30)
    846749014511809332450147
    >>> bell(4, Symbol('t'))
    t**4 + 6*t**3 + 7*t**2 + t
    >>> bell(6, 2, symbols('x:6')[1:])
    6*x1*x5 + 15*x2*x4 + 10*x3**2

    References
    ==========

    * http://en.wikipedia.org/wiki/Bell_number
    * http://mathworld.wolfram.com/BellNumber.html
    * http://mathworld.wolfram.com/BellPolynomial.html

    See Also
    ========

    euler, bernoulli
    """

    @staticmethod
    @recurrence_memo([1, 1])
    def _bell(n, prev):
        s = 1
        a = 1
        for k in range(1, n):
            a = a * (n-k) // k
            s += a * prev[k]
        return s

    @staticmethod
    @recurrence_memo([S.One, _sym])
    def _bell_poly(n, prev):
        s = 1
        a = 1
        for k in range(2, n+1):
            a = a * (n-k+1) // (k-1)
            s += a * prev[k-1]
        return expand_mul(_sym * s)

    @staticmethod
    def _bell_incomplete_poly(n, k, symbols):
        r"""
        The second kind of Bell polynomials (incomplete Bell polynomials).

        Calculated by recurrence formula:

        .. math:: B_{n,k}(x_1, x_2, \dotsc, x_{n-k+1}) =
                \sum_{m=1}^{n-k+1}
                \x_m \binom{n-1}{m-1} B_{n-m,k-1}(x_1, x_2, \dotsc, x_{n-m-k})

        where
            B_{0,0} = 1;
            B_{n,0} = 0; for n>=1
            B_{0,k} = 0; for k>=1

        """
        if (n==0) and (k==0):
            return S.One
        elif (n==0) or (k==0):
            return S.Zero
        s = S.Zero
        a = S.One
        for m in range(1, n-k+2):
            s += a*bell._bell_incomplete_poly(n-m, k-1, symbols)*symbols[m-1]
            a = a*(n-m)/m
        return expand_mul(s)

    @classmethod
    def eval(cls, n, k_sym=None, symbols=None):
        if n.is_Integer and n.is_nonnegative:
            if k_sym is None:
                return Integer(cls._bell(int(n)))
            elif symbols is None:
                return cls._bell_poly(int(n)).subs(_sym, k_sym)
            else:
                r = cls._bell_incomplete_poly(int(n), int(k_sym), symbols)
                return r

#----------------------------------------------------------------------------#
#                                                                            #
#                           Harmonic numbers                                 #
#                                                                            #
#----------------------------------------------------------------------------#

class harmonic(Function):
    r"""
    Harmonic numbers

    The nth harmonic number is given by 1 + 1/2 + 1/3 + ... + 1/n.

    More generally::

                   n
                  ___
                 \       -m
          H    =  )     k   .
           n,m   /___
                 k = 1

    As n -> oo, H_{n,m} -> zeta(m) (the Riemann zeta function)

    * harmonic(n) gives the nth harmonic number, H_n

    * harmonic(n, m) gives the nth generalized harmonic number
      of order m, H_{n,m}, where harmonic(n) == harmonic(n, 1)

    Examples
    ========

    >>> from sympy import harmonic, oo

    >>> [harmonic(n) for n in range(6)]
    [0, 1, 3/2, 11/6, 25/12, 137/60]
    >>> [harmonic(n, 2) for n in range(6)]
    [0, 1, 5/4, 49/36, 205/144, 5269/3600]
    >>> harmonic(oo, 2)
    pi**2/6

    References
    ==========

    * http://en.wikipedia.org/wiki/Harmonic_number

    """

    # Generate one memoized Harmonic number-generating function for each
    # order and store it in a dictionary
    _functions = {}

    nargs = (1, 2)

    @classmethod
    def eval(cls, n, m=None):
        if m is None:
            m = S.One
        if n == oo:
            return C.zeta(m)
        if n.is_Integer and n.is_nonnegative and m.is_Integer:
            if n == 0:
                return S.Zero
            if not m in cls._functions:
                @recurrence_memo([0])
                def f(n, prev):
                    return prev[-1] + S.One / n**m
                cls._functions[m] = f
            return cls._functions[m](int(n))

#----------------------------------------------------------------------------#
#                                                                            #
#                           Euler numbers                                    #
#                                                                            #
#----------------------------------------------------------------------------#

class euler(Function):
    r"""
    Euler numbers

    The euler numbers are given by::

                  2*n+1   k
                   ___   ___            j          2*n+1
                  \     \     / k \ (-1)  * (k-2*j)
          E   = I  )     )    |   | --------------------
           2n     /___  /___  \ j /      k    k
                  k = 1 j = 0           2  * I  * k

          E     = 0
           2n+1

    * euler(n) gives the n-th Euler number, E_n

    Examples
    ========

    >>> from sympy import Symbol, euler
    >>> [euler(n) for n in range(10)]
    [1, 0, -1, 0, 5, 0, -61, 0, 1385, 0]
    >>> n = Symbol("n")
    >>> euler(n+2*n)
    euler(3*n)

    References
    ==========

    * http://en.wikipedia.org/wiki/Euler_numbers
    * http://mathworld.wolfram.com/EulerNumber.html
    * http://en.wikipedia.org/wiki/Alternating_permutation
    * http://mathworld.wolfram.com/AlternatingPermutation.html

    See Also
    ========

    bernoulli, bell
    """

    nargs = 1

    @classmethod
    def eval(cls, m, evaluate=True):
        if not evaluate:
            return
        if m.is_odd:
            return S.Zero
        if m.is_Integer and m.is_nonnegative:
            from sympy.mpmath import mp
            m = m._to_mpmath(mp.prec)
            res = mp.eulernum(m, exact=True)
            return Integer(res)


    def _eval_rewrite_as_Sum(self, arg):
        if arg.is_even:
            k = C.Dummy("k", integer=True)
            j = C.Dummy("j", integer=True)
            n = self.args[0] / 2
            Em = (S.ImaginaryUnit * C.Sum( C.Sum( C.binomial(k,j) * ((-1)**j * (k-2*j)**(2*n+1)) /
                  (2**k*S.ImaginaryUnit**k * k), (j,0,k)), (k, 1, 2*n+1)))

            return Em


    def _eval_evalf(self, prec):
        m = self.args[0]

        if m.is_Integer and m.is_nonnegative:
            from sympy.mpmath import mp
            from sympy import Expr
            m = m._to_mpmath(prec)
            oprec = mp.prec
            mp.prec = prec
            res = mp.eulernum(m)
            mp.prec = oprec
            return Expr._from_mpmath(res, prec)

#----------------------------------------------------------------------------#
#                                                                            #
#                           Catalan numbers                                  #
#                                                                            #
#----------------------------------------------------------------------------#

class catalan(Function):
    r"""
    Catalan numbers

    The n-th catalan number is given by::

                 1   / 2*n \
          C  = ----- |     |
           n   n + 1 \  n  /

    * catalan(n) gives the n-th Catalan number, C_n

    Examples
    ========

    >>> from sympy import (Symbol, binomial, gamma, hyper, polygamma,
    ...             catalan, diff, combsimp, Rational, I)

    >>> [ catalan(i) for i in range(1,10) ]
    [1, 2, 5, 14, 42, 132, 429, 1430, 4862]

    >>> n = Symbol("n", integer=True)

    >>> catalan(n)
    catalan(n)

    Catalan numbers can be transformed into several other, identical
    expressions involving other mathematical functions

    >>> catalan(n).rewrite(binomial)
    binomial(2*n, n)/(n + 1)

    >>> catalan(n).rewrite(gamma)
    4**n*gamma(n + 1/2)/(sqrt(pi)*gamma(n + 2))

    >>> catalan(n).rewrite(hyper)
    hyper((-n + 1, -n), (2,), 1)

    For some non-integer values of n we can get closed form
    expressions by rewriting in terms of gamma functions:

    >>> catalan(Rational(1,2)).rewrite(gamma)
    8/(3*pi)

    We can differentiate the Catalan numbers C(n) interpreted as a
    continuous real funtion in n:

    >>> diff(catalan(n), n)
    (polygamma(0, n + 1/2) - polygamma(0, n + 2) + 2*log(2))*catalan(n)

    As a more advanced example consider the following ratio
    between consecutive numbers:

    >>> combsimp((catalan(n + 1)/catalan(n)).rewrite(binomial))
    2*(2*n + 1)/(n + 2)

    The Catalan numbers can be generalized to complex numbers:

    >>> catalan(I).rewrite(gamma)
    4**I*gamma(1/2 + I)/(sqrt(pi)*gamma(2 + I))

    and evaluated with arbitrary precision:

    >>> catalan(I).evalf(20)
    0.39764993382373624267 - 0.020884341620842555705*I

    References
    ==========

    * http://en.wikipedia.org/wiki/Catalan_number
    * http://mathworld.wolfram.com/CatalanNumber.html
    * http://geometer.org/mathcircles/catalan.pdf

    See Also
    ========

    sympy.functions.combinatorial.factorials.binomial
    """

    @classmethod
    def eval(cls, n, evaluate=True):
        if n.is_Integer and n.is_nonnegative:
            return 4**n*C.gamma(n + S.Half)/(C.gamma(S.Half)*C.gamma(n+2))

    def fdiff(self, argindex=1):
        n = self.args[0]
        return catalan(n)*(C.polygamma(0,n+Rational(1,2))-C.polygamma(0,n+2)+C.log(4))

    def _eval_rewrite_as_binomial(self,n):
        return C.binomial(2*n,n)/(n + 1)

    def _eval_rewrite_as_gamma(self,n):
        # The gamma function allows to generalize Catalan numbers to complex n
        return 4**n*C.gamma(n + S.Half)/(C.gamma(S.Half)*C.gamma(n+2))

    def _eval_rewrite_as_hyper(self,n):
        return C.hyper([1-n,-n],[2],1)

    def _eval_evalf(self, prec):
        return self.rewrite(C.gamma).evalf(prec)
