## Model 3: Spring-Mass System
## Author: M. Haley
from __future__ import division ## treat integers as real numbers in division
from visual import *
from visual.graph import *
scene.y = 400
scene.width = 600
scene.height = 850

## CONSTANTS

g = 9.8

## OBJECTS AND INITIAL VALUES

L0 = .5 ## equilibrium length
ks = 20 ## spring stiffness
ceiling = box(pos=(0,0,0), size = (0.2, 0.01, 0.2)) ## note location of origin
## define ball
ball = sphere(pos=(0,-L0,0), radius=0.025, color=color.yellow) ## note: spring initially compressed
ball.m = 1
ball.p = ball.m*vector(0,0,0)
## define spring
spring = helix(pos=ceiling.pos, color=color.orange, thickness=.003, coils=40, radius=0.015)
spring.axis = ball.pos - ceiling.pos
## set up time values
omega_analytic = sqrt(ks/ball.m)
T_analytic = 2*pi/omega_analytic
print 'The analytic period is', T_analytic
deltat = T_analytic/10000 ## make this 1/10000 of the period 
t = 0       ## start counting time at zero

## DISPLAY AND GRAPHS
trail = curve(color=ball.color)    ## craft trail: starts with no points
scene.center = vector(0,-L0,0)   ## move camera down to improve display visibility
positionWindow = gdisplay(width=600, height=400, title='vertical position vs. time', xtitle='t', ytitle='y')
ygraph = gcurve(color=color.yellow)
## energy graphs
energyWindow = gdisplay(x=600, width=600, height=400, title='K, U, and K+U vs. time',xtitle='t',ytitle='E (cyan), K (yellow), U(red), in Joules')
Kgraph = gcurve(color=color.yellow) ## create a gcurve for kinetic energy
Ugraph = gcurve(color=color.red) ## create a gcurve for potential energy
KplusUgraph = gcurve(color=color.cyan) ## create a gcurve for total energy

## CALCULATIONS

while t < 3*T_analytic:
    L = ball.pos - ceiling.pos
    Lhat = L/mag(L)
    s = mag(L) - L0
    ## calculate forces on the ball
    #FSpring = 
    Fgrav = ball.m*vector(0,-g,0)
    Fnet = FSpring + Fgrav
    ## apply momentum principle
    ball.p = ball.p + Fnet*deltat
    ## update position
    ball.pos = ball.pos + ball.p/ball.m*deltat
    ## update axis of spring
    spring.axis = ball.pos - ceiling.pos
    ## update kinetic and potential energies
    #K = 
    #Ugrav =  
    #Uspring = 
    U = Ugrav + Uspring
    Etotal = K + U
    ## update graphs
    ygraph.plot(pos=(t,ball.pos.y))
    Kgraph.plot(pos=(t,K)) ## add a point to the kinetic energy graph
    Ugraph.plot(pos=(t,U)) ## add a point to the potential energy graph
    KplusUgraph.plot(pos=(t,Etotal)) ## add a point to the K+U graph
    ## update time   
    t = t + deltat
    rate(3000)