from visual import *
import sys

def createSimpleCube(rside, nside):
    # Return a cube with length rside and nside atoms on each side
    cube = []
    spacing = [0,0,0]
    for n in range(3):
        if nside[n] == 1:
            spacing[n] = 0
        else:
            spacing[n] = float(rside[n]) / (nside[n] - 1)
    for x in range(nside[0]):
        for y in range(nside[1]):
            for z in range(nside[2]):
                cube.append((spacing[0] * x, spacing[1] * y, spacing[2] * z))
    return cube
def evalForce(atoms):
    # Evaluate the force on each atom
    force = []
    for i in range(len(atoms)):
        force.append((0.0, 0.0, 0.0))
    for i in range(len(atoms)):
        for j in range(i):
            dist2 = ((atoms[i][0] - atoms[j][0]) ** 2 +
                     (atoms[i][1] - atoms[j][1]) ** 2 +
                     (atoms[i][2] - atoms[j][2]) ** 2)
            dist6 = dist2 ** 3
            dist8 = dist2 * dist6
            dist14 = dist8 * dist6
            a = 12.0 / dist14 - 12.0 / dist8
            b0 = (a * (atoms[i][0] - atoms[j][0]))
            b1 = (a * (atoms[i][1] - atoms[j][1]))
            b2 = (a * (atoms[i][2] - atoms[j][2]))
            force[i] = (force[i][0] + b0, force[i][1] + b1, force[i][2] + b2)
            force[j] = (force[j][0] - b0, force[j][1] - b1, force[j][2] - b2)
    return force

try:
    tstep = float(sys.argv[3])
    nstep = int(sys.argv[4])
    r = float(sys.argv[1])
    rside = (r, r, r)
    n = int(sys.argv[2])
    nside = (n, n, n)
except (ValueError, IndexError):
    print "Usage: python %s rside nside tstep nstep" % sys.argv[0]
    sys.exit(1)

scene.autoscale = 0
scene.title = "Molecular dynamics simulator"

atoms = createSimpleCube(rside, nside)
veloc = []
for atom in atoms:
    veloc.append((0,0,0))
newForce = oldForce = evalForce(atoms)
balls = []
for atom in atoms:
    balls.append(sphere(pos=atom, color=color.blue, radius=0.3))

for istep in range(nstep):
    for i in range(len(atoms)):
        atoms[i] = (atoms[i][0] + tstep * (veloc[i][0] +
                        (tstep * newForce[i][0] / 2.0)),
                    atoms[i][1] + tstep * (veloc[i][1] +
                        (tstep * newForce[i][1] / 2.0)),
                    atoms[i][2] + tstep * (veloc[i][2] +
                        (tstep * newForce[i][2] / 2.0)))
    newForce = evalForce(atoms)
    for i in range(len(veloc)):
        veloc[i] = (veloc[i][0] + tstep *
                        (newForce[i][0] + oldForce[i][0]) / 2.0,
                    veloc[i][1] + tstep * 
                        (newForce[i][1] + oldForce[i][1]) / 2.0,
                    veloc[i][2] + tstep * 
                        (newForce[i][2] + oldForce[i][2]) / 2.0)
    oldForce = newForce
    for i in range(len(atoms)):
        balls[i].pos = atoms[i]