from numpy import zeros,empty
from pylab import plot,show

from visual import sphere, display, rate
from numpy import linspace, empty,sin

# Constants
N = 20
C = 1.0
m = 1.0
k = 6.0
omega = 2.0
alpha = 2*k-m*omega*omega

# Set up the initial values of the arrays
A = zeros([N,N],float)
for i in range(N-1):
    A[i,i] = alpha
    A[i,i+1] = -k
    A[i+1,i] = -k
A[0,0] = alpha - k
A[N-1,N-1] = alpha - k

v = zeros(N,float)
v[0] = C

# Perform the Gaussian elimination
for i in range(N-1):

    # Divide row i by its diagonal element
    A[i,i+1] /= A[i,i]
    v[i] /= A[i,i]

    # Now subtract it from the next row down
    A[i+1,i+1] -= A[i+1,i]*A[i,i+1]
    v[i+1] -= A[i+1,i]*v[i]

# Divide the last element of v by the last diagonal element
v[N-1] /= A[N-1,N-1]

# Backsubstitution
x = empty(N,float)
x[N-1] = v[N-1]
for i in range(N-2,-1,-1):
    x[i] = v[i] - A[i,i+1]*x[i+1]

# Make a plot using both dots and lines
#==============================================================================
# plot(x)
# plot(x,"ko")
# show()
#==============================================================================

eq_pos = linspace(0,2*N,N)
s = empty(N,sphere)
for i in range(N):
	s[i] = sphere()
	s[i].radius = 0.3
	
def plot_positions(positions):
	
	
	for i,pos in enumerate(positions):
		#s[i] = sphere()

		s[i].pos = [pos,0,0]

plot_positions(eq_pos)
d = display()
d.autoscale=False

t = 0
while True:
	
	rate(30)
	pos = eq_pos + x*sin(omega*t)
	plot_positions(pos)
	t +=1
