#Adapted with Significant Modifications, by T.L. Einstein, 11/03
#from Speed of Sound, Ellen Williams, 11/2/01

from visual.graph import *

N = 20
scolor = (1,1,0)
springs = []   #create springs as a list before putting data into it
atoms = []

def makespring(natom1, natom2, radius): # make spring from natom1 atom to natom2 atom
    if natom1 > natom2:
        springs.append(cylinder(pos=atoms[natom1].pos,
            axis=atoms[natom2].pos-atoms[natom1].pos,
            radius = radius, color=scolor))
        springs[-1].natom1 = natom1
        springs[-1].natom2 = natom2
    
def crystal(N=20, delta=1.0, R=None, sradius=None): #trimmed down, 
    #removing unneed methods from 3D version crystal.py
    if R == None:
        R = 0.2*delta
    if sradius == None:
        sradius = R/3.
    natom = 0
    for nx in range(N):
        hue =nx/(N+1.0)
        c = color.hsv_to_rgb((hue,1.0,1.0))
        atoms.append(sphere(pos=(nx*delta,0,0), radius=R, color=c))
        atoms[-1].natom = natom
        atoms[-1].indices = (nx,0,0)
        natom = natom+1
    for a in atoms:
        natom1 = a.natom
        nx, ny, nz = a.indices
 #       if nx == 0: # left
            # if this neighbor is the wall, save location:
 #           a.near[0] = None
        if nx != 0:
            natom2 = nx-1
            makespring(natom1, natom2, sradius)
 #       if nx == N-1: # right
 #       a.near[1] = None
        if nx != N-1:
            natom2 = nx+1
  #          a.near[1] = natom2
            makespring(natom1, natom2, sradius)
    return atoms


def forceHooke1D(r1, r2, k, req):
    #Hooke's law force, f in newtons
#Hooke's Law force f = -k(x-xeq), where xeq is the equilibrium separtion
#Force on atom1 at r1 due to atom 2 at r2 is f12
#f12 is positive (attractive force, to the right) if x2>x1 and x2-x1>xeq
#f12 is negative (repulsive force, to the left) if x2>x1 and x2-x1<req
#1-d case, f is a scalar, in newtons, with sign plus or minus one for direction
#input r1 and r2 and req are distances, units m
#k is the spring constant, units N/m
    dist=abs(r2-r1) #the separation between object 1 and object 2
    f=-k*(dist-req)  #the force on object 1 in N
    if r2>r1:
    	f=-f
    return f
      	   
    
#define constants  
#hypothetical atom of mass 30g/mole, spacing 0.5 nm
#use v ~ 5000 m/s ~ sqrt(k/m)*req to choose value of k
req=0.50e-9 #units m
k = 4.98  # units N/m
m = .03/6.02e23  # units kg per atom

# initial positions, speeds, and accelerations of n atoms
n = 20
xlist=[]
vlist=[]
alist=[]
for i in range(0,n):
    xlist.append(i*req)
    vlist.append(0.)
    alist.append(0.)
    	
#Choose dislacement of first atom to START SOUND PULSE
xlist[0]=-req/10.

#begin Euler calculation
dt = 1e-15  #time increment, units time,
                #initially chosen as 0.1*(req/5000ms-1)=1e-14
                #later reduced to improve energy conservation
timer=0.
icounter = 0
speedsound=[]
scene.autoscale=0
scene1=display(title='Motion of Chain',width=1200,height=200,\
center=(9.5*req,0,0),range=(11*req,11*req,11*req),background=(1,1,1))
#see VPython documentation for creating scenes with display; most entries
#should be familiar, but note the need to center the chain and the possibility 
#of extending the width

#Next we create an instance, called chain, of the object produced by crystal
#Chain is used rather than atoms to distinguish the local variables in crystal
#and in the main program.  Note that crystal is called only once.  Thereafter,
#its attributes, in particular positions, are modified.

chain=crystal(20, req)

#the while loop continues until the nth atom starts to move
while icounter < n and timer <(5000.*dt): #By increasing 5000 to 10000 or more, you can
    #watch the pulse propagate back through the chain!  But you must also remove the 
#icounter < n condition and put a way to bypass vlist[icounter] once icounter exceeds n-1  

#each atom (except for the ones at the ends) feels two forces
#I will use free end boundary conditions
#set up initial conditions for forces

#first atom
    alist[0] = (forceHooke1D(xlist[0],xlist[0+1], k, req))/m
    #atoms 2 through n-1
    for i in range(1,n-1):
            alist[i]=forceHooke1D(xlist[i],xlist[i-1], k, req)/m\
            +forceHooke1D(xlist[i],xlist[i+1], k, req)/m
    #last (nth) atom
    alist[n-1]=(forceHooke1D(xlist[n-1],xlist[n-2], k, req))/m
                     
    timer += dt
    rate(200)   #Can be speeded up, but low rate allows viewing of pulse reaching atoms.
    for i in range(0,n):
        xlist[i]=xlist[i]+vlist[i]*dt  #Can't use Eulerstep since 
    #                                                 not last in array
        vlist[i]=vlist[i]+alist[i]*dt
        chain[i].pos=(i*req +5.*(xlist[i]-i*req),0,0) #update positions, 
        #magnifying displacements 5-fold, as discussed in lab
        hue=float(i)/(n+1)  #reset original hue; HSV are not attributes of VPython objects;
        #could instead find RGB color of chain[i], convert to HSV, and save H
        saturatn=max(1.-30.*abs(i-xlist[i]/req),0.1) #set saturation; must be >=0, 
        #max displacement is req/10, but this exaggerates 3-fold for better visibility
        chain[i].color=color.hsv_to_rgb((hue,saturatn,1.)) #note ((  ,  ,  ))
#Note also that makespring is not called again; we just update the position and axis of 
#each spring, noting that the axes were created to point from i to i+1
    for i in range(0,n-1): 
        springs[i].pos = chain[i].pos
        springs[i].axis = chain[i+1].pos - chain[i].pos
#calculated sound speed "on the fly" by finding the time at which each atom first achieved max
#displacement compared with others.  would be more accurate to store all the atom positions
# vs. time and then find the absolute max for each list	
    if vlist[icounter]>0:
        speedsound.append(xlist[icounter]/timer)
        print('\nThe speed of sound at atom %d is %6.2e at timestep %d' ) \
        % (icounter, xlist[icounter]/timer, int(timer/dt))
        icounter += 1  #Still need to do this to end while loopicounter < n-1:
                            
print "\n The total time for the propagation is %6.2e \n" % timer
		
print '\nThe speed of sound (m/s) estimated from atoms 0 - %d is:' %(icounter-1) #since final increment
for i in range(0,icounter):
    print "%5d%12.3e" %(i, speedsound[i])