from visual import *

atom=[]
natom=10
ratom=.05
L=1.
g=9.8

scene.autoscale=0
scene.center=(0.5,0.5,0.5)
scene.range=(2,2,2)

ypos=[]
for i in range(natom):
	atom.append(sphere(pos=vector(0.05+i*0.1,0.05,0.5),radius=ratom))
	atom[i].v=vector(0,0,0)	
	ypos.append(atom[i].pos.y)

leftside=box(pos=(0,L,L/2),length=0.05,height=2*L,width=L,color=color.green)
rightside=box(pos=(L,L,L/2),length=0.05,height=2*L,width=L,color=color.green)
backside=box(pos=(L/2,L,0),length=L,height=2*L,width=0.05,color=color.green)
bottom=box(pos=(L/2,0,L/2),length=L,width=L,height=0.05,color=color.green)

t=0
dt=0.01

nlist=[]
ylist=[]
dy=0.2


while (t < 1000): 
	rate(100)

	U=0
	K=0

	for i in range(natom):
		# Do any of them collide
		if i > 0: 		
			for j in range (i-1): 
				dist=mag(atom[i].pos-atom[j].pos)
				if (dist < 2*ratom):
					rhat=(atom[i].pos-atom[j].pos)/dist
					v0=(atom[i].v+atom[j].v)/2.
					vi=-dot(atom[i].v-v0,rhat)
					vj=dot(atom[j].v-v0,rhat)
					if (vi > 0 and vj > 0):
						atom[i].v=atom[i].v+2*vj*rhat
						atom[j].v=atom[j].v-2*vi*rhat


		if (atom[i].pos.y < ratom and atom[i].v.y < 0):
			 atom[i].v.y=-atom[i].v.y
		if (atom[i].pos.x < ratom and atom[i].v.x < 0):
			atom[i].v.x=-atom[i].v.x
		if (atom[i].pos.x > L-ratom and atom[i].v.x > 0):
			atom[i].v.x=-atom[i].v.x
		if (atom[i].pos.z < ratom and atom[i].v.z < 0):
			atom[i].v.z=-atom[i].v.z
		if (atom[i].pos.z > L-ratom and atom[i].v.z > 0):
			atom[i].v.z=-atom[i].v.z
		

		atom[i].pos=atom[i].pos+atom[i].v*dt
		atom[i].pos.y=atom[i].pos.y-0.5*g*dt*dt
		atom[i].v.y=atom[i].v.y-g*dt	
		t=t+dt