#!/usr/bin/python -tO

from OpenGL.GL   import *
from OpenGL.GLU  import *
from OpenGL.GLUT import *
from math import *

import sys, types, random, copy

__date__    = "20041011"
__author__  = "Marcus Fritzsch <m _AT fritschy DOT_de>"
__license__ = "GPL"

class Mouse:
	"""my abstract mouse"""
	def __init__ (self):
		self.dx = self.dy = 0.0
		self.left = self.middle = self.right = False

class MouseMotion:
	"""my abstract mouse movement"""
	def __init__ (self):
		self.mouse = Mouse ()
		self.xRotate = self.yRotate = 0.0
		self.xTrans, self.yTrans, self.zTrans = 0.0, 0.0, -50.0

class Object3D:
	"""Base for 3d objects that can be rotated and/or translated"""
	def __init__ (self):
		self.xRot = self.yRot = self.zRot = 0.0
		self.xTrans = self.yTrans = self.zTrans = 0.0

	def doTransform (self):
		"""do simple transforms"""
		glTranslated (self.xTrans, self.yTrans, self.zTrans)
		glRotated (self.xRot, 1.0, 0.0, 0.0)
		glRotated (self.yRot, 0.0, 1.0, 0.0)
		glRotated (self.zRot, 0.0, 0.0, 1.0)

class BaseDrawMode:
	"""base of all drawing modes... like GL_POINTS or GL_LINES"""
	def __init__ (self, mode = 0):
		# XXX: need mode for fast distinguishement of DrawModes
		#      somewhat crappy... but works
		self.mode = mode
		pass
	def begin (self):
		"""for drawing modes that require glBegin/End calls
		should be overidden"""
		pass
	def end (self):
		"""for drawing modes that require glBegin/End calls
		should be overidden"""
		pass
	def draw (self, *a):
		"""draw the primitive
		should be overidden"""
		pass

class DrawPoints (BaseDrawMode):
	"""draw GL_POINTS primitive"""
	def __init__ (self):
		BaseDrawMode.__init__ (self, 1)
	def begin (self):
		glBegin (GL_POINTS)
	def end (self):
		glEnd ()
	def draw (self, *a):
		glVertex3fv (a [0])

class DrawLines (BaseDrawMode):
	"""draw GL_LINES primitive"""
	def __init__ (self):
		BaseDrawMode.__init__ (self, 2)
	def begin (self):
		glBegin (GL_LINES)
	def end (self):
		glEnd ()
	def draw (self, *a):
		glVertex3fv (a [1])
		glVertex3fv (a [0])

class DrawSpheres (BaseDrawMode):
	"""draw spheres using glutSolidSphere"""
	def __init__ (self):
		BaseDrawMode.__init__ (self, 3)
		# some default... silly settings
		self.radius, self.slices, self.stacks = 0.3, 8, 6
	def draw (self, *a):
		glTranslatef (a [0][0] - a [1][0], a [0][1] - a [1][1], a [0][2] - a [1][2])
		glutSolidSphere (self.radius, self.slices, self.stacks)

class Modes:
	"""valid Modes for selection purpose"""
	lines   = DrawLines
	points  = DrawPoints
	spheres = DrawSpheres

class Color:
	def __init__ (self, *a, **b):
		if len (a) == 3:
			self.r, self.g, self.b = a
		elif len (b.keys ()):
			self.r, self.g, self.b = b ['r'], b ['g'], b ['b']

	def listRed (self):
		for a in [Color (r=x, g=0, b=0) for x in xrange (0, 0x100)]:
			yield a

	def listRedMinus (self):
		for a in [Color (r=x, g=0, b=0) for x in xrange (0xff, -1, -1)]:
			yield a

	def listGreen (self):
		for a in [Color (r=0, g=x, b=0) for x in xrange (0, 0x100)]:
			yield a

	def listGreenMinus (self):
		for a in [Color (r=0, g=x, b=0) for x in xrange (0xff, -1, -1)]:
			yield a

	def listBlue (self):
		for a in [Color (r=0, g=0, b=x) for x in xrange (0, 0x100)]:
			yield a

	def listBlueMinus (self):
		for a in [Color (r=0, g=0, b=x) for x in xrange (0xff, -1, -1)]:
			yield a

	def iadd (self, col):
		self.r += col.r
		self.g += col.g
		self.b += col.b
		return self

	def listRed2Green2Blue (self):
		a = Color ()
		for r in a.listRed ():
			yield r
		b = a.listGreen ()
		for r in a.listRedMinus ():
			r.iadd (b.next ())
			yield r
		b = a.listBlue ()
		for r in a.listGreenMinus ():
			r.iadd (b.next ())
			yield r

