from visual import *
from visual.graph import *
scene.height = 800
scene.width = 200
#oscillations of a mass on a spring with mass
#   Jonathan Marr

start=0.0 #start with mass at origin, which is resting length of spring,
        #but won't be equil. pos. because of gravity

#create block to be the oscillating mass
block=box(pos=vector(0,start,0),size=(0.1,0.1,0.1),color=color.green)
block.m=0.01
block.p = vector(0,0,0)#start block at rest

#create spring, of natural length sprL, stiffness ks, and mass Mspr.
#One end wilk be fixed at top, a distance sprL from origin.
#and the other end will follow the block (as if it was attached)
ks=8.0
sprL=0.2
Mspr=0.03

#break spring up into Nspr pieces and give each piece a mass=Mspr/Nspr and stiffness Nspr*ks.
Nspr=12
i=1
spring = {}
Fspr = {}
Fnetspr = {}
print "Hello"
spring[1]=helix(pos=(0,sprL/Nspr,0),axis=(0,(block.y)-sprL/Nspr,0),radius=0.01,color=color.white,thickness=0.02,shineness=0.9)
spring[1].p=vector(0,0,0)
for i in range (2,Nspr+1):
    spring[i]=helix(pos=(0,i*sprL/Nspr,0),axis=(0,sprL/Nspr,0),radius=0.01,color=color.white,thickness=0.02,shineness=0.9)
    spring[i].p = vector(0,0,0)
    

#set up graphics for plotting of block's position and momentum vs. time
gdisplay(x=100,y=500,xtitle='time (sec)', ytitle='Y (cyan), Py (red)')
xcurve = gcurve(color=color.cyan)
pcurve = gcurve(color=color.red)

#"while" loop to create motion and plot data
deltat=0.001
t=0
Npd = 0 #counting index for number of periods 
tend = 0 # will calculate time duration for Npd number of pds
direction = -1 #we'll use change of direction at top to count periods, and will start moving downward
Fgblock=vector(0,-block.m*9.8,0) #gravity on block
Fgspr=vector(0,-(Mspr/Nspr)*9.8,0) #gravity on each piece of spring
while Npd<10:  # cound 10 periods
    rate(1000)
    xcurve.plot(pos=(t,block.y)) #plotting y-position data as it occurs
    pcurve.plot(pos=(t,block.p.y))  #likewise for the y-momentum data.

    Fspr[1]=-(Nspr*ks)*((spring[1].y-block.y)-sprL/Nspr)*vector(0,-1,0)  #spring force of bottom piece (resting length when block.y=0
    j=2
    while j<Nspr:
        Fspr[j]=-(Nspr*ks)*((spring[j].y-spring[j-1].y)-sprL/Nspr)*vector(0,-1,0)
            #spring force for each piece (ks of each piece in series = Nspr time ks of entire spring)
        j=j+1
    Fspr[Nspr]=-(Nspr*ks)*((sprL-spring[Nspr-1].y)-sprL/Nspr)*vector(0,-1,0)  #top piece is fixed at top
    Fnetblock=Fgblock+Fspr[1]
    k=1
    while k<Nspr:
        Fnetspr[k]=Fgspr-Fspr[k]+Fspr[k+1]  #net force acting on the top of each piece of spring, top of topmost doesn't move
        k=k+1

    block.p=block.p+Fnetblock*deltat # momentum update of block
    block.pos=block.pos+(block.p/block.m)*deltat #position update

    l=1
    while l<Nspr:
        spring[l].p=spring[l].p+Fnetspr[l]*deltat # momentum update of each piece of spring
        spring[l].pos=spring[l].pos+(spring[l].p/(Mspr/Nspr))*deltat
        l=l+1

    spring[1].axis=(0,block.y-spring[1].y,0)  #make sure end of spring follows cart in its motion
    m=2
    while m<Nspr+1:
        spring[m].axis=(0,spring[m-1].y-spring[m].y,0)  #make sure end of each spring follows top of piece below it
        m=m+1
              
    if direction<0 and block.p.y>0:  #change of direction at bottom
        direction=1
    if direction>0 and block.p.y<0:  #change of direction at top
        Npd=Npd+1 #counting number of periods by noting when change direction at one end
        tend=t #Note no. of seconds when last cycle ends.
        direction=-1
    t=t+deltat
print "number of periods = ",Npd, "over a duration of", tend, "seconds"
Period=tend/Npd
print "Period of one cycle = ", Period, "seconds"
#print "Delta t =", t, "Final momentum =",cart.p