from visual import *

# this function does the actual movement
# of an object, given the force on it
# and time it acts
# it also moves fake shadows around
def mover ( obj, force, dt, shadow ):

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

    # check limits on movement
    # the floor
    if ( obj.pos.y < 0.0 + obj.radius ):
        # next line prevents obj digging in on bounces
        obj.pos.y = - (obj.pos.y - obj.radius ) + obj.radius
        obj.velocity.y = - obj.velocity.y

    # left wall
    if ( obj.pos.x < -10.0 + obj.radius ):
        obj.velocity.x = - obj.velocity.x

    # right wall
    if ( obj.pos.x > 10.0 - obj.radius ):
        obj.velocity.x = - obj.velocity.x

    # back wall
    if ( obj.pos.z < -10.0 + obj.radius ):
        obj.velocity.z = - obj.velocity.z

    # front wall
    if ( obj.pos.z > 10.0 - obj.radius ):
        obj.velocity.z = - obj.velocity.z

    # move shadows
    shadow.pos.x = obj.pos.x
    shadow.pos.z = obj.pos.z

    return
    
# make two balls bounce around and collide
# set up balls
ballRad = 1.5
ball1 = sphere ( pos=( 8, ballRad, 0 ), radius = ballRad, color= color.red )
ball2 = sphere ( pos=(-8, ballRad, 2.5 ), radius = ballRad, color=color.yellow )
# initial velocities, masses
ball1.velocity = vector( -0.25, 0, 0 )
ball2.velocity = vector( 3, 0, 5 )

ball1.mass = 5.0
ball2.mass = 3.0

# ball shadows
shadow1 = cylinder ( pos=( ball1.pos.x,0,ball1.pos.z ),
                     radius=ballRad, length= 0.05,
                     color=color.black, axis= (0,0.05,0) )
shadow2 = cylinder ( pos=( ball2.pos.x,0,ball2.pos.z ),
                     radius=ballRad, length= 0.05,
                     color=color.black, axis= (0,0.05,0) )

# make world
floor = box ( size = (20, 0.5, 20 ))
floor.pos.y = -0.25
backwall = box ( size = ( 20,20, 0.5 ), pos= ( 0,10,-10.25 ))
scene.autoscale = 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

# cm is a cube
cm = box ( color= color.blue, size= (0.5,0.5,0.5))
cm.pos = (ball1.mass*ball1.pos + ball2.mass*ball2.pos)/(ball1.mass+ball2.mass)

# big loop
while 1:
    rate(100)
    
    # move balls- find forces and move them
    f1 = vector( 0,1,0 ) * ball1.mass * g
    mover ( ball1, f1, dt, shadow1 )

    f2 = vector( 0,1,0 ) * ball2.mass * g
    mover ( ball2, f2, dt, shadow2 )
    
    cm.pos = (ball1.mass*ball1.pos + ball2.mass*ball2.pos)/(ball1.mass+ball2.mass)

    #check for collisions
    n = ball2.pos - ball1.pos
    if ( mag(n) < (ball1.radius + ball2.radius) ):

        # they collided
        print "collision!"
        
        # figure out collision
        # set up some variables for collision
        # analyze collision in tiny time steps
        dt2 = dt/100
        nhat = norm(n)
        # closing velocities
        v1 = ball1.velocity
        v2 = ball2.velocity
        v1n = dot ( v1, nhat )
        v2n = dot ( v2, nhat )
        # force applied at each time step- a guess
        dNorm = (ball1.mass + ball2.mass)*(v1n-v2n)*10.0
        # number of time steps
        nHits = 0

        # have forces work on balls, then see if velocities are still
        # closing, or are they moving apart yet
        while ( v2n-v1n ) < 0:
            # find new closing velocities
            v1 = ball1.velocity
            v2 = ball2.velocity
            v1n = dot ( v1, nhat )
            v2n = dot ( v2, nhat )
            # scraping velocity
            v1s = ( v1-v2 ) - nhat * ( dot((v1-v2), nhat ) )
            v1shat = norm(v1s)
            # total force applied each time step to ball 2
            dF = nhat * dNorm + v1shat * muk * dNorm
            # apply dF to each ball for dt2 time
            mover ( ball1, -dF, dt2, shadow1 )
            mover ( ball2, dF, dt2, shadow2 )
            nHits = nHits + 1
            
        # balls have now stopped closing on each other
        # now they do forces as they spread apart
        print " collision took ", nHits, " steps"
        while ( nHits > 0 ):
            mover ( ball1, -dF*epsilon, dt2, shadow1 )
            mover ( ball2, dF*epsilon, dt2, shadow2 )
            nHits = nHits -1