#class BaseColorMode:
#	def __init__ (self):
#		pass
#
#class ColorRed2Green2Blue (BaseColorMode):
#	pass

class StrangeAttractor (Object3D):
	"""Base for strange attractors"""
	def __init__ (self, name, params):
		Object3D.__init__ (self)

		# set attractor attributes... some general
		self.a = self.b = self.c = self.h = 0.0
		self.steps = 0

		self.colors = Color ()

		self.mode = DrawPoints ()
		self.name = name										# the attractor's name
		self.list = 0
		self.setParameters (**params)
		self.setupDisplayList ()						# try setting up the disp-list

	def setMode (self, mode):
		"""update self.mode to mode"""
		m = mode ()
		if self.mode.mode != m.mode:
			self.mode = mode ()
			self.anythingChanged = True

	def setParameters (self, *a, **args):
		"""set attractor parameters
		ex:
		setParameters (a = 10.0, b = 20.0, c = 23.0, h = 10, steps = 10000)"""
		params = args
		if len (a) > 0 and len (a [0]) == 2:
			params = dict (a)
		elif len (a) > 0 and len (a) % 2 == 0:
			for i in xrange (len (a)):
				if i % 2 == 0:
					params [a [i]] = a [i+1]
		for k, v in params.items ():				# parse parameters
			try:
				p = getattr (self, k)
				print self.name, " setting parameter  ", k, '=', p, "to", v
			except:
				print self.name, " adding parameter  ", k, '=', v
			self.__dict__ [k] = v
		self.anythingChanged = True					# for list-check

	def draw (self):
		"""draw attractor"""
		self.doTransform ()
		if self.doWithoutList:
			self.display ()
		else:
			if self.anythingChanged:
				self.setupDisplayList ()
			else:
				glCallList (self.list)

	def setupDisplayList (self):
		"""if s.t. changed, recompile display list"""
		if glIsList (self.list):
			glDeleteLists (self.list, 1)

		self.genList ()

		if self.doWithoutList:
			return

		glNewList (self.list, GL_COMPILE)
		glColor3d (1.0, 1.0, 1.0)
		self.mode.begin ()
		self.display ()
		self.mode.end ()
		glEndList ()

		self.anythingChanged = False

	def displayAtomic (self, old, new):
		"""draw a point, line or s.t. alse..."""
		self.mode.draw (new, old)

	def genList (self):
		"""generate display list, if possible"""
		self.list = glGenLists (1)
		if self.list == 0:
			self.doWithoutList = True
			print >> sys.stderr, "[!]  Display lists are strangely not available!!"
		else:
			self.doWithoutList = False

	def display (self):
		pass

class Lorenz (StrangeAttractor):
	"""Specialized attractor, Lorenz"""
	def __init__ (self, *a, **b):
		del a		# unused
		StrangeAttractor.__init__ (self, 'Lorenz', b)
		self.xRot = -90.0
		self.yTrans = -25.0

	def display (self):
		"""general display method, will be caled by the superclass"""
		dx = dy = dz = x1 = y1 = z1 = 0.0
		x0 = y0 = z0 = 1e-31
		for S in xrange (self.steps + 1):
			dx, dy, dz = self.h * self.a * (y0 - x0), self.h * (x0 *(self.b - z0) - y0), self.h * (x0 * y0 - self.c * z0)
			x0 += dx
			y0 += dy
			z0 += dz
			if S > 100:
				self.displayAtomic ((x1, y1, z1), (x0, y0, z0))
			x1, y1, z1 = x0, y0, z0

class Roessler (StrangeAttractor):
	"""Specialized attractor, Roessler"""
	def __init__ (self, *a, **b):
		del a		# unused
		StrangeAttractor.__init__ (self, 'Roessler', b)
		# some custom transformations...
		self.xRot = -90.0
		self.yTrans -= 17.0

	def display (self):
		"""general display method, will be caled by the superclass"""
		x1 = y1 = z1 = x0 = y0 = z0 = 1e-4
		dx = dy = dz = 0.0
		for S in xrange (self.steps + 1):
			dx, dy, dz = self.h * (-y0 - z0), self.h * (x0 + self.a * y0), self.h * (self.b + z0 * (x0 - self.c))
			x0 += dx
			y0 += dy
			z0 += dz
			if S > 100:
				self.displayAtomic ((x1, y1, z1), (x0, y0, z0))
			x1, y1, z1 = x0, y0, z0

