#!/usr/bin/env python
"""
 OneDLDS.py: One-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 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
 -C #  Use circular colormap; Default: simple colorful colormap
"""

	# Import modules
	#
import getopt,sys
from random import random
from time import clock
from numpy import pi, sin
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 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 range(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 range(nSites) ]
	map = [ LinearCircleMap(r,couple[i]) for i in range(nSites) ]
	return map

	# Logistic lattice
	#
def LLIterate(state,nSites,r,c):
	couple = [ NearestNeighborhoodCoupling(state,nSites,i,c)
		   for i in range(nSites) ]
	map = [ LogisticMap(r,couple[i]) for i 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

	# Space-time diagram display, wrap around
	#
def SpTmDisplayState(state,nSites,t,CellSize):
	for i in range(nSites):
		r,g,b = ColorMap(state[i])
		glColor3f(r,g,b)
		glRecti(i*CellSize,(nSteps-t-1)*CellSize,
			(i+1)*CellSize,(nSteps-t)*CellSize)
	if (t % nSteps) == 0: screen.flip()

	# Space-time diagram display, scrolling
	#
def SpTmDisplayStateScroll(state,nSites,nSteps,CellSize):
	# Copy front buffer to back buffer:
		# Set buffer to read from. Destination is GL_BACK; this never changes.
	glReadBuffer(GL_FRONT)
		# Offset copy by CellSize pixels to implement vertical scroll
	glRasterPos2i(0,CellSize)
	glCopyPixels(0,0,nSites*CellSize,(nSteps-1)*CellSize,GL_COLOR)
		# If glCopyPixels and related commands elsewhere, use:
#	glReadBuffer(GL_BACK)
	# Write new state at bottom row
	for i in range(nSites):
		r,g,b = ColorMap(state[i])
		glColor3f(r,g,b)
		glRecti(i*CellSize,0,(i+1)*CellSize,CellSize)
	screen.flip()

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

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

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

	# Make spacetime diagram by iterating LDS
	#
screen = window.Window(CellSize*nSites,CellSize*nSteps,
	caption = '1D Lattice Dynamical System Simulator',resizable = False)
screen.set_location(100,50)
screen.clear()

t = 0
FirstTime = 1
LastTime = 0.0
while not screen.has_exit:
	screen.dispatch_events()
		# Iterate state
	state = LDSIterate(state,nSites,r,c)
		# Display state in spacetime diagram
	if DisplayFormat == 'Scroll':
		if (FirstTime == 1):
			SpTmDisplayState(state,nSites,t,CellSize)
			screen.flip()
			if (t == nSteps - 1):
				FirstTime = 0
		else:
			SpTmDisplayStateScroll(state,nSites,nSteps,CellSize)
	else:
		SpTmDisplayState(state,nSites,t,CellSize)
	t += 1
	t %= nSteps
	if t == 0:
		print "%d steps in %f secs; %f sps" \
			% (nSteps, clock()-LastTime, nSteps/(clock()-LastTime))
		LastTime = clock()