# ulam spiral by <silveiraneto@gmail.com>
# see http://en.wikipedia.org/wiki/Ulam_spiral
# or see http://www.eupodiatamatando.com/2007/05/17/espiral-de-ulam/

# This program generates a BMP file with the ulam spiral.
# Change the SIZE variable to set the size of the image.

import time # for benchmarks

# We try to open pygame
try:
	import pygame
except ImportError:
	print 'You need pygame installed to this program works.'
	print 'please visit http://www.pygame.org for instructions'

# Constants
SIZE = (5000, 5000)
WHITE = (255,255,255)
BLACK = (0,0,0)
BOXSIZE = 1

# Receives a number (small) and return True if prime, False otherwise.
# Beautiful but slow.
def isPrime(n):
	return not [x for x in xrange(2,int(n**0.5)+1) if n%x == 0] and n > 1

# Receives a number and return True if prime, False otherwise.
# A little better, aprox. 5 times faster.
def isPrime2(n):
	if n % 2 == 0 and n != 2:
		return False
	sqrt = int((n**0.5)+1)
	i = 3
	while(i<sqrt):
		if n % i == 0:
			return False
		i+=2
	return True

# Draw the ulam spiral on it
def draw_spiral(surface):
	step = 1
	x = SIZE[0]/2
	y = SIZE[1]/2
	movement = [(BOXSIZE,0), (0,-BOXSIZE), (-BOXSIZE, 0), (0,BOXSIZE)]
	k = 0
	n = 1
	i = 0
	while x > 0 and y > 0 and x < SIZE[0] and y < SIZE[1]:
		for j in xrange(step):			
  			if isPrime2(n):				
				if BOXSIZE == 1:
					pygame.draw.line(surface, BLACK, (x,y), (x,y))
				else:
 					pygame.draw.rect(surface, BLACK, (x,y,BOXSIZE,BOXSIZE) )

			orientation = movement[k]
			x += orientation[0]
			y += orientation[1]
			n += 1
		k = (k+1)%4
		if (i%2)==1:
			step += 1
		i += 1

# The main function
def main():
	img = pygame.Surface(SIZE, depth=8)
	img.fill(WHITE)
	start = time.clock()
 	draw_spiral(img)
	end = time.clock()
	output = str(SIZE[0])+'x'+str(SIZE[1])+'.bmp'
	pygame.image.save(img, output)
	print 'Output saved in '+output
	print 'Tested '+str(SIZE[0]*SIZE[1])+' numbers'
	print 'in '+str(end - start)+' seconds.'

if __name__ == "__main__":
    main()