'''Module for working with prime numbers
'''
import math
import pickle

class PrimeGen():
	'''Main point of access for prime operations.'''
	def __init__(self):
		self.primelist = [1,2,3]
		self.at = 3
		self.highest = 3
		self.count = 3
		self.iternumber = None
		self.itercount = None
		self.iterstate = None


	def is_prime(self,n):
		'''calls is_prime_lazy'''
		return self.is_prime_lazy(n)

	def is_prime_lazy(self,number):
		'''Returns True if a number is prime.  Returns as soon as factor is found.  Tries to evaluate as few prime numbers as possible'''
		if number == 2 or number == 1:
			return True
		bound = math.sqrt(number)
		i = 1
		while self.primelist[i] <= bound:
			n = self.primelist[i]
			remainder = number % n
			if not remainder:
				return False
			if i >= (self.count - 1):
				self.get_next()
			i += 1
		return True

	def is_prime_greedy(self,n):
		'''Returns True if a number is prime.  Returns as soon as factor is found.  Makes sure it has a big enough list of primes before beginning'''
		i = 1
		if n == 2:
			return True
		bound = math.sqrt(n)
		if max(self.primelist) <= bound:
			self.go_till_num(bound)
		while self.primelist[i] <= bound:
			if n % self.primelist[i] == 0:
				return False
			i += 1
		return True

	def __call__(self,number=None,count=None):
		if number is not None:
			self.iternumber = number
		if count is not None:
			self.itercount = count
			self.iterstate = 0

		return self

	def next(self):
		'''Used in iterators. Example:
		p = prime.PrimeGen()
		for primenumber in p: print primenumber # Keeps generating numbers
		for primenumber in p(number=100): print primenumber # Generates remaining primes up to 100.  If already past 100, does nothing
		for primenumber in p(count=10): print primenumber # Generates next 10 prime numbers
		'''
		if self.itercount is not None:
			if self.iterstate < self.itercount:
				self.iterstate += 1
				return self.get_next()
			self.itercount = None
			raise StopIteration()
		elif self.iternumber is not None:
			if max(self.primelist) < self.iternumber:
				self.get_next()
				if self.at < self.iternumber:
					return self.at
			self.iternumber = None
			raise StopIteration()
		else:
			return self.get_next()

	def __iter__(self):
		return self

	def get_next(self):
		'''Finds next prime starting at self.at'''
		found = False
		while not found:
			self.at += 1
			if self.is_prime_greedy(self.at):
				self.primelist.append(self.at)
				self.count += 1
				found = True
		return self.at

	def go_till_count(self,number):
		'''Generate until the total number of primes (including already gen'd) is number'''
		while self.count < number:
			self.get_next()
			
	def go_till_num(self,number):
		'''Generate primes until self.at reaches number'''
		while max(self.primelist) < number:
			self.get_next()

	def __str__(self):
		return str(self.primelist[-10:])
		
	def find_prime_factors(self,number):
		'''Returns the prime factorization of number.  Will generate needed prime list if it doesn't exist'''
		factors = []
		bound = math.sqrt(number)
		#if max(self.primelist) <= bound:
			#self.go_till_num(bound)
		i = 1
		while self.primelist[i] <= bound:
			n = self.primelist[i]
			if n > bound: 
				return factors
			remainder = number % n
			if not remainder:
				factors.append(n)
				div = number/n
				subfactors = self.find_prime_factors(div)
				if subfactors:
					factors += subfactors
				else:
					factors.append(div)
				return factors
			if i >= (self.count -1):
				self.get_next()
			i += 1

	def save(self,filename='primelist.dct'):
		f = open(filename,'w')
		pickle.dump(self.primelist,f)
		f.close()
		
	def load(self,filename='primelist.dct'):
		f = open(filename,'r')
		self.primelist = pickle.load(f)
		f.close()
		self.at = max(self.primelist)


if __name__ == '__main__':
	import random
	import time
	TESTRUNS = 10
	TESTLENGTH = 1000
	MAX = 1000000
	times = []

	for t in xrange(TESTRUNS):
		t0 = time.time()

		p = PrimeGen()
		for i in xrange(TESTLENGTH):
			n = random.randrange(1,MAX)
			p.find_prime_factors(n)

		times.append((t0,time.time()))

	times = map(lambda x:x[1]-x[0],times)
	print 'Factorization Time:',sum(times)/len(times)


	p = PrimeGen()
	arry = []
	for primenumber in p:
		arry.append(primenumber)
		if primenumber> 30: break # Keeps generating numbers, make sure you exit manually
	print "Generate continuously, but break after 30 manually:", arry
	arry = []
	for primenumber in p(number=100): arry.append(primenumber) # Generates remaining primes up to 100.  If already past 100, does nothing
	print "Generate till 100:", arry
	arry = []
	for primenumber in p(count=10): arry.append(primenumber)# Generates next 10 prime numbers
	print "Generate next 10:", arry
	arry = []