#Andrew Eshelman
#Section 1
#Assignment 4 Part A
#Simple 2-mass, 1-spring simulation... with comments

#import visual library to display things
from visual import *

#import random library for cool effects
from random import *

#constants

x0=1     #equilibrium length of spring
d0=.1    #initial displacement of atom2 from equilibrium
k=10     #spring constant of spring
m=1      #mass of atoms
dt=.001  #timestep
tmax=5   #maximum time elapsed
t=0      #time counter

#generate atom 1
atom1 = sphere(pos = vector(0,0), radius = .2, p = vector(0,0))

#generate atom 2
atom2 = sphere(pos = vector(1+d0,0), radius = .2, p = vector(0,0))

#generate spring connecting atoms
spring = helix(pos = atom1.pos, axis = atom2.pos-atom1.pos, coils=4)

#stops the scene from auto-scaling
scene.autoscale = 0

#main loop: loops until t reaches tmax
while t<tmax:

    #update spring position
    spring.pos=atom1.pos

    #update spring axis
    spring.axis=atom2.pos-atom1.pos

    #generate vector between atoms
    s=norm(atom2.pos-atom1.pos)*(mag(atom2.pos-atom1.pos)-x0) 

    #update momentum of atom 2
    atom2.p+=dt*-k*s

    #update momentum of atom 1
    atom1.p+=dt*-k*-s

    #update position of atom 2
    atom2.pos+=dt*atom2.p/m

    #update position of atom 1
    atom1.pos+=dt*atom1.p/m

    #uncomment for psychadelic party-town
    #atom1.color=(random(),random(),random())
    #atom2.color=(random(),random(),random())
    #spring.color=(random(),random(),random())

    #increment time counter
    t+=dt

    #rate: increased so it isn't as slow
    rate(300)