"""
primes.py  V.1.4, July 31 2006, by leonardo maffi.

Functions:
  primeQ(n): return True is n is prime.
  primes(n): return the list of prime numbers <= n.
  nextPrime(n): return the next prime number.
"""


def primes(n):
    """primes(n): return a list of prime numbers <= n.

    This code is fast but uses O(n) memory instead of O(sqrt(n)) like an usual Sieve.
    Code by Wensheng Wang (V.1.1, Sep 19 2005):
    http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/366178
    Then improved and debugged."""
    assert isinstance(n, (int, long))
    if n == 2:
        return [2]
    elif n<2:
        return []
    s = range(3, n+2, 2)
    mroot = n ** 0.5
    half = len(s)
    i = 0
    m = 3
    while m <= mroot:
        if s[i]:
            j = (m*m - 3) / 2
            s[j] = 0
            while j < half:
                s[j] = 0
                j += m
        i += 1
        m = 2 * i + 3
    if s[-1] > n:
        s[-1] = 0
    return [2] + filter(None, s)


def SPRP(N, B):
    if min(N, B) < 2:
        return False

    # Find d and s that satisfy N - 1 = 2^s * d, where d is odd and s >= 0.
    d = N - 1
    s = 0
    while not (d & 1):
        d >>= 1
        s += 1

    # Use right-to-left binary exponentiation to Calculate B^d mod N, result in i.
    i = 1
    r = B
    while d > 1:
        if d & 1:
            i = (i * r) % N
        d >>= 1
        r = (r * r) % N
    i = (i * r) % N

    # If i = 1 or N - 1, then N is a SPRP base B.
    result = (i == 1) or (i == N - 1)

    # Otherwise, see if i^(2^r) mod N = N - 1, where 0 <= r < s.
    r = 0
    while (not result) and (r + 1 <= s):
        i = (i * i) % N
        result = (i == N - 1)
        r += 1

    return result


def nextPrime(n):
    """
    nextPrime(n): return the prime number successive to n.
    It doesn't work with n >= 341.550.071.728.321

    Sources:
      http://primes.utm.edu/glossary/page.php?sort=StrongPRP
      if n < 1,373,653 is both a 2 and 3-SPRP, then n is prime.
      if n < 25,326,001 is a 2, 3, and 5-SPRP, then n is prime.
      if n < 118,670,087,467 is a 2, 3, 5, and 7-SPRP, then either n = 3,215,031,751 or n is prime.
      if n < 2,152,302,898,747 is a 2, 3, 5, 7, and 11-SPRP, then n is prime.
      if n < 3,474,749,660,383 is a 2, 3, 5, 7, 11, and 13-SPRP, then n is prime.
      if n < 341,550,071,728,321 is a 2, 3, 5, 7, 11, 13, and 17-SPRP, then n is prime.
    """
    assert isinstance(n, (int, long))
    if n < 2: return 2
    if n < 3: return 3

    while True:
        if n & 1:
            n += 2
        else:
            n += 1

        if n == 3215031751: continue

        if SPRP(n, 2) and SPRP(n, 3):
            if n < 1373653: return n
            if SPRP(n, 5):
                if n < 25326001: return n
                if SPRP(n, 7):
                    if n < 118670087467: return n
                    if SPRP(n, 11):
                        if n < 2152302898747: return n
                        if SPRP(n, 13):
                            if n < 3474749660383: return n
                            if SPRP(n, 17):
                                if n < 341550071728321: return n
                                else:
                                    raise "Too much big number for this function."


def primeQ(n):
    """primeQ(n): return True if n is prime.
    It doesn't work with n >= 341.550.071.728.321
    Sources: see nextPrime function."""
    assert isinstance(n, (int, long))
    if n < 2: return False
    if n == 2 or n == 3: return True

    if not(n & 1):
        return False

    if not SPRP(n, 2): return False
    if not SPRP(n, 3): return False
    if n < 1373653: return True
    if not SPRP(n, 5): return False
    if n < 25326001: return True
    if n == 3215031751: return False
    if not SPRP(n, 7): return False
    if n < 118670087467: return True
    if not SPRP(n, 11): return False
    if n < 2152302898747: return True
    if not SPRP(n, 13): return False
    if n < 3474749660383: return True
    if not SPRP(n, 17): return False
    if n < 341550071728321: return True
    else:
        raise "Too much big number for this function."


