"""
 Routines for working with primes and their divisors.

   >>> p = Primes(1000)
   >>> len(p)
   168
   >>> p[-1]
   997
   >>> 101 in p
   True
   >>> p.divisors(336)
   [1, 2, 3, 4, 6, 7, 8, 12, 14, 16, 21, 24, 28, 42, 48, 56, 84, 112, 168, 336]
   >>> p[0:4]
   [2, 3, 5, 7]
   >>> p.number2powers(336)  
   [4, 1, 0, 1]
   >>> p.powers2number([4, 1, 0, 1])     # 2**4 + 3**1 + 5**0 + 7**1
   336
   >>> p.sum_divisors(336)               # sum of factors less than 336
   656

                                         # Euler's totient function :
   >>> p.phi(25)                         # phi(n) = count relatively prime < n
   20

   # much faster than repeated p.phi() :
   >>> phi_list(10)                      # [phi(0), phi(1), ... phi(n-1)]
   [0, 0, 1, 2, 2, 4, 2, 6, 4, 6]

 Defining a Primes object with, say,"p=Primes(1000)" doesn't actually
 calculate anything.  The first time that p is accessed in any way, all
 the primes up to 1000 are found and stored.  All subsequent methods
 are fast.

 The Primes objects are saved so that subsequent objects can re-use
 the previous calculations if possible.

 The string representation gives a summary of how many primes were
 found and how long it took to do the calcultion.  For example, on a
 2.4GHz MacBookPro, str(Primes(2e6)) gives

   '148933 primes up to 2000000 found in 00:00:18.22'.

 This file aso has an interface to a probabilistic prime test
 as defined in miller_rabin.py.

   >>> from primes import is_prime
   >>> is_prime(101)
   True
   >>> is_prime(100)
   False

Jim Mahoney | mahoney@marlboro.edu | Nov 2008 v0.1 | GPL
"""

from math import *
from elapsed_time import ElapsedTime

saved_primes = []

def is_prime(n, s=40):
  """Return True if n is probably prime with probability > (1-2**-s).
     >>> is_prime(101)
     True
     >>> is_prime(100)
     False
  """
  from miller_rabin import MillerRabin
  if (int(n)<2): return False
  return MillerRabin(int(n), s)

def test_phi_list():
  """
    a million euler totient values in 11 sec :
    > test_phi_list()
    phi_list(100) in 00:00:00.00 
    phi_list(1000) in 00:00:00.01 
    phi_list(10000) in 00:00:00.04 
    phi_list(100000) in 00:00:00.53 
    phi_list(1000000) in 00:00:10.95 
  """
  from elapsed_time import ElapsedTime
  for n in [1e2, 1e3, 1e4, 1e5, 1e6]:
    timer = ElapsedTime()
    p = phi_list(n)
    print " phi_list(%i) in %s " % (n, timer)

def phi_list(n):
  """Return a list of euler's totient function [phi(1), phi(2), ..., phi(n-1)
     using a modified Sieve of Eratosthenes.
     >>> phi_list(9) # is [phi(0), phi(1), ..., phi(8)]
     [0, 0, 1, 2, 2, 4, 2, 6, 4]
  """
  n=int(n)
  numbers = range(2,n+1)
  stop_value = sqrt(n)  # above this the sieve has found all primes
  phis = range(0,n)
  phis[1]=0
  while numbers:
    np = p = numbers.pop(0)
    # print "  p = %i " % p
    while np < n:                       # for each multiple np of this prime,
      # print "   n*p = %i " % np
      # print "   phi[%i] was %i " % (np,phis[np])
      phis[np] = phis[np]*(p-1)/p       # reduce phi(np) by (1-1/p).
      # print "   phi[%i] is now %i " % (np,phis[np])
      np += p
    if p < stop_value:  # past this, numbers[] list is all primes.
      numbers = filter(lambda m: m % p != 0, numbers)
  return phis

