###################################################################################################
# Based on the szglorenz.cpp program of Wayne City High School (Illinois). Thanks Stan Blank!
# This version is a generic dynamical system machine, by Kyle Wilkinson
# Updated May 8, 2006
###################################################################################################

from PySZG import *
import sys
import time
import random
from math import sin, cos, pi
#font = arTexFont()
widgetList = []

wIdentity = ar_identityMatrix()
mIdentity = ar_identityMatrix()
nIdentity = ar_identityMatrix()
tIdentity = ar_identityMatrix()

class navigator:
	def __init__(self):
		self.node = None
		self.lastTime = -1
		# The speed in units per second.
		self.speed = 5.0
                # Threshold on control (-1,1)
		self.threshold = 0.4

	def update(self, framework, time):
		# If this is the first time calling the method, just set the time called and return.
		if self.lastTime < 0:
			self.lastTime = time
			return
	
		wandDir = ar_ERM(framework.getMatrix(1))*arVector3(0,0,-1)
		wandTransDir = ar_ERM(framework.getMatrix(1))*arVector3(1,0,0)
		deltaX = framework.getAxis(0)
		deltaY = framework.getAxis(1)
		if abs(deltaX) < self.threshold:
			deltaX = 0
		else:
			deltaX = (deltaX - self.threshold) / (1.0 - self.threshold)
		if abs(deltaY) < self.threshold:
			deltaY = 0
		else:
			deltaY = (deltaY - self.threshold) / (1.0 - self.threshold)
		# t is the egocentric translation vector.
		t = (wandDir * deltaY + wandTransDir * deltaX) * (time - self.lastTime) * self.speed 

		self.node.set(ar_TM(-t)*self.node.get())
		self.lastTime = time
	
class manipulator:
	# The constructor.
	def __init__(self):
		# The transform node to be manipulated.
		self.node = None
                self.worldNode = None
                self.switchSys = False
                self.switchView = False
                self.updatePerspective = False
	        self.wandMatrixAtGrab = None
                self.objMatrixAtGrab = None

                
        # Start grabbing the object.
        def attemptGrab(self, matrix):
            if not self.node:
                return
            self.objMatrixAtGrab = self.node.get()
            self.wandMatrixAtGrab = matrix

        # Release a grabbed object (if any is currently grabbed).
        def attemptRelease(self):   
            self.objMatrixAtGrab = None
            # If an object is grabbed, change its transform in response to user interface events.
        
        def attemptDrag(self, matrix):
            # Check to see if the object is currently grabbed.
            if self.objMatrixAtGrab:
                tmp = ar_ETM(self.wandMatrixAtGrab)
                trans = ar_ETM(matrix)*tmp.inverse()*ar_ETM(self.objMatrixAtGrab)
                rot = ar_ERM(matrix)*ar_ERM(self.wandMatrixAtGrab.inverse())*ar_ERM(self.objMatrixAtGrab)
                self.node.set(trans*rot)
         
        # Processes the event stream. We also need access to the framework object (to get current event values).
	def update(self, framework, e):
            if e.getType() == AR_EVENT_BUTTON:
                if e.getButton() == 1:
                    if e.getIndex() == 1:
                        if self.updatePerspective == False:
                            self.updatePerspective = True
                            self.switchView = True
                        else:
                            self.updatePerspective = False
                            self.switchView = True
                    elif e.getIndex() == 0:    
                        self.switchSys = True
                    elif e.getIndex() == 2:
                        self.attemptGrab(framework.getMatrix(1))
                else:
                    self.attemptRelease()
            self.attemptDrag(framework.getMatrix(1)) 
                    