try:
    import psyco # Use Psyco if available.
except ImportError:
    pass
else:
    psyco.bind(SPRP)
    psyco.bind(primes)
    psyco.bind(nextPrime)
    psyco.bind(primeQ)

__all__ = ["primes", "nextPrime", "primeQ"]



# TESTS ================================================================
"""
# Code used to test nextPrime against Mathematica:

def gen_data():
    # this saves some numbers
    from random import randint, seed
    seed(1)
    thresholds = (1373653, 25326001, 3215031751, 118670087467, 2152302898747, 3474749660383, 341550071728321)

    testdata = list(thresholds[:-1])
    last = thresholds[-1]
    testdata.extend( randint(0, last) for x in xrange(10000) )

    for n in thresholds[:-1]:
        inf = int(n * 0.9)
        sup = int(n * 1.1)
        testdata.extend( randint(inf, sup) for x in xrange(5000) )

    inf = int(thresholds[-1] * 0.9)
    sup = thresholds[-1]
    testdata.extend( randint(inf, sup) for x in xrange(5000) )

    print "Numbers saved:", len(testdata)
    file("data", "w").write( "\n".join(map(str, testdata)) )


def test_data():
    # this takes the saved numbers and saves their successive primes into another file
    numbers = map(int, file("data"))
    print "Numbers:", len(numbers)
    succ = []
    for i, n in enumerate(numbers):
        if i % 500 == 0: print i
        succ.append( nextPrime(n))
    file("succ", "w").write( "\n".join(map(str, succ)) )


# Mathematica code used to test it:

succPrime[n_] := Block[{m = n + 1},
      While[Not[PrimeQ[m]], m++ ];
    m]

data = ReadList["c:\\datas\\data", Number];
succ = ReadList["c:\\datas\\succ", Number];

For[i = 1, i <= 45006, i++,
  If[Mod[i, 500] == 0, Print["*"]];
  If[
    succPrime[ data[[i]] ] != succ[[i]],
    Print[i]
    ]
  ]

"""



"""
# To test primeQ against nextPrime, very slow!

def _primeQ(n): # Just for test, do not use!
    return nextPrime(n-1) == n

from random import randint
thresholds = (1373653, 25326001, 3215031751, 118670087467, 2152302898747, 3474749660383, 341550071728321)

testdata = list(thresholds[:-1])
last = thresholds[-1]
testdata.extend( randint(0, last) for x in xrange(10000) )

for n in thresholds[:-1]:
    inf = int(n * 0.9)
    sup = int(n * 1.1)
    testdata.extend( randint(inf, sup) for x in xrange(5000) )

inf = int(thresholds[-1] * 0.9)
sup = thresholds[-1]
testdata.extend( randint(inf, sup) for x in xrange(5000) )

for n in testdata:
    assert _primeQ(n) == primeQ(n)
"""


"""
# To test primeQ against gmpy
import gmpy
from random import randint
for i in xrange(100):
    n = randint(2, 341550071728321)
    print n, gmpy.is_prime(n), primeQ(n)
    if gmpy.is_prime(n) != int(primeQ(n)):
        print n
"""


"""
def small_primeQ_test(n=200000):
    l = primes(n)
    sl = set(l)
    for i in xrange(1, l[-1]+1):
        if primeQ(i) != (i in sl):
            print "Error:", i
psyco.bind(small_primeQ_test)
small_primeQ_test()
"""
# TESTS END ================================================================


if __name__ == '__main__': # Demo
    print primes(200)
    print

    from time import clock
    pmax = 21
    for p in xrange(12, pmax):
        n = 2 ** p
        t1 = clock()
        primes(n)
        t2 = clock()
        print "primes(2^%s= %s):" % (p, n), round(t2-t1, 3), "s"