#!/usr/bin/env python
"""
 OneDLDS.py: One-dimensional lattice dynamical system simulator
 	Using PyGame/SDL graphics

 Options:
 -t #  Number of time steps
 -n #  Size of lattice
 -s #  Size of cell
 -I #  Single "on" site, rather than random initial configuration
 -w #  Wrap-around spacetime diagram; Default: scrolling sptm diagram
 -r #  Map parameter value
 -c #  Coupling strength
 -L #  Simulate Logistic Lattice; Default: Dripping Handrail
"""

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

	# 1D Map LDS routines
	# Initialize floating point lattice, load in initial configuration
def LDSInit(nSites,dorandom):
	if dorandom:   # Random IC
		first_row = [random() for i in xrange(nSites)]
	else:          # Single ON site
		first_row = [0.5]*nSites
		first_row[nSites/2] = 0.1
	return first_row

	# Logistic Map
def LogisticMap(r,x):
	return r * x * (1.0 - x)

	# Linear Circle Map
def LinearCircleMap(w,x):
	y = w + 0.95 * x
	while y > 1.0: y -= 1.0
	return y

	# Couple nearest neighborhoods:
	#   Impose spatially periodic boundary conditions.
def NearestNeighborhoodCoupling(state,nSites,i,c):
	return c*state[(i-1)%nSites] + (1.0-2.0*c)*state[i] + c*state[(i+1)%nSites]

	# Iterate the LDS lattice:
	#   New value of site is a function of preceding values of
	#      itself and two neighbors.
	# Dripping Handrail
def DHRIterate(state,nSites,r,c):
	couple = [ NearestNeighborhoodCoupling(state,nSites,i,c)
		   for i in xrange(nSites) ]
	map = [ LinearCircleMap(r,couple[i]) for i in xrange(nSites) ]
	return map

	# Logistic lattice
def LLIterate(state,nSites,r,c):
	couple = [ NearestNeighborhoodCoupling(state,nSites,i,c)
		   for i in xrange(nSites) ]
	map = [ LogisticMap(r,couple[i]) for i in xrange(nSites) ]
	return map

	# Colormap: Site value in [0,1] to RGB color
def SimpleColorMap(x):
	return int(255*x),int(255*x),int(255*(1.0-x))

	# Circular Colormap: Site value in [0,1] to RGB color periodically
def CircularColorMap(x):
	x = 1.0 + N.sin(2.0*N.pi*x)
	x *= 0.5
	return int(255*x),int(255*x),int(255*x)

	# Space-time diagram display, wrap around
def SpTmDisplayState(state,nSites,surface,t,CellSize):
	for i in xrange(nSites):
		surface.fill(ColorMap(state[i]),[i*CellSize,t*CellSize,CellSize,CellSize])

	# Space-time diagram display, 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
	for i in xrange(nSites):
		surface.fill(ColorMap(state[i]),[i*CellSize,(nSteps-1)*CellSize,CellSize,CellSize])

# Set defaults
CellSize = 4
nSteps   = 125
nSites   = 107
size = (CellSize*nSites,CellSize*nSteps)
dorandom = 1
	# Default system: Dripping Handrail
LDSIterate = DHRIterate
ColorMap = CircularColorMap
r = 0.1
c = 1./3.
	# Two display formats
DisplayFormat = 'Scroll'

# Get command line arguments, if any
opts,args = getopt.getopt(sys.argv[1:],'t:n:s:IwLr:c:')
for key,val in opts:
	if key == '-t': nSteps   = int(val)
	if key == '-n': nSites   = int(val)
	if key == '-s': CellSize = int(val)
	if key == '-I': dorandom = 0
	if key == '-w': DisplayFormat = 'Wrap'
	if key == '-L':
		LDSIterate = LLIterate
		ColorMap = SimpleColorMap
		r = 3.4
		c = 0.01
	if key == '-r': r        = float(val)
	if key == '-c': c        = float(val)

	# Set initial configuration
state = LDSInit(nSites,dorandom)

pygame.init()
black = 0, 0, 0
white = 255, 255, 255
	# Get display surface
screen = pygame.display.set_mode(size)
pygame.display.set_caption('1D Lattice Dynamical System Simulator')
	# Clear display
screen.fill(white)
pygame.display.flip()

t = 0
while 1:
	for event in pygame.event.get():
		if event.type == pygame.QUIT: sys.exit()
		# Iterate state
	state = LDSIterate(state,nSites,r,c)
		# 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