#!/usr/bin/env python
"""
 TwoDLDS.py: Two-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 tandom initial configuration
 -r #  Map parameter value
 -c #  Coupling strength
 -L #  Simulate Logistic Lattice; Default: Dripping Handrail
"""

	# Import modules
import getopt,sys
import time
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,randomIC):
	state = N.zeros((nSites,nSites),N.Float)
	if randomIC:   # Random IC
		for i in xrange(nSites):
			for j in xrange(nSites): state[i][j] = random()
	else:          # Single "on" site
		for i in xrange(nSites):
			for j in xrange(nSites): state[i][j] = 0.5
		state[nSites/2][nSites/2] = 0.1
	return state

    # 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 (Moore) neighborhoods:
    #   Impose spatially periodic boundary conditions.
def NearestNeighborhoodCoupling(state,nSites,i,j,c):
	return ( c*state[(i-1)%nSites][j] + (1.0-4.0*c)*state[i][j]
			+ c*state[(i+1)%nSites][j]
			+ c*state[i][(j-1)%nSites] + c*state[i][(j+1)%nSites] )

	# Iterate the LDS lattice:
	#   New value of site is a function of previous values of
	#      itself and four (Moore) neighbors.
	#   Impose spatially periodic boundary conditions.
    # Dripping Handrail
def DHRIterate(state,nSites,r,c):
	couple = [ [ NearestNeighborhoodCoupling(state,nSites,i,j,c)
			for i in xrange(nSites) ]
			for j in xrange(nSites) ]
	map = [ [ LinearCircleMap(r,couple[i][j])
			for i in xrange(nSites) ]
			for j in xrange(nSites) ]
	return map

    # Logistic lattice
def LLIterate(state,nSites,r,c):
	couple = [ [ NearestNeighborhoodCoupling(state,nSites,i,j,c)
			for i in xrange(nSites) ]
			for j in xrange(nSites) ]
	map = [ [ LogisticMap(r,couple[i][j])
			for i in xrange(nSites) ]
			for j 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)
	return 255-int(255*x),150,int(255*x)

	# Spatial configuration display
def SpDisplayState(state,nSites,surface,CellSize):
	for i in xrange(nSites):
		for j in xrange(nSites):
			surface.fill(ColorMap(state[i][j]),
				[i*CellSize,j*CellSize,CellSize,CellSize])

# Set defaults
CellSize = 12
nSteps = 100
nSites = 47
randomIC = 1
    # Default system: Dripping Handrail
LDSIterate = DHRIterate
ColorMap = CircularColorMap
r = 0.065
c = 1./5.

# Get command line arguments, if any
opts,args = getopt.getopt(sys.argv[1:],'t:n:s:ILr: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': randomIC = 0
	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)

size = (CellSize*nSites,CellSize*nSites)
	# Set initial configuration
state = LDSInit(nSites,randomIC)

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

t = 0
t0 = time.clock()
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
	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