from visual import *
from visual.graph import *
Tmax=100.0
MAX=0.001

pic1=gdisplay(x=0,y=0,width=400,height=400,title="Dtheta vs t",
             xtitle="t",ytitle="theta-diff",xmax=Tmax,xmin=0,ymin=-MAX,ymax=MAX,
             foreground=color.black, background=color.white)
pic2=gdisplay(x=400,y=0,width=400,height=400,title="Domega vs t",
             xtitle="t",ytitle="omega-diff",xmax=Tmax,xmin=0,ymin=-MAX,ymax=MAX,
             foreground=color.black, background=color.white)
plot1=gcurve(gdisplay=pic1,color=color.blue)
plot2=gcurve(gdisplay=pic2,color=color.red)

## don't display the ball - boring
scene.visible=0
def force(pos,vel,t):
    result = vector(0,0,0)
    g=9.81
    L=9.81
    k=0.5
    F=0.5
    alpha=2.0/3.0
    result.x=-(g/L)*sin(pos.x)-k*vel.x+F*sin(alpha*t)
    result.y=0.0
    result.z=0.0
    return result
    
    
ball1=sphere(pos=vector(0.2,0.0,0),vel=vector(0,0,0))
ball2=sphere(pos=vector(0.201,0,0),vel=vector(0,0,0))
dt=0.01
t=0

count=0
while (t<Tmax) :
#    rate(1000)
    t=t+dt
# leapfrog alg.
    a1_initial=force(ball1.pos,ball1.vel,t)
    ball1.pos=ball1.pos+ball1.vel*dt+a1_initial*0.5*dt*dt
    a2_initial=force(ball2.pos,ball2.vel,t)
    ball2.pos=ball2.pos+ball2.vel*dt+a2_initial*0.5*dt*dt

# theta angular variable - restrict to -pi->+pi    
    if (ball1.pos.x>math.pi):
        ball1.pos.x=ball1.pos.x-2*math.pi
    if (ball1.pos.x<-math.pi):
        ball1.pos.x=ball1.pos.x+2*math.pi
    
    if (ball2.pos.x>math.pi):
        ball2.pos.x=ball2.pos.x-2*math.pi
    if (ball2.pos.x<-math.pi):
        ball2.pos.x=ball2.pos.x+2*math.pi
        
    a1_final=force(ball1.pos,ball1.vel,t)
    ball1.vel=ball1.vel+(a1_initial+a1_final)*dt*0.5
    a2_final=force(ball2.pos,ball2.vel,t)
    ball2.vel=ball2.vel+(a2_initial+a2_final)*dt*0.5
    diff=ball1.pos.x-ball2.pos.x
    if (diff>math.pi):
        diff=diff-2*math.pi
    if (diff<-math.pi):
        diff=diff+2*math.pi

    #don't plot every data point since the gcurve has limitations
    #
    count += 1
    if (count%10==0):
        plot1.plot(pos=(t,diff))
        plot2.plot(pos=(t,ball1.vel.x-ball2.vel.x))