##
## Name:     prime.py
## Purpose:  Implements the Sieve of Eratosthenes.
## Author:   M. J. Fromberger <http://www.dartmouth.edu/~sting/>
##
## To work with primes, create an instance of the Primes class,
##   p = Primes()
##
## You can now index primes by position:
##   p[0] ==> 2, p[1] ==> 3, etc.
##
## You can slice:
##   p[:5]  ==> [2, 3, 5, 7, 11]
##   p[4:7] ==> [11, 13, 17]
##
## Test for primality as inclusion:
##   65537 in p   ==> True
##   45 in p      ==> False
##
## Find the index of a prime:
##   p.find(25)   ==> -1
##   p.find(5)    ==> 2
##
## Return the factorization of a value:
##   p.factor(25) ==> <generator> yielding [5, 5]
##
## Iterate over the collection:
##   for z in p:
##     if z < 1000: print z
##     else: break
##
## Clear out the cache
##   p.clip()
##
## Factoring and primality testing are done using trial division, but
## with with a few efficiencies: Primes are only generated up to the
## square root of the target value, and only odd candidates are tried,
## which reduces the loop overhead.
##
## Note, however, that the .find() must generate all the primes less
## than or equal to the candidate in order to find the index, so it is
## going to be time-consuming for large candidates.  Indexing and
## slicing have the same issue, although to a slightly lesser extent.
##

class Primes ( object ):
    """An indexable generator of primes."""
    
    def __init__(self):
        """Create a new instance of the generator."""

        # The cache is seeded with 2 and 3 so that the sieve can skip
        # even candidates easily.  Do not remove them, lest it fail.
        self._primes = [2, 3]
        self._pmap = {}
        self._pmap.update((q, p) for (p, q) in enumerate(self._primes))
    
    def next(self):
        """Generate the next prime not already known."""

        # This is mutually recursive with .factor(), but it is safe in
        # this case because all the candidate values are greater than
        # the largest prime already in the sieve at all times.
        
        c = self._primes[-1] + 2
        while self.factor(c).next() <> c:
            c += 2
        
        self._pmap[c] = len(self._primes)
        self._primes.append(c)
        return c
    
    def __contains__(self, prm):
        """Answers True if the given value is prime, otherwise False."""

        return self.factor(prm).next() == prm
    
    def __getitem__(self, pos):
        """Get the pos'th prime, indexed from 0.  This may be
        time-consuming for large indices."""
        
        if isinstance(pos, int):
            if pos < 0:
                raise ValueError("index must be non-negative")
            
            while len(self._primes) <= pos:
                self.next()
        elif isinstance(pos, slice):
            while len(self._primes) < max(pos.start, pos.stop):
                self.next()
        
        return self._primes.__getitem__(pos)
    
    def __iter__(self):
        """Obtain a generator over all the primes."""

        pos = 0
        while True:
            while pos >= len(self._primes):
                self.next()

            yield self._primes[pos]
            pos += 1
    
    def find(self, value):
        """Returns the prime index of the value, or -1 if the value is
        not prime.  This operation may be time consuming for large
        values."""
        
        while self._primes[-1] < value:
            self.next()
        
        return self._pmap.get(value, -1)
    
    def factor(self, value):
        """Return a generator for the prime factorization of value.
        Factors are returned in nondecreasing order.  Factors by
        trial division, but efficiently."""
        
        if value == 1 or value in self._pmap:
            yield value
            return
        
        for p in self:
            while value % p == 0:
                yield p
                value /= p
            else:
                if p ** 2 > value:
                    if value > 1:
                        yield value
                    return
    
    def clip(self, size = 0):
        """Trim the cache of primes to no more than size elements."""
        
        if size >= 0:
            self._primes = self._primes[:max(size, 2)]
            self._pmap = {}
            self._pmap.update((q, p) for (p, q) in enumerate(self._primes))
        else:
            raise ValueError("size must be non-negative")

# Here there be dragons