class particle:
	# The constructor.
	def __init__(self):
		self.x = arVector3(0,5,20)
		self.oldX = arVector3(0,0,0)
		self.useTail = 1
		self.lineWidth = 1.0
		self.tailPosition = 0
		self.tailLength = 100
		self.pointsNode = None 
		self.drawable = None
		self.mat = arMaterial()
		self.mat.diffuse = arVector3(1,0,0)
		self.mat.specular = arVector3(.3, .3, .3)
		self.mat.emissive = arVector3(.3,.3,.3)
		self.mat.exponent = 100
		self.type = 'cube'
		self.trans = None
		self.aCount = 0
		self.xeq = None
                self.yeq = None
                self.zeq = None
                self.h = None

	def attach(self, n):
		m = n.new("material")
		m.set(self.mat)
		if self.useTail:
			self.pointsNode = m.new("points")
			state = self.pointsNode.new("graphics state")
			state.set(("line_width", self.lineWidth))
			self.drawable = state.new("drawable")
			self.drawable.set(("lines",self.tailLength))
		self.trans = m.new("transform")
		self.trans.set(ar_TM(self.x))
		scl = self.trans.new("transform")
		scl.set(ar_RM('y', 6.283*random.random())*ar_SM(.5))
		if self.type == 'cube':
			s = arCubeMesh()
			s.attachMesh(scl)	

		elif self.type == 'sphere':
			s = arSphereMesh()
			# The number of divisions in the sphere. It will have about divisions*divisions*2 triangles.
			s.setAttributes(8)
			s.attachMesh(scl)
			
	def advance(self):
		x = self.x[0]
		y = self.x[1]
		z = self.x[2]
		
                self.oldX[0] = x
		self.oldX[1] = y
		self.oldX[2] = z
		h = self.h
		for i in range(5):
			x = self.xeq(x,y,z,h)
			y = self.yeq(x,y,z,h)
			z = self.zeq(x,y,z,h) 
		    
		# Demarshall the position into object storage.
		self.x[0] = x
		self.x[1] = y
		self.x[2] = z
		# Change the position of the sphere/cube.
		self.trans.set(ar_TM(self.x))
		# Update the points used to plot the lines in the trail. The trail is stored in the scene graph
		# as a line set. This means that every even indexed point is the beginning of a line segment, with the
		# following odd indexed point being the end of that segment. Consequently, at each stage we just replace
		# the oldest line segment, giving the effect of a trail of finite length.
		if self.pointsNode:
			self.pointsNode.set((self.x, self.oldX), (2*self.tailPosition, 2*self.tailPosition+1))
			self.tailPosition  += 1
			if self.tailPosition >= self.tailLength:
				self.tailPosition = 0

        #Change the particle's system of equations
        def change(self, xeq, yeq, zeq, h):
            self.xeq = xeq
            self.yeq = yeq
            self.zeq = zeq
            self.h = h

# Argument parsing. The following are valid invocations of the program.
#  python lorenz.py
#    100 particles, no tails, spheres
#  python lorenz.py 200
#    200 particles, no tails, spheres
#  python lorenz.py 100 1
#    100 particles. tails, spheres
#  python lorenz.py 500 0 cube
#    500 particles, no tails, cubes
def processCommandLine(argv, w):
    if len(argv) > 1 :
            number = int(sys.argv[1])
    else:
            number = 100
    if len(sys.argv) > 2 and int(sys.argv[2]) == 0:
            useTail = 0
    else:
            useTail = 1
    if len(sys.argv) > 3 and sys.argv[3] == "cube":
            shape = "cube"
    else:
            shape = "sphere"
    plist = []
    for i in range(number):
            p = particle()
            p.type = shape
            p.useTail = useTail
            # The particles are staggered slightly from the default initial position. 
            # This demonstrated sensitive dependence on initial conditions (one form of chaos).
            p.x[2] = p.x[2] + i*0.03
            p.x[0] = p.x[0] + i*0.03
            p.x[1] = p.x[1] + i*0.03
            # It's a better visualization if the particles start out as a rainbow hue.
            p.mat.diffuse = arVector3(random.random(),random.random(),random.random())
            # Add nodes to the scene graph. 
            p.attach(w)
            plist.append(p)

    return plist  

def addLights(g):
	r = g.getRoot()
	l = r.new("light")
	light = arLight()
	light.lightID = 0
	light.position = arVector4(0,0,-1,0)
        light.ambient = arVector3(.3,.3,.3)
        light.diffuse = arVector3(0.5, 0.5, 0.5)
	l.setLight(light)
	l = r.new("light")
	light = arLight()
	light.lightID = 1
	light.position = arVector4(0,0,1,0)
        light.ambient = arVector3(.3,.3,.3)
        light.diffuse = arVector3(0.5, 0.5, 0.5)
	l.setLight(light)


def eventProcessing(framework, eventQueue):
	while not eventQueue.empty():
		e = eventQueue.popNextEvent()
                for w in widgetList:
			w.update(framework, e)
	return 1

