from visual import *

# this function does the actual movement
# of an object, given the force on it
# and time it acts
# this version doesnt do shadows
#
# 10/31/02- it also does rotation around
# the y-axis. the R argument is the
# distance from the center of mass of the object

def moveAndRotate ( obj, force, R, dt ):

    # find acceleration of CM
    acc = vector (0,0,0)
    acc = force / obj.mass
    
    # move object CM
    obj.velocity = obj.velocity + acc * dt
    obj.pos = obj.pos + obj.velocity * dt

    # now find torque and rotate it
    torque = cross ( R, force )
    alpha = torque / obj.I
    obj.angVel = obj.angVel + alpha *dt
    obj.angPos = obj.angPos + obj.angVel *dt

    # note rotate needs change to theta, not new theta
    obj.rotate ( angle = obj.angVel.y * dt, axis=(0,1,0) )

    return 
    
# make a stick slide around and bounce off a wall
# set up stick

stickL = 3
stickH = 1
stickW = 1

stick = box ( pos=( 0, 0.5, 5 ), length = stickL, height=stickH, width=stickW, color=color.red )

# initial velocities, masses
stick.velocity = vector( 0, 0, -0.2 )
stick.mass = 5.0
stick.I = stick.mass * ( stickL*stickL + stickW*stickW )/12.0

stick.angVel = vector( 0, 0.1, 0 )
stick.angPos = vector( 0, 0.3, 0 )
stick.rotate ( angle = stick.angPos.y, axis=(0,1,0) )

cornerR = math.sqrt( stickL*stickL/4.0 + stickW*stickW/4.0 )
cornerPhi = math.atan(stickW/stickL)

# make world
floor = box ( size = (20, 0.5, 20 ), pos = ( 0, -0.25, 10 ), color= color.green)

scene.autoscale = 0
scene.background = ( 0.8,0.8,1.0)
scene.lights = [ vector ( 0,1,0 ) ]

#collision epsilon- fraction of velocity conserved after collision
epsilon = 0.9

# mu-sub-k
muk = 1.0

# world vars
g = -2
dt =0.04
drag = 0.01

# big loop
while 1:
    rate(100)
    
    # move stick- find forces and move them
    force = vector(0,0,0)
    # vector from CM to corner of stick (needed for torque in moveAndRotate )
    stickR = cornerR * vector( math.sin(stick.angPos.y),0,math.cos(stick.angPos.y) )
    
    moveAndRotate ( stick, force, stickR, dt)

    # has it hit the wall at z=0? have any of its corners?
    # lets find the corners
    phis = [ cornerPhi, -cornerPhi ]
    thetas = [ 0, 3.14159 ]

    # cycle thru corners
    for phi in phis:
        for theta in thetas:
            
            # find vector from CM to corner
            relCornerPos = cornerR * vector( math.cos( stick.angPos.y+theta+phi ),0, math.sin( stick.angPos.y+theta+phi ) )
            # add to cm pos to get abs position of corner
            cornerPos = stick.pos + relCornerPos

            # check if this is less than z=0
            if cornerPos.z < 0:

                # it is, we have a corner inside the wall
                print "Hit the wall!!!"
                energy = 0.5 * stick.mass * mag(stick.velocity) * mag(stick.velocity)
                energy = energy + 0.5 * stick.I * stick.angVel.y * stick.angVel.y
                print "energy before collision is ",energy

                # do a mini force loop until it has bounced out
                # n is perp to surface
                n = vector( 0,0,1 )
                
                dt2 = dt /100
                # guess at force
                dF = stick.mass * mag (stick.velocity)/dt2 * n /100

                # find closing speed
                # it is stick vcm +
                vCorner = stick.velocity
                # omega R tangent to circle
                tan2Circle = vector( math.sin(stick.angPos.y+theta+phi),0,
                                     math.cos(stick.angPos.y+theta+phi))
                vCorner = vCorner + cornerR * stick.angVel.y * tan2Circle
                # closing speed is vCorner-dot-normal
                closingSpeed = dot( vCorner, n )

                #now loop until no more closing v
                nHits = 0
                while ( closingSpeed < 0 ):
                    
                    # do some force and torque on stick
                    # find vector from CM to corner for torque
                    relCornerPos = cornerR * vector( math.cos( stick.angPos.y+theta+phi ),0,
                                                     math.sin( stick.angPos.y+theta+phi ) )
                    # do force and torque on stick
                    moveAndRotate ( stick, dF, -relCornerPos, dt2)
                    nHits = nHits + 1

                    # find new closing speed
                    vCorner = stick.velocity
                    tan2Circle = vector( math.sin( stick.angPos.y + theta + phi),0,
                                     math.cos( stick.angPos.y + theta + phi))
                    vCorner = vCorner + cornerR * stick.angVel.y * tan2Circle
                    closingSpeed = dot( vCorner, n )

                # okay time to rebound
                print "collision took ", nHits, " hits"
                while ( nHits > 0):
                    moveAndRotate ( stick, dF*epsilon, -relCornerPos, dt2)
                    nHits = nHits - 1
                    
                # print energy after collision
                energy = 0.5 * stick.mass * mag(stick.velocity) * mag(stick.velocity)
                energy = energy + 0.5 * stick.I * stick.angVel.y * stick.angVel.y
                print "energy after collision is ",energy