from numpy import *
from matplotlib.pyplot import *

# DEFINE FUNCTIONS
def compute_lj_force(x0, x1):
    # TODO
    # .
    # .
    # .

def compute_lj_potential(x0, x1):
    # TODO
    # .
    # .
    # .

def compute_forces(x):
    """Compute and return the forces acting onto the different
    planets, depending on the positions x."""
    global epsilon, sigma
    _, N = x.shape
    f = zeros_like(x)
    for i in range(1,N):
        for j in range(i):
            # distance vector
            fij = compute_lj_force(x[:,i], x[:,j])
            f[:,i] -= fij
            f[:,j] += fij
    return f

def compute_energies(x, v):
    _, N = x.shape
    E_pot = 0.0
    E_kin = 0.0
    # sum up potential energies
    for i in range(1,N):
        for j in range(i):
            E_pot += compute_lj_potential(x[:,i], x[:,j])
    # sum up kinetic energy
    for i in range(N):
        E_kin += 0.5 * dot(v[:,i],v[:,i])
    return E_pot + E_kin
    
def step_vv(x, v, f, dt):
    # update positions
    x += v*dt + 0.5*f * dt*dt
    # half update of the velocity
    v += 0.5*f * dt
        
    # compute new forces
    f = compute_forces(x)
    # we assume that m=1 for all particles

    # second half update of the velocity
    v += 0.5*f * dt

    return x, v, f

# constants
dt = 0.01
tmax = 20.0

# running variables
t = 0.0

# particle positions
x = zeros((3,5))
x[:,0] = [0.0, 0.0, 0.0]
x[:,1] = [5.0, 0.3, 0.0]
x[:,2] = [8.0, 1.8, 0.0]
x[:,3] = [11.0, 0.0, -1.0]
x[:,4] = [12.0, 9.0, 0.0]

# particle velocities
v = zeros((3,5))
v[:,0] = [2.0, 0.0, 0.0]
v[:,1] = [0.0, 0.0, 0.0]
v[:,2] = [0.0, 0.0, 0.0]
v[:,3] = [0.0, 0.0, 0.0]
v[:,4] = [0.0, 0.0, 0.0]

f = compute_forces(x)

# variables to cumulate data
traj = []

# number of particles
N = x.shape[1]

# open the trjectory file
vtffile = open('ljbillards.vtf', 'w')
# write the structure of the system into the file: 
# N particles ("atoms") with a radius of 0.5
vtffile.write('atom 0:%s radius 0.5\n' % (N-1))

# main loop
while t < tmax:
    x, v, f = step_vv(x, v, f, dt)
    t += dt

    traj.append(x.copy())

    # write out that a new timestep starts
    vtffile.write('timestep\n')
    # write out the coordinates of the particles
    for i in range(N):
        vtffile.write("%s %s %s\n" % (x[0,i], x[1,i], x[2,i]))

vtffile.close()

traj = array(traj)

# plot the trajectory
for i in range(N):
    plot(traj[:,0,i], traj[:,1,i], label='%s'%i)
legend()
show()