class PeterDeJong (StrangeAttractor):
	def __init__ (self, *a, **b):
		del a		# unused
		StrangeAttractor.__init__ (self, 'PeterDeJong', b)
	def display (self):
		# xn+1 = sin(a yn) - cos(b xn)
		# yn+1 = sin(c xn) - cos(d yn)
		x0 = y0 = x1 = y1 = dx = dy = 0.0
		for S in xrange (self.steps + 1):
			dx, dy = (sin (self.a * y0) - cos (self.b * x0)), (sin (self.c * x0) - cos (self.d * y0))
			x0, y0 = dx, dy
			self.displayAtomic ((x1, y1, 0), (x0, y0, 0))
			x1, y1 = x0, y0

class PeterDeJong (StrangeAttractor):
	def __init__ (self, *a, **b):
		del a		# unused
		StrangeAttractor.__init__ (self, 'PeterDeJong', b)
	def display (self):
		# xn+1 = sin(a yn) + c cos(a xn)
		# yn+1 = sin(b xn) + d cos(b yn)
		x0 = y0 = z0 = x1 = y1 = z1 = dx = dy = dz = 0.0
		for S in xrange (self.steps + 1):
			dx, dy, dz = sin (self.a * y0) - z0 * cos (self.b * x0), z0 * sin (self.c * x0) - cos (self.d * y0), self.e * sin (x0)
			x0, y0, z0 = dx, dy, dz
			self.displayAtomic ((x1, y1, z1), (x0, y0, z0))
			x1, y1, z1 = x0, y0, z0
#	def display (self):
#		# xn+1 = sin(a yn) + c cos(a xn)
#		# yn+1 = sin(b xn) + d cos(b yn)
#		x0 = y0 = x1 = y1 = dx = dy = 0.0
#		for S in xrange (self.steps + 1):
#			dx, dy = (sin (self.a * y0) - cos (self.b * x0)), (sin (self.c * x0) - cos (self.d * y0))
#			x0, y0 = dx, dy
#			self.displayAtomic ((x1, y1, 0), (x0, y0, 0))
#			x1, y1 = x0, y0

