from visual import *

print """
Bruce Sherwood Fall 2000
Click to cycle through different angular momentum cases:
   rotation == 0: barbell orientation doesn't change;
      (mounted on frictionless axle, no torque acts on it)
   1: barbell rotates at same rate as rod
   2: barbell rotates a lot
"""
scene.width=scene.height=800
scene.x = scene.y = 0
scene.background = color.white
rmass = 0.2
L = 2.
rod=cylinder(pos=(0,0,0), axis=(4,0,0), radius=0.03, color=(1,.9,0))
barbell = frame()
s1=sphere(frame = barbell, pos=(0,L/2.,0), radius=rmass, color=(1,0,0))
s1.mass=0.01
##s1.trail=curve(color=s1.color)
s2=sphere(frame = barbell, pos=(0,-L/2.,0), radius=rmass, color=(0,0,1))
s2.mass=0.01
##s2.trail=curve(color=s2.color)
rd=cylinder(frame = barbell, pos=s1.pos, axis=(s2.pos-s1.pos),
            color=(1,1,1), radius=0.04)
barbell.trail=curve(color=(.6,.6,.6))
##cc=sphere(pos=barbell.pos+barbell.axis, radius=0.2, color=color.yellow)
##dd=sphere(pos=barbell.pos-barbell.axis, radius=0.2, color=color.green)

barbell.Icm = 2*s1.mass*(L/2.0)**2
barbell.pos = rod.pos+rod.axis
barbell.Iorig = 2*s1.mass*(mag(rod.axis)**2)

omegaCM = vector(0,0,pi)
omega = vector(0,0,pi/5.)
##omegaCMmag = mag(omegaCM)
dt = 0.05
t = 0.0
scene.range=5
Lscale = 2.0/0.1

LT=arrow(pos=rod.pos, axis=(0,0,0), color=color.cyan,
         shaftwidth = 0.2)
LR=arrow(pos=barbell.pos, axis=(0,0,0), color=color.green,
         shaftwidth = 0.2)
##LTOT = arrow(pos=rod.pos, axis=(0,0,0), color=color.green,
##             shaftwidth = 0.3)
print scene.lights


rotation = 0
direction = 1
while 1:
    while 1:
        if scene.mouse.clicked:
            scene.mouse.getclick()
            break
        rate(30)
        dtheta = mag(omegaCM)*dt
        dphi = mag(omega)*dt
        rod.rotate(angle=dphi, axis=omega, origin=(0,0,0))
        barbell.pos = rod.pos+rod.axis
        Ltrans = barbell.Iorig*omega
        Lrot = vector(0,0,0)
        if rotation == 1:
            barbell.rotate(angle=dphi, axis=omegaCM, origin=(barbell.pos))
            Lrot = barbell.Icm*omega 
        if rotation == 2:
            barbell.rotate(angle=direction*dtheta, axis=omegaCM, origin=(barbell.pos))
            Lrot = direction*barbell.Icm*omegaCM
        LT.axis = Ltrans*Lscale
        LR.pos = barbell.pos
        LR.axis = Lrot*Lscale
    ##    LTOT.axis = (Ltrans+Lrot)*Lscale
    ##    s1.trail.append(pos=barbell.pos-barbell.axis/2.0)
    ##    s2.trail.append(pos=barbell.pos+barbell.axis/2.0)
    ##    if t < 10*dt: print barbell.pos, barbell.axis
        t = t+dt
    ##    if t < 2*pi/mag(omega):
    ##        barbell.trail.append(pos=barbell.pos)
    rotation = rotation+1
    if rotation > 2:
        rotation = 0
        direction = -direction