from visual import *
from random import random
from time import clock

N = 15
Ntotal = N
scolor = (1,1,0)
springs = []
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=5, delta=1.0, R=None, sradius=None):
    if R == None:
        R = 0.2*delta
    if sradius == None:
        sradius = R/5.
    xmin = -(N-1.0)/2.
    sradius = R/5.
    natom = 0
    for nx in range(N):
        x = xmin+nx*delta
        hue = (nx)/(N+1.0)
        c = color.hsv_to_rgb((hue,1.0,1.0))
        atoms.append(sphere(pos=(x,0,0), radius=R, color=c))
        atoms[-1].p = vector()
        atoms[-1].near = range(2)
        atoms[-1].wallpos = range(2)
        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
            a.wallpos[0] = a.pos-vector(L,0,0)
        else:
            natom2 = nx-1
            a.near[0] = natom2
            makespring(natom1, natom2, sradius)
        if nx == N-1: # right
            a.near[1] = None
            a.wallpos[1] = a.pos+vector(L,0,0)
        else:
            natom2 = nx+1
            a.near[1] = natom2
            makespring(natom1, natom2, sradius)
    return atoms
m = 1.
k = 1.
L = 1.
R = 0.4*L
sradius = R/3.
vrange = 0.1*L*sqrt(k/m)
dt = 2.*pi*sqrt(m/k)/50.
#dt = 2.*pi*sqrt(m/k)/10.
atoms = crystal(N=N, delta=L, R=R, sradius=sradius)
scene.autoscale = 0

ptotal = vector()
for a in atoms:
    px = m*(-vrange/2+vrange*random())
    py = 0
    pz = 0
    a.p = vector(px,0,0)
    ptotal = ptotal+a.p

for a in atoms:
    a.p = a.p-ptotal/(N**2)

tt = clock()
Nsteps = 0
while 1:
    for a in atoms:
        F = vector()
        pos = a.pos
        near = a.near
        wallpos = a.wallpos
        for nn in range(2):
            natom = near[nn]
            if natom == None: # if this nearest neighbor is the wall
                r = wallpos[nn]-pos
            else:
                r = atoms[natom].pos-pos
            F = F+norm(r)*(mag(r)-L)
        a.p = a.p + k*F*dt

    for a in atoms:
        a.pos = a.pos+(a.p/m)*dt

    for s in springs:
        s.pos = atoms[s.natom1].pos
        s.axis = atoms[s.natom2].pos-s.pos
        s.radius = sradius*sqrt((L-2*R)/(mag(s.axis)-2*R))

    if Nsteps == 100:
        tt = clock()-tt
        print '%0.1f' % tt, 'sec for', Nsteps, 'steps with', N, 'on a side'
    Nsteps = Nsteps+1