# This program uses the Euler-Cromer method to update the position of a ball in response to
# a force, here defined as a spring force in the x direction.  It has three types of output;
# it annimates a ball and traces it's trajectory, it plots the x component of the ball's
# positon against time, and it prints the final momentum and position of the ball.
# This current version gives the ball a mass of 1kg so that the ball's velocity is numerically
# equal to its momentum.  This should be a useful 'shell' for you to modify to address a
# variety of computational problems.

from __future__ import division
from visual import *
from visual.graph import *

# Constants
k=5
b=0
t=0
tmax = 60
dt=0.001

# Objects
ballec=sphere(pos = (10,0,0), p=vector(0,0,0), m=1, radius = 0.1, color=color.red)
                        # the above line defines a ball with initial position (10,0,0)m, 
                        # initial momentum (0,0,0)kg m/s, a radius of 0.1 m, and a mass of 1 kg.

balle=sphere(pos = (10,0,0), p=vector(0,0,0), m=1, radius = 0.1, color=color.cyan)

trail1=curve(color = ballec.color)     # This sets up the program for creating a curve which we will,
                                    # inside the loop below, choose to trace the path of the ball
trail2 = curve(color = balle.color)
funct1 = gcurve(color=ballec.color)
funct2 = gcurve(color = balle.color)
funct1.plot(pos=(t,ballec.x))      # Plot the initial point
funct2.plot(pos =(t,balle.x))
trail1.append(pos=ballec.pos)  # Start's tracing the ball's trail at it's initial position
trail2.append(pos=balle.pos)
Fe = -k*vector(balle.x,0,0)
#Execution              
while t<tmax:               # Observe that the loops is reentered only as long as t<tmax,
                            # so final time at which it enters will be between tmax-dt and tmax.
    Fec=-k*vector(ballec.x,0,0)                     # Calculate Force using the current x component
    balle.pos = balle.pos + (balle.p/balle.m)*dt
    ballec.p = ballec.p + Fec*dt                      # Use the force to update the momentum
    balle.p = balle.p + Fe*dt
    ballec.pos = ballec.pos + (ballec.p/ballec.m)*dt    # Use the updated momentum to update the position
    Fe = -k*vector(balle.x,0,0)
    t=t+dt                                      # Advance the time by a step dt
    funct1.plot(pos=(t,ballec.x))                 # Plot the new point
    funct2.plot(pos=(t,balle.x))
    trail1.append(pos=ballec.pos)
    trail2.append(pos=balle.pos)
    #rate(10000)                                  # This sets the rate with which the loop is executed;
                                                # to calculate as quickly as possible, comment it out,
                                                # but to slow down the simulation, decrease the rate value
print "at t = ",t," the momentum is ", ballec.p, "kg m/s, and the position is ", ballec.pos, "m."