'''
Name: Sajjan Singh Mehta
Recitation Assignment #2
01/30/2007
PHYS 114-002
Recitation Class #2
'''


from random import random
from visual import *


# Constants
N = 8                         # Number of particles in the box
RUN_TIME = 2.5                # Amount of time to run before reversing velocity
dt = 0.0005                   # Timestep
sigma = 0.1                   # Sigma value for the Lennard-Jones potential
sigma_6 = sigma**6            # Sigma to the 6th power - to prevent repeated calculation
epsilon = 1                   # Epsilon value for the Lennard-Jones potential
GRAVITATIONAL_FORCE = True    # Tells whether or not to incorporate gravity with the Lennard-Jones potential
g = vector(0, -9.8, 0)        # Gravitational acceleration

# Runtime variables
particles = []                # List of particles in the box
t = 0                         # Current time during the run of the program
runStage = 0                  # Keeps track of the current run stage


# Updates particle accelerations based on Lennerd-Jones potential
def updateAcceleration(s):
    global particles, epsilon, sigma_6, GRAVITATIONAL_FORCE, g
    F = vector(0, 0, 0)
    
    if GRAVITATIONAL_FORCE:
        F = g * s.m
    
    for i in particles:
        if s != i:
            r = mag(s.pos - i.pos)
            F += norm(s.pos - i.pos) * (24 * epsilon) * ((2 * sigma_6**2 / r**13) - (sigma_6 / r**7))
    s.a = F / s.m


# Generate the particles
for i in xrange(N):
    s = sphere(m = 1, radius = 0.025, a = vector(0, 0, 0))
    s.color = (random(), random(), random())
    
    # Generates evenly spaced positions for the balls in a circular ring
    theta = 2.0 * pi / N
    s.pos = vector(0.25 * cos(theta * i), 0.25 * sin(theta * i), 0)
    
    # Generates random velocity with magnitude 1
    # The method used generates a random z value and uses that to calculate phi (since z = cos(phi)
    # in spherical coordinates).  sin(phi) then scales the radius of the circle on which
    # x = cos(theta) and y = sin(theta) lie on, resulting in a point (x, y, z) that exists on the
    # surface of a sphere with radius 1.  
    theta = (2 * pi) * random()
    z = -1 + (2 * random())
    phi = arccos(z)
    s.v = vector(cos(theta) * sin(phi), sin(theta) * sin(phi), z)

    particles.append(s)


# Create the walls showing the edge of the box
box(pos = (0.5, 0, 0), length = 0.01, height = 1, width = 1)
box(pos = (-0.5, 0, 0), length = 0.01, height = 1, width = 1)
box(pos = (0, 0.5, 0), length = 1, height = 0.01, width = 1)
box(pos = (0, -0.5, 0), length = 1, height = 0.01, width = 1)
box(pos = (0, 0, -0.5), length = 1, height = 1, width = 0.01)

# Create a label to display the orientation of velocity
orientation = label(pos = (-0.5, 0.5, 0), text = 'Velocity: Forward', color = color.blue,
                    opacity = 0, box = 0)
time = label(pos = (-0.5, 0.6, 0), text = 'Time: 0.0', color = color.blue,
             opacity = 0, box = 0)


# Set scene properties
scene.autoscale = 0
scene.range = (0.75, 0.75, 0.75)


while runStage < 2:
    rate(1 / dt)    
    
    for s in particles:
        # Numeric integration using velocity verlet - final result is the same as Euler's method
        # since there for constant accelerations.  Start by updating the position from the velocity
        # found in the last iteration - then half-update velocity using the current acceleration.
        s.pos += (s.v * dt) + (0.5 * s.a * dt**2)
        s.v += (0.5 * s.a * dt)
    
    # Update the acceleration for every particle
    for s in particles:
        updateAcceleration(s)
    
    for s in particles:
        # Half-update velocity again using the new acceleration
        s.v += (0.5 * s.a * dt)
        
        # Reverse the individual momentum components of the balls if they
        # hit any of the walls
        if abs(s.pos.x) + s.radius >= 0.5:
            s.v.x *= -1
        if abs(s.pos.y) + s.radius >= 0.5:
            s.v.y *= -1
        if abs(s.pos.z) + s.radius >= 0.5:
            s.v.z *= -1
    
    # Update run time and check if reveral of velocities is required
    t += dt
    if runStage == 0:
        time.text = 'Time: '+ str(t)
    else:
        time.text = 'Time: '+ str(round(RUN_TIME - t))
    
    if t >= RUN_TIME:
        if runStage == 0:
            runStage += 1
            t = 0.0
            orientation.text = 'Velocity: Reversed'
            
            for s in particles:
                s.v *= -1
        
        elif runStage == 1:
            runStage += 1