#Andrew Eshelman
#Section 1
#Assignment 4
#Bound Pyramid

from visual import *
from random import *

#constants
seed(672)          #seed for random color
x0=1               #equilibrium length of all springs
d0=vector(0,0,0)   #initial displacement of atom at tip of pyramid
k=10.              #spring constant for all springs
m=1                #mass of all atoms
dt=.01             #timestep
tmax=50            #maximum elapsed time
t=0                #time counter
base=3             #length of pyramid base
dis=(.5*x0)**.5    #distance between pyramid levels
                   #set such that inter-level springs are in equilibrium
airc=1             #momentum-reducing air resistence coefficent, Sug: .995
init=0             #boolean for initial displacment from mouse

class spring:

    #stores atoms being connected and generates visual
    def __init__(self,a,b):
        self.a=a
        self.b=b
        self.vis = helix(pos=a.pos, axis=b.pos-a.pos, coils=4, radius=.05,
                         color=(random(),random(),random()), thickness=.02)
    #updates momentums of a and b, and visual component
    def up(self): 
        s=norm(self.b.pos-self.a.pos)*(mag(self.b.pos-self.a.pos)-x0)
        self.a.p+=dt*k*s
        self.b.p+=dt*-k*s
        self.vis.pos=self.a.pos
        self.vis.axis=self.b.pos-self.a.pos

class atom:

    #sets initial position and momentum; generates sphere for visual
    def __init__(self, pos, radius, p):
        self.pos=pos
        self.p=p
        self.vis=sphere(pos = self.pos, radius = .2)

    #updates position and visual and moentum if friction exists
    def up(self):
        self.p*=airc
        self.pos+=dt*self.p/m
        self.vis.pos=self.pos

#imaginary atoms; they don't move
class iatom:
              
    def __init__(self,pos):
        self.pos=pos
        #dummy p variable to avoid errors in spring.up()
        self.p=vector(23,23)
    
atoms=[]
iatoms=[]
springs=[]

#adjusts scene for optimum viewing
scene.range=(base,base,base)
scene.center=((base-1)/2.,(base-1)/2.,(base-1)/2.)

#generates two binding walls
box(pos=(-x0,(base-1)/2.,(base-1)/2.), length=.1, height=base, width=base)
box(pos=((base-1)*dis+x0,(base-1)/2.,(base-1)/2.), length=.1, height=base, width=base)


# 3-D list to store atoms
for x in range(base):
    atoms.append([])
    for y in range(base-x):
        atoms[x].append([])
        for z in range(base-x):
           atoms[x][y].append(atom(vector(x*dis,y+x/2.,z+x/2.),.2,vector(0,0)))

#Initial displacment
atoms[base-1][0][0].pos+=d0
atoms[base-1][0][0].up()

#2-D list of iatoms at base; used to bound pyramid
iatoms.append([])
for y in range(base):
    iatoms[0].append([])
    for z in range(base):
        iatoms[0][y].append(iatom(vector(-x0,y,z)))

#iatom at tip
iatoms.append(iatom(vector((base-1)*dis+x0, (base-1)/2.,( base-1)/2.)))

#next 2: intra-level connections
for x in range(base-1):
    for y in range(base-x):
        for z in range(base-x-1):
            springs.append(spring(atoms[x][y][z], atoms[x][y][z+1]))

for x in range(base-1):
    for y in range(base-x-1):
        for z in range(base-x):
            springs.append(spring(atoms[x][y][z], atoms[x][y+1][z]))

#inter-level connections
for x in range(base-1):
    for y in range(base-x-1):
        for z in range(base-x-1):
            springs.append(spring(atoms[x][y][z], atoms[x+1][y][z]))
            springs.append(spring(atoms[x][y+1][z], atoms[x+1][y][z]))
            springs.append(spring(atoms[x][y][z+1], atoms[x+1][y][z]))
            springs.append(spring(atoms[x][y+1][z+1], atoms[x+1][y][z]))

#iatom connections
for y in range(base):
    for z in range(base):
        springs.append(spring(iatoms[0][y][z], atoms[0][y][z]))

springs.append(spring(iatoms[1], atoms[base-1][0][0]))

#displace atom with mouse-drag
while not init:
    if scene.mouse.events:
        m1 = scene.mouse.getevent()
        for thing1 in atoms:
            for thing2 in thing1:
                for thing3 in thing2:
                    if m1.drag and m1.pick == thing3.vis:
                        j=thing3
                    elif m1.drop:
                        j.pos=m1.pos
                        init=1

#Main loop        
while t<tmax:
    
    #updates springs
    for thing in springs:
        thing.up()

    #updates atoms
    for thing1 in atoms:
        for thing2 in thing1:
            for thing3 in thing2:
                thing3.up()
    t+=dt

print 'done'