from visual import *
import sys
import math

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
def evalPotential(atoms):
    # Evaluate the total potential energy
    potential = 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
            dist12 = dist6 ** 2
            potential = potential + ((1.0 / dist12) - (2.0 / dist6))
    return potential
def nextAtoms(atoms, alpha, searchDirection):
    newAtoms = []
    for i in range(len(atoms)):
        newAtoms.append((atoms[i][0] + alpha * - searchDirection[i][0],
                         atoms[i][1] + alpha * - searchDirection[i][1],
                         atoms[i][2] + alpha * - searchDirection[i][2]))
    return newAtoms
def goldenSectionSearch(atoms, searchDirection, tol):
    t = (math.sqrt(5) - 1) / 2
    a = -t - 2
    b = t + 1
    x1 = a + (1 - t) * (b - a)
    f1 = evalPotential(nextAtoms(atoms, x1, searchDirection))
    x2 = a + (t * (b - a))
    f2 = evalPotential(nextAtoms(atoms, x2, searchDirection))
    while (abs(f1 - f2) > tol):
        if f1 > f2:
            a = x1
            x1 = x2
            f1 = f2
            x2 = a + (t * (b - a))
            f2 = evalPotential(nextAtoms(atoms, x2, searchDirection))
        else:
            b = x2
            x2 = x1
            f2 = f1
            x1 = a + (1 - t) * (b - a)
            f1 = evalPotential(nextAtoms(atoms, x1, searchDirection))
    return x1

try:
    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" % sys.argv[0]
    sys.exit(1)

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

tol = 0.01
atoms = createSimpleCube(rside, nside)
force = evalForce(atoms)
totalPotential = evalPotential(atoms)
forcetol = 0.01 * len(atoms)
totalForce = forcetol + 1
balls = []
for atom in atoms:
    balls.append(sphere(pos=atom, color=color.blue, radius=0.3))
while totalForce > forcetol:
    alpha = goldenSectionSearch(atoms, force, tol)
    atoms = nextAtoms(atoms, alpha, force)
    oldPotential = totalPotential
    force = evalForce(atoms)
    totalPotential = evalPotential(atoms)
    tol = abs((totalPotential - oldPotential) / totalPotential) * 10
    for i in range(len(atoms)):
        balls[i].pos = atoms[i]
    totalForce = 0
    for f in force:
        totalForce = totalForce + f[0] ** 2 + f[1] ** 2 + f[2] ** 2
print "The minimum potential is %e" % totalPotential