def mainLoop(f, manip, nav, t, w,  plist, xeqs, yeqs, zeqs, hs):
    index = 0
    while 1:
            #switch views
            if manip.switchView:
                manip.switchView = False
                switchView(m = manip.node, n = nav.node, w = w, t = t)

            if manip.updatePerspective:
                updatePerspective(plist[0], t)

            f.processEventQueue()
            nav.update(f, time.clock())
            f.setViewer()
            
            #switch dynamical systems
            if(manip.switchSys):
                manip.switchSys = False
                index = (index + 1) % len(xeqs)
                changeFunctions(plist, xeqs[index], yeqs[index], zeqs[index], hs[index])
                    
            #advance particles one unit    
            for i in plist:
                i.advance()
            
            f.swapBuffers()

def updatePerspective(p, node):
    xPos = arVector3(-p.x[0], -p.x[1], -p.x[2])
    node.set(ar_TM(xPos))

def switchView(w, m, n, t):
    global wIdentity
    tmp = w.get()
    w.set(wIdentity)
    wIdentity = tmp

    global mIdentity
    tmp = m.get()
    m.set(mIdentity)
    mIdentity = tmp

    global nIdentity
    tmp = n.get()
    n.set(nIdentity)
    nIdentity = tmp

    global tIdentity
    tmp = t.get()
    t.set(tIdentity)
    tIdentity = tmp

def initDSG():
    #initialze and create scene graph
    f = arDistSceneGraphFramework()
    f.setEventQueueCallback(eventProcessing)
    f.setAutoBufferSwap(0)
    
    if f.init(sys.argv) != 1:
            sys.exit()
    #g is DSG
    g = f.getDatabase()
    addLights(g)
    #root attach navigational transform
    t = g.getRoot().new("transform")
    
    addStars(5000, t)
    n = t.new("transform")
    
    #attach world transfrom
    w = n.new("transform")
    w.set(ar_TM(0,5,-40)*ar_SM(1.0))

    #attach manipulator transform
    m = w.new("transform")

    #tie navigator to navigational transform
    nav = navigator()
    nav.node = n

    #tie manipulator to manipulator transform
    manip = manipulator()
    manip.node = m
    manip.worldNode = w

    #allows manipulator to process events
    widgetList.append(manip)

    # Starts all the services, windows, etc. 
    if f.start() != 1:
            sys.exit()
    return (f, t, nav, w, manip)

def createFunctions():
    #Lorenz
    xLor = lambda x,y,z,h:(10*y - 10*x)*h + x
    yLor = lambda x,y,z,h:(28*x - y - x*z)*h + y
    zLor = lambda x,y,z,h:(-(2.66667)*z + x*y)*h + z     
    hLor = .004
    
    #Rossler
    xRos = lambda x,y,z,h: -(y+z)*h + x
    yRos = lambda x,y,z,h: (x + .2*y)*h + y
    zRos = lambda x,y,z,h:(0.2+z*(x-4.7))*h + z
    hRos = .04

    xeqs = (xLor, xRos)
    yeqs = (yLor, yRos)
    zeqs = (zLor, zRos)
    hs = (hLor, hRos)
    return (xeqs, yeqs, zeqs, hs)

def changeFunctions(plist, xeq, yeq, zeq, h):
    for i in plist:
        i.change(xeq, yeq, zeq, h)

def addStars(number, pNode):
    addSizedStars(pNode, 4*number/5, 1.0)
    addSizedStars(pNode, 3*number/20, 2.0)
    addSizedStars(pNode, number/20, 3.0)

def addSizedStars(node, number, size):
    state = node.new("graphics state") 
    state.set(("point_size", size)) 
    p = state.new("points") 
    c = p.new("color4") 
    d = c.new("drawable") 
    pointList = [] 
    colorList = [] 
    for i in range(number):
        pointList.append(getStarCoord())
        colorList.append(arVector4(1,1,1,1)) 
    p.set(pointList)
    c.set(colorList) 
    d.set(("points", number))

def getStarCoord():
   random.seed()
   frac = random.random()
   theta = 2*pi*frac
   frac = random.random()
   phi = (frac - 0.49)*pi
   return arVector3(100*cos(theta)*cos(phi), 100*sin(theta)*cos(phi),100*sin(phi)) 

##########################################################################
#                           Main Program                                 #
##########################################################################


f, t, nav, w, manip = initDSG()

#create list of particles based on command line switches
plist = processCommandLine(sys.argv, manip.node)

#create array of functions
xeqs, yeqs, zeqs, hs = createFunctions()

#attach functions
changeFunctions(plist, xeqs[0], yeqs[0], zeqs[0], hs[0])

mainLoop(f, manip, nav, t, w, plist, xeqs, yeqs, zeqs, hs)