#!/usr/bin/env python
"""
 OneDCA.py: One-dimensional elementary cellular automaton 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
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):
	if dorandom:   # Random IC
		first_row = [randint(0,1) for i in xrange(nSites)]
	else:          # Single ON site
		first_row = [0]*nSites
		first_row[nSites/2] = 1
	return first_row

	# Convert rule number to map (look-up table)
	#   from neighborhoods to next values
def RuleNumberToLUT(rulenum):
	lut = [(rulenum/pow(2,i)) % 2 for i in xrange(8)]
	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[4*state[(j-1)%nSites]+2*state[j]+state[(j+1)%nSites]]
		   for j in xrange(nSites)]
	return new

	# Space-time diagram display, wrap around
def SpTmDisplayState(state,nSites,surface,t,CellSize):
	RowRect = pygame.Rect(0,t*CellSize,nSites*CellSize,CellSize)
	surface.fill(white,RowRect)
	for i in xrange(nSites):
		if state[i] > 0:
			surface.fill(black,[i*CellSize,t*CellSize,CellSize,CellSize])

	# Space-time diagram display, wrap scrolling
def SpTmDisplayStateScroll(state,nSites,surface,nSteps,CellSize):
	# Scroll surface up CellSize pixels
	surface.unlock()
	surface.blit(surface,(0,0),[0,CellSize,nSites*CellSize,nSteps*CellSize])
	# Write new state at bottom row
	RowRect = pygame.Rect(0,(nSteps-1)*CellSize,nSites*CellSize,CellSize)
	surface.fill(white,RowRect)
	for i in xrange(nSites):
		if state[i] > 0:
			surface.fill(black,[i*CellSize,(nSteps-1)*CellSize,CellSize,CellSize])

# Set defaults
CellSize = 4
nSteps   = 125
nSites   = 207
size = (CellSize*nSites,CellSize*nSteps)
dorandom = 1
rule     = 18
	# Two display formats
DisplayFormat = 'Scroll'

# Get command line arguments, if any
opts,args = getopt.getopt(sys.argv[1:],'t:n:rwR:')
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 = 0
	if key == '-w': DisplayFormat = 'Wrap'
	if key == '-R': rule     = int(val)

	# Set initial configuration
state = CAInit(nSites,dorandom)
	# Translate rule number to look-up table
LUT = RuleNumberToLUT(rule)
	# Make spacetime diagram by iterating CA

pygame.init()
black = 0, 0, 0
white = 255, 255, 255
	# Get display surface
screen = pygame.display.set_mode(size)
pygame.display.set_caption('1D Elementary Cellular Automaton Simulator')
	# Clear display
screen.fill(white)
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
while 1:
	for event in pygame.event.get():
		if event.type == pygame.QUIT: sys.exit()
		# Iterate state
	state = CAIterate(state,nSites,LUT)
		# Display state in spacetime diagram
	if DisplayFormat == 'Scroll':	
		SpTmDisplayStateScroll(state,nSites,screen,nSteps,CellSize)
		pygame.display.flip()
	else:
		SpTmDisplayState(state,nSites,screen,t,CellSize)
		if (t % nSteps) == 0: pygame.display.flip()
	t += 1
	t %= nSteps