#!/usr/bin/env python

"""
Some useful functions for number theory
"""

__all__ = ['divisors', 'is_prime', 'primes', 'pi', 'first_primes', 'prime_divisors']

def divisors(N):
    """divisors(N) -> tuple of divisors of the integer N
	Example:
	>>> divisors(2009)
	(7, 41, 49, 287, 2009)
	>>> divisors(2011)
	(2011,)
	>>>
"""
    l = list()
    for i in range(2, N+1):
        if N%i == 0:
            # i is a divisor!
            l.append(i)
    return tuple(l)

def is_prime(N):
    """is_prime(N) -> primality test of N
	Example:
	>>> is_prime(7)
	True
	>>> is_prime(9)
	False
	>>>
"""
    if N < 2:
        return False
    elif N == 2:
        return True
    for j in range(2, N):
        if N % j == 0:
            # i is not prime!
            return False
    else:
        return True

def next_prime(n):
    """next_prime(n) -> returns the smallest prime number p > n
	Example:
	>>> next_prime(0)
	2
	>>> next_prime(2)
	3
	>>> next_prime(7)
	11
	>>> next_prime(1000)
	1009
"""
    p = n + 1
    while not is_prime(p):
        p += 1
    return p

def prime_divisors(N):
    """prime_divisors(N) -> tuple of prime divisors of the integer N
	Example:
	>>> prime_divisors(2009)
	(7, 41)
	>>> prime_divisors(2011)
	(2011,)
	>>>
"""
    l = list()
    for d in divisors(N):
        if is_prime(d):
            l.append(d)
    return tuple(l)

def primes():
    """primes() -> generate all prime numbers
	Example:
>>> for i in primes():
...     if i > 10:
...          break
...     print i
...
2
3
5
7
>>>
"""
    yield 2
    i = 3
    while True:
        if is_prime(i):
            yield i
        i += 2

def pi(N):
    """pi(N) -> number of primes less than N
	Example:
>>> print pi(10)
4
>>>
"""
    count = 0
    for p in primes():
        if p >= N:
            break
        count += 1
    return count

def repeat(n, v=None):
    i = 0
    while i < n:
        yield v
        i += 1

def first_N_of(iterable, N):
    """first_N(N) -> iterate over the first N elements of iterable
	Example:
	>>> for p in first_N_of(primes(), 4):
	...     print p
	...
	2
	3
	5
	7
"""
    for dummy, v in zip(xrange(N), iterable):
        yield v

def first_primes(N):
    """first_primes(N) -> iterate over the first N primes
	Example:
	>>> for i in first_primes(4):
	...     print i
        ...
        2
	3
	5
	7
	>>>
"""
    return first_N_of(primes(), N)
    #for dummy, p in zip(repeat(N), primes()):
    #    yield p

def twin_primes():
    """twin_primes() -> generate all twin primes
	Example:
>>> for tp in first_N_of(twin_primes(), 6):
...     print tp
...
(3, 5)
(5, 7)
(11, 13)
(17, 19)
(29, 31)
(41, 43)
>>>
"""
    return de_polignac_primes(k=1)

def de_polignac_primes(k):
    """de_polignac_primes() -> generate all primes (p1, p2) | p2 - p1 = 2*k
	Example:
>>> for tp in first_N_of(de_polignac_primes(1), 6):
...     print tp
...
(3, 5)
(5, 7)
(11, 13)
(17, 19)
(29, 31)
(41, 43)
>>> for c, tp in enumerate(de_polignac_primes(7)):
...     if c > 5:
...          break
...     print tp
...
(113, 127)
(293, 307)
(317, 331)
(773, 787)
(839, 853)
(863, 877)
>>>
"""
    p1 = None
    for p2 in primes():
        if p1 is not None and p2 - p1 == 2 * k:
            yield (p1, p2)
        p1 = p2
    

if __name__ == "__main__":
    import sys
    import doctest
    doctest.testmod()
    sys.exit(1)