class StrangeApp:
	"""Application main class... does everything from OpenGL to Arguments..."""

	doc = """Keys:
	'+'   increase number of points (1000)
	'-'   decrease number of points (1000)
	'?'   show this help message
	'r'   select Roessler attractor
	'l'   select Lorenz attractor
	'L'   select displaymode lines
	'P'   select displaymode points (default)
	'S'   select displaymode spheres (really slow)
	'F'   toggle fog on/off (on is default)
	'a'   decrease attractor parameter a
	'A'   increase attractor parameter a
	'b'   decrease attractor parameter b
	'B'   increase attractor parameter B
	'c'   decrease attractor parameter c
	'C'   increase attractor parameter c
	'h'   decrease attractor parameter h
	'H'   increase attractor parameter h

Mouse:  
	right button pressed:  pan the attractor
	left button pressed:   rotate the attractor
	middle button pressed: zoom
	mousewhele up/down:    zoom

Special keys:
	up,down,right,left: pan the attractor
	pg up/down:         zoom
	"""

	def __init__ (self):
		self.motion = MouseMotion ()
		self.screenWidth = self.screenHeight = 400

		self.initGL ()		# before creating gl-objects, create gl status
		#self.attractor = Lorenz (a = 10.0, b = 28.0, c = 8.0 / 3.0, h = 0.01, steps = 10000)
		#self.attractor = PeterDeJong (a = 1.4, b = -2.3, c = 2.4, d = -2.1, e = 2.0, steps = 10000)
		#-0.546116
		#-0.219956
		#0.727927
		#-1.58878
		#-1.31701
		self.attractor = PeterDeJong (a = -0.546116, b = -0.219956, c = 0.727927, d = -1.58878, e = -1.31701, steps = 50000)
		self.hasFog = True

		self.rand = random.Random ()

	def run (self):
		"""start glutMainLoop"""
		glutMainLoop ()

	def initGL (self):
		"""init OpenGL state machine"""
		glutInitWindowSize (self.screenWidth, self.screenHeight)
		glutInitDisplayMode (GLUT_RGBA | GLUT_DOUBLE | GLUT_DEPTH)
		glEnable (GL_DEPTH_TEST)
		glDepthFunc (GL_GEQUAL)
		glutInit (sys.argv)
		glutCreateWindow ("strange Attractors")

		glutDisplayFunc (self.display)
		glutSpecialFunc (self.keyboardSpecial)
		glutKeyboardFunc (self.keyboard)
		glutReshapeFunc (self.reshape)
		glutMotionFunc (self.motionFunc)
		glutMouseFunc (self.mouse)

		self.setupFog ()
		#self.setupLighting ()
		glShadeModel (GL_SMOOTH)

		self.reshape (self.screenWidth, self.screenHeight)

	def setupLighting (self):
		glMatrixMode (GL_MODELVIEW)
		glEnable (GL_LIGHTING)
		glEnable (GL_LIGHT1)

		glLightfv (GL_LIGHT1, GL_POSITION, (5, 10, -10))
		glLightfv (GL_LIGHT1, GL_AMBIENT, (0.1, 0.1, 0.1))
		glLightfv (GL_LIGHT1, GL_DIFFUSE, (0.5, 0.5, 0.5))
		glLightfv (GL_LIGHT1, GL_SPECULAR, (0.9, 1.0, 0.9))

		glMaterialfv (GL_FRONT_AND_BACK, GL_SHININESS, (0.1,))

	def setupFog (self):
		"""setup OpenGL fog... well, somewhat crappy, cause I do not
		fully understand OpenGL fog"""
		fogColor = (0.0, 0.0, 0.2, 1.0)

		glEnable (GL_FOG)
		glFogi (GL_FOG_MODE, GL_EXP2)
		glFogfv (GL_FOG_COLOR, fogColor)
		glFogf (GL_FOG_DENSITY, 0.25)
		glHint (GL_FOG_HINT, GL_DONT_CARE)
		glFogf (GL_FOG_START, 0.1)
		glFogf (GL_FOG_END, 2.0)
	
		glClearColor (fogColor [0], fogColor [1], fogColor [2], fogColor [3])

	def mouse (self, button, state, x, y):
		"""general mouse callback, sets self.motion and self.mouse stati"""
		x, y = x * 0.125, y * 0.125
		if button == GLUT_LEFT_BUTTON:
			if state == GLUT_DOWN:
				self.motion.mouse.left = True
				self.motion.mouse.dx = x - self.motion.yRotate
				self.motion.mouse.dy = y - self.motion.xRotate
			else:
				self.motion.mouse.left = False
		elif button == GLUT_MIDDLE_BUTTON:
			if state == GLUT_DOWN:
				self.motion.mouse.middle = True
				self.motion.mouse.dy = y - self.motion.zTrans
			else:
				self.motion.mouse.middle = False
		elif button == GLUT_RIGHT_BUTTON:
			if state == GLUT_DOWN:
				self.motion.mouse.right = True
				self.motion.mouse.dx = x - self.motion.xTrans
				self.motion.mouse.dy = y - self.motion.yTrans
			else:
				self.motion.mouse.right = False
		# XXX: buttons 3 and 4 are on my mouse whela down, up
		elif button in (3, 4) and state == GLUT_DOWN:
			if button == 3: self.motion.zTrans -= 2.5
			else: self.motion.zTrans += 2.5
			glutPostRedisplay ()

	def motionFunc (self, x, y):
		"""move scene by translating/rotating by mouse movement"""
		x, y = x * 0.125, y * 0.125
		if self.motion.mouse.left:
			self.motion.yRotate = x - self.motion.mouse.dx
			self.motion.xRotate = y - self.motion.mouse.dy
			if self.motion.yRotate > 359.0: self.motion.yRotate -= 360.0
			if self.motion.xRotate > 359.0: self.motion.xRotate -= 360.0
		elif self.motion.mouse.middle:
			self.motion.zTrans = y - self.motion.mouse.dy
		elif self.motion.mouse.right:
			self.motion.xTrans = x - self.motion.mouse.dx
			self.motion.yTrans = y - self.motion.mouse.dy
		glutPostRedisplay ()

	def idle (self):
		"""glutIdle callback"""
		glutPostRedisplay ()

	def reshape (self, w, h):
		"""glut reshape callback"""
		if h == 0: h = 1
		glViewport (0, 0, w, h)
		glMatrixMode (GL_PROJECTION)
		glLoadIdentity ()
		gluPerspective (45.0, float (w) / float (h), 0.1, 10.0)
		glMatrixMode (GL_MODELVIEW)
		self.screenWidth, self.screenHeight = w, h

	def display (self):
		"""glut display callback"""
		glClear (GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
		glLoadIdentity ()

		glScaled (0.05, 0.05, 0.05)
		glTranslated (self.motion.xTrans, -self.motion.yTrans, self.motion.zTrans)

		glRotatef (self.motion.xRotate, 1.0, 0.0, 0.0)
		glRotatef (self.motion.yRotate, 0.0, 1.0, 0.0)

		self.attractor.draw ()

		# show axes
		glBegin (GL_LINES);

		# x-axis
		glColor3f (1.0, 0.0, 0.0)
		glVertex3f (0.0, 0.0, 0.0)
		glVertex3f (1.0, 0.0, 0.0)

		# y-axis
		glColor3f (0.0, 1.0, 0.0)
		glVertex3f (0.0, 0.0, 0.0)
		glVertex3f (0.0, 1.0, 0.0)

		# z-axis
		glColor3f (0.0, 0.0, 1.0)
		glVertex3f (0.0, 0.0, 0.0)
		glVertex3f (0.0, 0.0, 1.0)

		glEnd ();

		glutSwapBuffers ()

	def keyboardSpecial (self, key, x, y):
		"""glut spacial callback"""
		del x, y	# unused
		# arrow keys for panning
		if key == GLUT_KEY_UP:
			self.motion.yTrans -= 0.25
		elif key == GLUT_KEY_DOWN:
			self.motion.yTrans += 0.25
		elif key == GLUT_KEY_LEFT:
			self.motion.xTrans -= 0.25
		elif key == GLUT_KEY_RIGHT:
			self.motion.xTrans += 0.25
		# page up/down for zooming translation
		elif key == GLUT_KEY_PAGE_UP:
			self.motion.zTrans -= 0.5
		elif key == GLUT_KEY_PAGE_DOWN:
			self.motion.zTrans += 0.5
		# finally, redisplay
		glutPostRedisplay ()

	def keyboard (self, key, x, y):
		"""glut keyboard callback"""
		del x, y	# unused
		if ord (key) == 27 or key == 'q':
			sys.exit (0)

		if key in ('+', '-', 'l', 'r', 'L', 'P', 'S', 'F', '?', 'n', 'N'):
			if key == 'r' and self.attractor.name != 'Roessler':
				self.attractor = Roessler (a = 0.2, b = 0.2, c = 5.7, h = 0.05, steps = 10000)
			elif key == 'l' and self.attractor.name != 'Lorenz':
				self.attractor = Lorenz (a = 10.0, b = 28.0, c = 8.0 / 3.0, h = 0.01, steps = 10000)
			elif key == '+':
				self.attractor.setParameters ("steps", int(self.attractor.steps * 1.10))
			elif key == '-':
				self.attractor.setParameters ("steps", int(self.attractor.steps * 0.90))
			elif key == 'L':
				self.attractor.setMode (Modes.lines)
			elif key == 'P':
				self.attractor.setMode (Modes.points)
			elif key == 'S':
				self.attractor.setMode (Modes.spheres)
			elif key == 'F':
				if self.hasFog:
					glDisable (GL_FOG)
				else:
					glEnable (GL_FOG)
				self.hasFog = not self.hasFog
			elif key == 'n':
				self.oldAttr = copy.deepcopy (self.attractor)
				for i in [k for k in self.attractor.__dict__.keys () if k in "abcdefg"]:
					self.attractor.setParameters (i, self.rand.uniform (0.1, 3.5))
			elif key == 'N' and getattr (self, "oldAttr"):
				self.attractor = self.oldAttr
				del self.oldAttr
			elif key == '?':
				self.showDoc ()
				return
			glutPostRedisplay ()
			return

		param = key.lower ()

		# check attractor parameters
		if param not in [i for i in dir (self.attractor) if not callable (getattr (self.attractor, i)) and len (i) == 1]:
			print >> sys.stderr, "%s: no such parameter for %s!" % (param, self.attractor.name)
			return
		pref = getattr (self.attractor, param)
		if type (pref) != types.IntType and type (pref) != types.LongType and type (pref) != types.FloatType:
			print >> sys.stderr, "%s: is no number prameter in %s!" % (param, self.attractor.name)
			return

		if param == key: # key.lower () == key
			self.attractor.setParameters ((param, pref - (pref * 0.1)))
		else:
			self.attractor.setParameters ((param, pref + (pref * 0.1)))

		glutPostRedisplay ()

	def showDoc (self):
		print self.doc

if __name__ == '__main__':
	app = StrangeApp ()
	app.run ()