# ulamSpiral.py
# Make a Ulam Spiral highlighting all the prime numbers.
'''
https://en.wikipedia.org/wiki/Ulam_spiral
'''

import math
import matplotlib.pyplot as plt
import matplotlib.cm as cm
from prime import *


# N = 11  # number of rows/cols in arr[], and odd 
N = 201

def spiralArr(start, n):
  # Create an array of integers spiraling 
  # outward from the center beginning with start

  # arr[x][y]
  arr = [[0 for y in range(n)] for x in range(n)]

  counter = start
  dir = 'down'
  pos = (math.floor(n/2), math.floor(n/2))
   
  for i in range(n*n):
    northPos = pos[0] - 1
    southPos = pos[0] + 1
    eastPos = pos[1] + 1
    westPos = pos[1] - 1
  
    northval = 0 if (northPos < 0) else \
               arr[northPos][pos[1]]
  
    southval = 0 if (southPos >= n) else \
               arr[southPos][pos[1]]
  
    eastval = 0 if (eastPos >= n) else \
               arr[pos[0]][eastPos]
  
    westval = 0 if (westPos < 0) else \
              arr[pos[0]][westPos] 
  
    adjs = int(northval == 0) + int(southval == 0) + int(eastval == 0) + int(westval == 0)
    if adjs >= 3:
      if dir == 'down':
        dir = 'right'
      elif dir == 'right':
        dir = 'up'
      elif dir == 'up':
        dir = 'left'
      elif dir == 'left':
        dir = 'down'
  
    arr[pos[0]][pos[1]] = counter
    counter += 1
  
    if dir == 'left':
      pos = (pos[0], westPos)
    elif dir == 'down':
      pos = (southPos, pos[1])
    elif dir == 'right':
      pos = (pos[0], eastPos)
    elif dir == 'up':
      pos = (northPos, pos[1])

  return arr

arr = spiralArr(41, N)  # or 11, 17, 41 
                       # Euler's lucky numbers
# Create an array of boolean values: 
# T for prime, F for composite
primes = [[False for y in range(N)] for x in range(N)]
for i in range(N):
  for j in range(N):
    if isPrime(arr[i][j]):
      primes[i][j] = True
plt.matshow(primes, cmap=cm.binary)
plt.axis('off')
plt.show()