def sieve(n):
  """Return a list of primes up to n via the Sieve of Eratosthenes.
     (The algorithm here was based on its wikipedia article.)
     >>> sieve(11)
     [2, 3, 5, 7, 11]
  """
  numbers = range(2,n+1)
  stop_value = sqrt(n)
  primes = []
  while True:
    p = numbers.pop(0)
    primes.append(p)
    # print " %8i : %i " % (len(primes), p)
    if p > stop_value:
      return primes + numbers
    ### cross out multiples of p from the list of numbers.
    ## first version
    # copy = []
    # for n in numbers:
    #   if n % p != 0:
    #     copy.append(n)
    # numbers = copy
    ## second version; faster
    numbers = filter(lambda m: m % p != 0, numbers)

def mapcount(baselist, func):
  """Given a baselist like [3,2,1],
     return [func([0,0,0], func([0,0,1]), func([0,1,0],...]
     where the successive lists count up to each of the values in baselist.
     (The function will passed lists which are the same size as baselist;
      whatever it returns will be put into the created list.)
     >>> mapcount([2,1], lambda x: x[:])
     [[0, 0], [1, 0], [2, 0], [0, 1], [1, 1], [2, 1]]
  """
  if len(baselist)==0:
    return [func([])]
  current = [0]*len(baselist)
  howmany = reduce(lambda y,x:y*(x+1), baselist, 1)
  answer = []
  for i in range(howmany):
    answer.append(func(current))
    i = 0
    current[i] += 1
    while current[i] > baselist[i] and i+1 < len(current):
      current[i] = 0
      i += 1
      current[i] += 1
  return answer


