#! /usr/bin/env python
"""
newton-strings.py is a slight alteration of cradle.py in which now we 
include a dampening coefficient so that the cradle approximately
stops after about 10 seconds.


PHZ4151C
Mini-Exam 1, Problem 5
Feb 1, 2019
"""
from __future__ import division, print_function
import sys
import vpython as vs
from math import sin, cos, sqrt, exp

if len(sys.argv) == 2:
	mu = float(sys.argv[1])
else:
	mu = 0


g = 9.81	#acceleration due to gravity, m/s^2
L = 1		#length of the pendulums, m
theta_0 = 0.5	#max angular displacement, radians
r = 0.1		#ball radius, m
t,dt = 0, 0.01	#initialize time and set time evolution, s
#mu = .25	#dampening coefficient

c1 = vs.vector(-r,L/2,0) #these are the points that each ball
c2 = vs.vector(r,L/2,0)	 #is swinging about

#initial position of each ball, relative to c1/c2 respectively
p1 = vs.vector(0,-L,0)
p2 = vs.vector(L*sin(theta_0), -L*cos(theta_0), 0)

#create two balls and position them according to the initial 
#positions above
b1 = vs.sphere(pos=p1+c1, radius=r, color=vs.color.red)
b2 = vs.sphere(pos=p2+c2, radius=r, color=vs.color.blue)

#create two "ropes" that swing about c1/c2 and whose other end 
#matches the position of their respective ball
r1 = vs.cylinder(pos=c1, axis=p1, radius=r/10)
r2 = vs.cylinder(pos=c2, axis=p2, radius=r/10)

#animate the Newton's cradle
while(True):
	vs.rate(1/dt)
	t += dt

	#use equation for angular position of a simple pendulum
	#to determine theta at time t
	#now theta_t is now exponentially decaying with dampening
	#coefficient mu
	theta_t = theta_0*cos(sqrt(g/L)*t)*exp(-mu*t)

	p = L*vs.vector(sin(theta_t), -cos(theta_t), 0)
	
	#for negative theta_t swing left ball
	if theta_t < 0:
		b1.pos = p+c1
		r1.axis = p

	#for positive swing right ball
	else:
		b2.pos = p+c2
		r2.axis = p