#!/usr/bin/env python
"""
 TwoDLDS.py: Two-dimensional lattice dynamical system simulator
 	Using PyGlet graphics
	2008 (C) JPC

 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
 -C #  Use circular colormap; Default: simple colorful colormap
"""

	# Import modules
	#
import getopt,sys
import time
from random import random
from numpy import pi, sin, zeros
try:
		# Disable error checking for increased performance
	from pyglet import options
	options['debug_gl'] = False
	from pyglet import window
	from pyglet.gl import *
except ImportError:
	raise ImportError, "PyGlet package required."

	# 1D Map LDS routines
	# Initialize floating point lattice, load in initial configuration
	#
def LDSInit(nSites,randomIC):
	state = zeros((nSites,nSites),dtype=float)
	if randomIC:   # Random IC
		for i in range(nSites):
			for j in range(nSites): state[i][j] = random()
	else:          # Single "on" site
		for i in range(nSites):
			for j in range(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 range(nSites) ]
			for j in range(nSites) ]
	map = [ [ LinearCircleMap(r,couple[i][j])
			for i in range(nSites) ]
			for j in range(nSites) ]
	return map

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

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

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

	# Spatial configuration display
	#
def SpDisplayState(state,nSites,CellSize):
	for i in range(nSites):
		for j in range(nSites):
			r,g,b = ColorMap(state[i][j])
			glColor3f(r,g,b)
			glRecti(i*CellSize,j*CellSize,(i+1)*CellSize,(j+1)*CellSize)
	screen.flip()

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

	# Get command line arguments, if any
	#
opts,args = getopt.getopt(sys.argv[1:],'t:n:s:ICLr: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 == '-C': ColorMap = CircularColorMap
	if key == '-L':
		LDSIterate = LLIterate
		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)

	# Get display surface
screen = window.Window(CellSize*nSites,CellSize*nSites,
	caption = '2D Lattice Dynamical System Simulator',resizable = False)
screen.set_location(100,50)
screen.clear()

t = 0
t0 = time.clock()
while not screen.has_exit:
	screen.dispatch_events()
		# Iterate state
	state = LDSIterate(state,nSites,r,c)
		# Display state
	SpDisplayState(state,nSites,CellSize)
	t += 1
	if (t % nSteps) == 0:
		t1 = time.clock()
		print "Iterations per second: ", float(nSteps) / (t1 - t0)
		t0 = t1