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

Ruth Chabay Fall 2009 added flags to make arrows and trails visible / invisible
"""
scene.width=scene.height=800
scene.x = scene.y = 0
scene.background = color.white

def world_space_pos(frame, local):
    x_axis = norm(frame.axis)
    z_axis = norm(cross(frame.axis, frame.up))
    y_axis = norm(cross(z_axis, x_axis))
    return frame.pos+local.x*x_axis+local.y*y_axis+local.z*z_axis

# set flag to 1 to show arrow or trail; 0 not to show it
SHOW_LTrans = 1
SHOW_LRot = 1
SHOW_trail1 = 1
SHOW_trail2 = 1

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, visible = SHOW_trail1)
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, visible = SHOW_trail2)
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))

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.)
dt = 0.05
t = 0.0
scene.range=5.5
Lscale = 2.0/0.1

LT=arrow(pos=rod.pos, axis=(0,0,0), color=(.7,.5,0),
         shaftwidth = 0.2, visible = SHOW_LTrans)
LR=arrow(pos=barbell.pos, axis=(0,0,0), color=(.7,.5,0),
         shaftwidth = 0.2, visible = SHOW_LRot)
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
        s1.trail.append(world_space_pos(barbell,s1.pos))
        s2.trail.append(world_space_pos(barbell,s2.pos))
        t = t+dt
    rotation = rotation+1
    s1.trail.pos = []
    s2.trail.pos = []
    if rotation > 2:
        rotation = 0
        direction = -direction