class Primes:

  def __init__(self, maximum):
    self.maximum = int(maximum)
    self.is_setup = False
    saved_primes.append(self)

  def __str__(self):
    self.setup()
    return "%i primes up to %i found in %s" \
           % (len(self.primes), self.maximum, self.setup_time)

  def phi(self, n):
    """Euler's 'totient' function
       >>> p = Primes(100)
       >>> p.phi(17)        # phi(prime) = prime-1
       16
       >>> p.phi(25)        # phi(prime**k) = (prime-1)*prime**(k-1)
       20
       >>> p.phi(77)        # phi(a*b) = (a-1)*(b-1) when a,b prime
       60
    """
    if n==1: return 0
    # print "phi(%i) = ... " % n
    if n in self: return n - 1
    index = 0
    result = m = n
    while True:
      p = self[index]
      # print "   At prime %i: result=%i, m=%i" % (p, result, m)
      if p > m:
        return result
      elif m % p == 0:
        m = m/p
        result = (result*(p - 1))/p
      index += 1

  def setup(self):
    if self.is_setup:
      return True
    self.factorization_hash = {}
    for p in saved_primes:
      if not id(p)==id(self) and p.maximum >= self.maximum and p.is_setup:
        self.maximum = p.maximum
        self.primes = p.primes
        self.prime_hash = p.prime_hash
        self.setup_time_sec = 0
        self.setup_time = '0 sec (previously calculated)'
        self.is_setup = True
        return True
    timer = ElapsedTime()
    self.primes = sieve(self.maximum)
    self.prime_hash = {}
    for prime in self.primes:
      self.prime_hash[str(prime)] = True
    self.setup_time_sec = timer
    self.setup_time = str(timer)
    self.is_setup = True
    return True

  def __len__(self):
    self.setup()
    return len(self.primes)

  def __contains__(self, n):
    self.setup()
    if n > self.maximum:
      raise Exception, "(%i in primes) out of bounds; max=%i" \
            % (n, self.maximum)
    return self.prime_hash.get(str(n), False)

  def __getitem__(self, k):
    self.setup()
    return self.primes[k]

  def clear_history(self):
    """Delete history of previously created Primes objects."""
    del saved_primes[0:len(saved_primes)]

  def number2powers(self, n):
    """Return prime factor powers b[:] where n = p[0]**b[0] + ...
       and where p are the prime numbers [2,3,5,...].
       >>> p = Primes(100)
       >>> p.number2powers(63)
       [0, 2, 0, 1]
    """
    # uses primes[] as defined above
    self.setup()
    answer = []
    i = 0
    while n > 1:
      answer.append(0)
      while n % self.primes[i] == 0:
        n = n/self.primes[i]
        answer[i] += 1
      i += 1
    return answer

  def powers2number(self, powers):
    """Return n given b[:] where n = p[0]**b[0] + ...
       >>> p = Primes(100)
       >>> p.powers2number([0, 2, 0, 1])
       63
    """
    self.setup()    
    answer = 1
    for i in range(len(powers)):
      answer *= self.primes[i]**powers[i]
    return answer

  def n_divisors(self, n):
    """Return the number of divisors of n.
       >>> p = Primes(100)
       >>> p.n_divisors(24)    # 1, 2, 3, 4, 6, 8, 12, 24 : 8 divisors.
       8
    """
    # The method here is quicker than actually finding the divisors themselves.
    # Instead, find powers of primes such that n = primes[i]**b[i],
    # and from that find product((1+b[i])) which is the way that
    # the prime components can be rearranged and hence the number of factors.
    self.setup()
    return reduce(lambda y,x:y*(x+1), self.number2powers(n), 1)

  def n_prime_factors(self, n):
    """Return the number of prime factors of n.
       >>> p = Primes(100)
       >>> p.n_prime_factors(20) # number2powers is [2,0,1] ; 2 prime factors
       2
    """
    self.setup()
    return len(filter(lambda x:x>0, self.number2powers(n)))

  def factorizations(self, n):
    """Return all factorizations of n. (other than 1*n)
       >>> p = Primes(100)
       >>> p.factorizations(24)
       [(2, 2, 2, 3), (2, 2, 6), (2, 3, 4), (2, 12), (3, 8), (4, 6)]
    """
    self.setup()
    if self.factorization_hash.has_key(n):
      return self.factorization_hash[n]
    #print "factorizations({})".format(n)
    if n in self:        # if prime,
      answer = []        # no factorizations.
      return []
    else:
      d = self.divisors(n)
      #print " d=",d
      result = set()
      for i in d[1:]:
        j = n/i
        if j < i:
          break
        result.add((i,j))
      for (x,y) in set(result):
        #print " (x,y)=", (x,y)
        for z in self.factorizations(y):
          result.add(tuple(sorted((x,) + z)))
        for z in self.factorizations(x):
          result.add(tuple(sorted((y,) + z)))
      answer = list(result)
      answer.sort()
    self.factorization_hash[n] = answer
    return answer

  def divisors(self, n, do_sort=True):
    """Return a list of the divisors of n, e.g. [1, ..., n].
       >>> p = Primes(100)
       >>> p.divisors(24)
       [1, 2, 3, 4, 6, 8, 12, 24]
    """
    # The sort slows this down a bit, but gives a more intuitive result.
    #
    # The approach here is to loop by "counting over the prime factor powers,
    # which doesn't give them in numerical order.
    #
    # For example, for n=24, the prime decomposition is 2**3 * 3**1
    # which means powers = [3,1] = exponents primes[:] list.
    # Treating this as a 2-digit number and counting gives
    #                [0,0], [1,0], [2,0], [3,0], [0,1], [1,1], [2,1], [3,1]
    # which becomes  1      2      4      8      3      6      12     24
    #
    self.setup()    
    result = mapcount(self.number2powers(n), lambda x: self.powers2number(x))
    if do_sort:
      result.sort()
    return result

  def sum_divisors(self, n):
    """Return sum of divisors of n less than n.
       >>> p = Primes(10)
       >>> p.sum_divisors(12)         # 1+2+3+4+6
       16
    """
    self.setup()
    return sum(self.divisors(n)[0:-1])

def _test():
  from doctest import testmod
  testmod()
if __name__=="__main__": _test()