#!/usr/bin/env python
"""
 TwoDECA.py: Two-dimensional CA simulator
 	Using PyGame/SDL graphics

 Options:
 -t #  Number of time steps
 -n #  Size of lattice
 -c #  Size of cell
 -r #  Random initial configuration, rather than a single ON site
 -R #  Use rule # for the simulation
"""

	# Import modules
import getopt,sys
import time
from random import randint
try:
	import Numeric as N
	import pygame
	import pygame.surfarray as surfarray
except ImportError:
	raise ImportError, "Numeric and PyGame required."

	# Celullar automata routines
	# Initialize CA lattice, load in initial configuration
def CAInit(nSites,dorandom):
	state = N.zeros((nSites,nSites))
	if dorandom:   # Random IC
		for i in xrange(nSites):
			for j in xrange(nSites): state[i][j] = randint(0,1)
	else:          # Single ON site
		state[nSites/2][nSites/2] = 1
	return state

	# Convert rule number to map (look-up table)
	#   from neighborhoods to next values
	#	Binary CA, radius 1.
	#   von Neumann neighborhood
	#            a
	#           bcd
	#            e
	#	LUT index = abcde
def RuleNumberToLUT(rulenum):
	nStates = 2
	Radius  = 1
	nEntries = pow(nStates,4*Radius+1)
	lut = [(rulenum/pow(2,i)) % 2 for i in xrange(nEntries)]
	return lut

	# Sum mod 2 of local states in neighborhood
def SumMod2RuleLUT():
	nStates = 2
	Radius  = 1
	nNeighbors = 4 * Radius + 1
	nEntries = pow(nStates,nNeighbors)
	lut = [0]*nEntries
	for i in xrange(nEntries):
		sum = 0
		for j in xrange(nNeighbors):
			if i & (1<<j): sum += 1
		lut[i] = sum % 2
	return lut

	# Random rule
def RandomRuleLUT():
	nStates = 2
	Radius  = 1
	nNeighbors = 4 * Radius + 1
	nEntries = pow(nStates,nNeighbors)
	lut = [0]*nEntries
	for i in xrange(nEntries): lut[i] = randint(0,1)
	return lut

	# Iterate the CA lattice, producing a spacetime diagram:
	#   New value of site is a function of previous values of
	#      itself and two neighbors.
	#   Form integer of the previous values, use as index into look-up table.
	#   Impose spatially periodic boundary conditions.
def CAIterate(state,nSites,rule):
	new = [[
		rule[ 8*state[(i-1)%nSites][j] + 4*state[i][j]
			+ 2*state[(i+1)%nSites][j]
			+ 16*state[i][(j-1)%nSites] + state[i][(j+1)%nSites] ]
			for i in xrange(nSites)]
			for j in xrange(nSites)]
	return new

	# Spatial configuration display
def SpDisplayState(state,nSites,surface,CellSize):
	ConfigRect = pygame.Rect(0,0,nSites*CellSize,nSites*CellSize)
	surface.fill(OffColor,ConfigRect)
	for i in xrange(nSites):
		for j in xrange(nSites):
			if state[i][j] > 0:
				surface.fill(OnColor,[i*CellSize,j*CellSize,CellSize,CellSize])

# Set defaults
CellSize = 4
nSteps = 100
nSites   = 107
dorandom = 0
rule     = 18

# Get command line arguments, if any
opts,args = getopt.getopt(sys.argv[1:],'t:n:c:rR:')
for key,val in opts:
	if key == '-t': nSteps   = int(val)
	if key == '-n': nSites   = int(val)
	if key == '-c': CellSize = int(val)
	if key == '-r': dorandom = 1
	if key == '-R': rule     = int(val)

size = (CellSize*nSites,CellSize*nSites)
	# Set initial configuration
state = CAInit(nSites,dorandom)
	# Build rule's look-up table
LUT = SumMod2RuleLUT()
# LUT = RandomRuleLUT()

pygame.init()
OnColor = 255, 0, 0
OffColor = 0, 0, 255
	# Get display surface
screen = pygame.display.set_mode(size)
pygame.display.set_caption('2D Cellular Automaton Simulator')
	# Clear display
screen.fill(OffColor)
pygame.display.flip()
	# Create RGB array whose elements refer to screen pixels
	# Strangeness: sptmdiag is no longer used, but if this line
	#	is removed, then display updates slow down considerably!
sptmdiag = surfarray.pixels3d(screen)

t = 0
t0 = time.clock()
while 1:
	for event in pygame.event.get():
		if event.type == pygame.QUIT: sys.exit()
		# Iterate state
	state = CAIterate(state,nSites,LUT)
		# Display state
	SpDisplayState(state,nSites,screen,CellSize)
	pygame.display.flip()
	t += 1
	if (t % nSteps) == 0:
		t1 = time.clock()
		print "Iterations per second: ", float(nSteps) / (t1 - t0)
		t0 = t1