# -*- coding: utf-8 -*-
# vi:ts=4:et
#
# $Date: 2003/08/15 12:29:29 $
# $Revision: 1.4 $
# =====================================================

"""library for permutation and combination"""

import sys
import operator


__all__ = [
           'factorial',
           'permutation',
           'combination',
           'catalan',
           ]

def factorial(first, stop=0):
    """Return the factorial.

    if second argument is given,
    return the multiply of sequence from n to stop.
    n x (n-1) x ... x stop
    """
    #factorial(5,2) -> 5*4*3*2
    #factorial(4,3) -> 4*3

    # 4! -> 4*3*2*1 = 24
    # 3! -> 3*2*1 = 6

    if stop:
        assert isinstance(stop, int) and stop >= 0, "second argument must be a non-negative integer."

    assert isinstance(first, int) and first >= 0, "first argument must be a non-negative integer."

    assert first >= stop, "first argument must be greater than or equal to second argument."

    if first < 2:
        return 1

    if not stop:
        # *1 is trivial, so set the second argument 2.
        stop = 2

    return _mul_of_sequence(range(stop, first + 1), 1)

def _mul_of_sequence(seq, initial=None):
    """Return the multipy of sequence.
    """
    # [sample]

    # if seq is [2,4,6]
    # _mul_of_sequence(seq) = 2 * 4 * 6 = 48

    if initial:
        return reduce(operator.mul, seq, initial)
    else:
        return reduce(operator.mul, seq)

def combination(first, second):
    """Return the combination.

    first argument can be a sequence or an integer.
    second argument must be an integer.

    if the first argument is an integer,
    return the combination of (n, m)

    if the first argument is a sequence,
    out of the sequence, take n elements.
    """
    assert isinstance(second, int) and second >= 0, "second argument must be a non-negative integer."

    if isinstance(first, int):
        return _combination_of_integer(first, second)

    # Now we know that the first argument is not an integer.
    # the first needs to be a sequence.

    try:
        len(first)
    except TypeError, e:
        raise TypeError, "bad operand type was found in the argument : %s"%first
    else:
        return _combination_of_sequence(first, second)

def _combination_of_sequence(seq, num):
    """out of the sequence, take num elements
    """
    if len(seq) > num:
        num = len(seq)
    #assert len(seq) >= num,  seq

    return list(com_gen(seq, num))

def _combination_of_integer(big, small):
    """combination of two integers.
    """
    # (4,2)  -> 4*3 / 2*1 = 6
    # (n, m) -> n!/(m! * (n-m)!)
    # = n*(n-1) * ... * (n-m+1)/m!
    assert big >= small >=0, "invalid argument (%s, %s)\n"%(big, small) + \
    " "*4 + "first argument must be greater than or equal to second argument."

    return factorial(big, big - small +1)/factorial(small)

def permutation(*seq):
    """Return the permutation of sequence
    """
    # XXX
    # by using arbitrary argument lists,
    # you can use expressions both
    # permutation(('a','b','c'))    # -> single argument of tuple
    # and
    # permutation('a', 'b', 'c')    # -> three arguments of string

    if len(seq) == 1:
        seq = seq[0]

    # list method, 'pop', is used in perm_gen
    if not isinstance(seq, list):
        try:
            seq = list(seq)
        except TypeError, e:
            raise TypeError, "can't convert to list: %r"%(repr(seq))

    return list(perm_gen(seq))

def perm_gen(seq):
    """generate all permutations from sequence.
    """
    # largely owes the idea from comp.lang.python community.

    def permhalf(seq):
        # permutation has a symmetric structure.
        # symmetric group <--> permutation
        pop, insert, append = seq.pop, seq.insert, seq.append
        llen = len(seq)
        if llen <= 2:
            yield seq
        else:
            aRange = range(llen)
            v = pop()
            for p in permhalf(seq):
                for j in aRange:
                    insert(j, v)
                    yield seq
                    del seq[j]
            append(v)

    ph = permhalf(seq)

    for h in ph:
        p = h[:]

        # XXX
        # don't forget to add [:] at yield statement.
        yield p[:]
        p.reverse()
        yield p[:]

# NOTE
# next function is included in the distribution.
# see the following script.
# dist/src/Lib/test/test_generators.py
# the original name was 'gcomb'
def com_gen(seq, k):
    """Generate all combinations of k elements from the seq.
    """
    if k > len(seq):
        return
    if k == 0:
        yield []
    else:
        first, rest = seq[0], seq[1:]
        # A combination does or doesn't contain first.
        # If it does, the remainder is a k-1 comb of rest.
        for c in com_gen(rest, k-1):
            c.insert(0, first)
            yield c
        # If it doesn't contain first, it's a k comb of rest.
        for c in com_gen(rest, k):
            yield c


def catalan(num):
    """Return the Catalan number
    """
    # just for fun

    # C_n = (1 / (n+1)) * (2n, n)
    #
    #     = (1 / (n+1)) * (2n)! / (n!)^2
    #
    #     = (2n)! / ((n+1)! * n!)
    #
    #     = (2n * (2n-1) * ... * (n+1) / (n+1) !
    #
    #     = (2n * (2n-1) * ... * (n+2) / n !

    assert isinstance(num, int) and num > 0, "must be a positive integer."
    if num == 1:
        return 1

    return factorial(2 * num, num+2) / factorial(num)

