#!/bin/python
from vpython import *


size = .25
theta = pi/4
init_pos = vec(-15, size, 0)
init_v = vec(20*cos(theta), 20*sin(theta), 0)

dt = 1e-4
C_drag = 0.9
g = 9.8

bgc = vec(59/256, 114/256, 165/256)

scene = canvas(center = vec(0,5,0), width=600, background=bgc)
floor = box(length=50, height=0.01, width=4, color=color.blue)
ball = sphere(radius = size, color=color.red, make_trail = True, trail_radius = 0.05) # the ball

table = graph(width = 300, align = 'left')
fv = gcurve(graph = table, color=color.blue, width=4)

a1 = arrow(color = color.green, shaftwidth = 0.1)
a1.axis = vec(1, -1, 0)
ball.pos = init_pos
ball.v = init_v

t = 0
cnt_g = 0
pcnt = 0
sum_dis = 0
max_h = 0

def hypot(a, b):
    return (a*a+b*b)**.5

def movetheball(dt, ball):
    global sum_dis
    global max_h
    max_h = max(max_h, ball.pos.y)
    ball.pos += ball.v*dt
    sum_dis += mag(ball.v) * dt
    ball.v += vec(0, -g, 0) * dt - C_drag*ball.v*dt
    if ball.pos.y <= size and ball.v.y < 0: # new: check if ball hits the ground
        ball.v.y *= -1
        return 1;
    return 0;

def putarrow():
    a1.axis = ball.v / 5
    a1.pos = ball.pos - a1.axis
    a1.axis += vec(0, 0, 1)
    a1.pos += vec(0, 0, 1)

while cnt_g < 3:
    rate(1/dt)
    t += dt
    pcnt += 1
    cnt_g += movetheball(dt, ball)
    if pcnt == 1000:
        fv.plot( pos=(t, hypot(ball.v.x, ball.v.y)) )
        pcnt = 0
    putarrow()


show_pos = vec(-10, 14, 0)
diff_pos = vec(0, 2, 0)
text(text="Displacement:" + str(mag(ball.pos)), pos = show_pos)
text(text="Distance Traveled:" + str(sum_dis), pos = show_pos-diff_pos)
text(text="Maximum Height:" + str(max_h), pos = show_pos-diff_pos*2)