#!/usr/bin/python
# LorenzODE.py: The Lorenz chaotic attractor in interactive 3D.
#       Note the use of the Euler integration method.

from visual import *
from Agent import *

print """
Right button drag to rotate camera to view scene.
Middle button drag up or down to zoom in or out
"""
scene2 = display(title="scene 2")
#scene2.center = vector(0,0,25)
#scene2.lights = [vector(-0.25,-0.5,-1.0),vector(0.25,0.5,1.0)]
#scene2.height = 600
#scene2.width = 600
#scene2.background = (.5,.5,.5)
#scene2.autocenter = 1
#scene2.forward = (0,1,-0.5)
scene.title = "Lorenz differential equation"
#scene.center = vector(0,0,25)
# If you can cross your eyes to see stereo, uncomment the next two lines.
#scene.stereo = 'crosseyed'
#scene.stereodepth = 2
#scene.autocenter = 1
#scene.height = 600
#scene.width = 600
#scene.background = (1,1,1)
#scene.forward = (0,1,-0.5)
#scene.lights = [ vector(-0.25, -0.5, -1.0) ,vector(0.25, 0.5, 1.0) ]

# Draw axes
#curve( pos = [ (-25,0,0), (25,0,0) ], color = (0,0,0), radius = 0.2 )
#curve( pos = [ (0,-25,0), (0,25,0) ], color = (0,0,0), radius = 0.2 )
#curve( pos = [ (0,0,0), (0,0,50) ], color = (0,0,0), radius = 0.2 )

lorenz = curve( color = color.blue, radius=0.2,display=scene )
lorenzz = curve(color=color.yellow,radius=0.2,display=scene2)

dt = 0.01
y = vector(10,24,33)
sigma = 10.0
R = 28.0
b = -2.5
theta = 3.0 * pi / 4.0
z = vector(1,1,1) 

#define a function that returns a lorenz
#field with sigma, R and b fixed
def lorenzField(sigma,R,b):
	return lambda v: vector(sigma*(v[1]-v[0]),R*v[0]-v[0]*v[2]-v[1],b*v[2]+v[0]*v[1])

#create the function to use in the integrator
lz = lorenzField(sigma,R,b)
integrator = NewEulerIntegrator(0.9)
#create an agent with initial position vector (-10,-7,35)
orbiter1 = Agent(y,lz,integrator)
#create an agent with initial position vector (1,1,1)
orbiter2 = Agent(z,lz,integrator)
#set the neighbor attribute so the integrator can take 
#accound of neighbor's field
orbiter1.neighbor = orbiter2
orbiter2.neighbor = orbiter1

for t in arange(0,100,dt):
  rate(200)
  dydt = vector(lz(y))
  orbiter1.move(dt)
  orbiter2.move(dt)
  orbiter1.update()
  orbiter2.update()
  y = orbiter1.pos
  z = orbiter2.pos
  # Draw lines colored by speed
  #c = clip( [mag(dydt) * 0.005], 0, 1 )[0]
  #xt = cos(theta) * y[0] - sin(theta) * y[1]
  #yt = sin(theta) * y[0] + cos(theta) * y[1]
  lorenz.append(y)
  lorenzz.append(z)
  #lorenz.append( pos=y, color=(c,0.1, 1-c) )