#Andrew Eshelman
#Section 1
#Assignment 4 Part B
#Arbitrary number of masses/springs

from visual import *
from random import *

#Constants
x0=1       #equilibrium length of all springs
d0=0       #initial distance from equilibrium of right-most atom
k=10       #spring constant of all springs
m=.5       #mass of all atoms
dt=.001    #timestep
tmax=7     #maximum time elapsed
t=0        #time counter
natoms=10  #number of atoms
time=0     #stores the time when left-most atom first moves, used to calculate speed of sound
init=0     #boolean for inital displacment from mouse drag
sound=0    #boolean: set to non-zero for speed-of-sound calc; must also set d0 to non-zero (.5)

#Lists to store atoms and springs
atoms=[]
springs=[]

#Generate atoms, starting at origin, 1 unit of distance between each succesive atom
for f in range(natoms):
    atoms.append(sphere(pos = vector(f,0), radius = .2, p = vector(0,0)))

#Optional (for fun): uncomment to connect last atom to first atom
#atoms.append(atoms[0])
#natoms+=1

#displaces right-most atom
atoms[natoms-1].pos+=vector(d0,0)

#Generates springs between atoms
for f in range(natoms-1):
    springs.append(helix(pos = atoms[f].pos, axis = atoms[f+1].pos-atoms[f].pos, coils=4))

#adjusts scene for optimum viewing
scene.range=(natoms/2*1.1,natoms/2*1.1,natoms/2*1.1)
scene.center=(natoms/2,0,0)

#Given script to input initial displacment with mouse;
#skips if d0 is defined or speed of sound is to be calculated
if not d0 and not sound:
    while init == 0:
        if scene.mouse.events:
            m1 = scene.mouse.getevent()
            for i in range(natoms):
                if m1.drag and m1.pick == atoms[i]:
                    j=i
                elif m1.drop:
                    atoms[j].pos=m1.pos
                    init=1

#Main Loop
while t<tmax:

#optional while condition for speed of sound; must comment preceding condition
#while not time:

    #update momentums of atoms
    for f in range(natoms-1):
        s=norm(atoms[f+1].pos-atoms[f].pos)*(mag(atoms[f+1].pos-atoms[f].pos)-x0)
        atoms[f].p+=dt*k*s
        atoms[f+1].p+=dt*k*-s

    #update atom positions
    for thing in atoms:
        thing.pos+=dt*thing.p/m

    #update springs
    for f in range(natoms-1):
        springs[f].pos=atoms[f].pos
        springs[f].axis=atoms[f+1].pos-atoms[f].pos

    t+=dt

    #store time that left-most atom first moves
    if mag(atoms[0].pos)>(d0/2) and not time:
        time=t

#print speed of sound if sound mode is on
if sound:
    print "Experimental speed of sound: ", (natoms-1+d0)/time
    print "Theoretical speed of sound: ", x0*((k/m)**.5)