from visual import *
from visual.graph import *

scene.width = 200
scene.height = 600
gray = (0.7,0.7,0.7)
scene.background = gray
ball1 = sphere(pos=(0,100,0), radius = 2, color = color.red)
ball2 = sphere(pos=(5,100,0), radius = 2, color = color.yellow)

scene.autoscale = 0
#params for air resistance calc
rho = 1.29 #density of air at sea level
g = 9.81 #acceleration of gravity
#ball1
ball1.mass = .1110
ball1.R = .04
ball1.A = pi*(ball1.R)*(ball1.R)
#ball2
ball2.mass = .7110
ball2.R = .04
ball2.A = pi*(ball2.R)*(ball2.R)
Cp1 = 1. #drag coefficient
Cp2 = 1.
graph1 = gdisplay(x=0, y=0, width=600, height=550, 

             title='Velocity vs. Time', xtitle='t(s)', ytitle='V (m/s)', 
             xmax=5., xmin=0., ymax=50, ymin=0, 
             foreground=color.white, background=color.black)


velPlot1 = gcurve(color=color.red)
velPlot2 = gcurve(color=color.yellow)

#initialize the program
ball1.velocity = vector(0,0,0) #start from rest
ball2.velocity = vector(0,0,0) #start from rest

ball1.acceleration = vector(0,-9.81,0)
ball2.acceleration = vector(0,-9.81,0)
time = 0
dt = 0.005
rate(100)
while time <5.:
    if scene.mouse.clicked:
        ball1.pos = ball1.pos + ball1.velocity*dt
        ball1.velocity = ball1.velocity+ball1.acceleration*dt

        ball2.pos = ball2.pos + ball2.velocity*dt
        ball2.velocity = ball2.velocity+ball2.acceleration*dt
    
        if ball1.velocity.y < 0:
            velPlot1.plot(pos=(time,mag(ball1.velocity)))
            ball1.acceleration.y = -9.81 + 0.5*Cp1*ball1.A*dot(ball1.velocity,ball1.velocity)/(ball1.mass)
        else:
            ball1.acceleration.y = 0

        if ball2.velocity.y < 0:
            velPlot2.plot(pos=(time,mag(ball2.velocity)))
            ball2.acceleration.y = -9.81 + 0.5*Cp2*ball2.A*dot(ball2.velocity,ball2.velocity)/(ball2.mass)
        else:
            ball2.acceleration.y = 0
        
        time = time + dt