from __future__ import division
from visual import *

print """
Click a level to change the number of quanta of energy
in one of the x, y, or z quantized oscillators.
"""

scene.width = 800
scene.height = 600
scene.range = 120
scene.background = color.white
scene.foreground = color.black
oa = vector(5,0,0)

yoffset = -50
wells=[]        ## create three harmonic oscillator potential wells
xx = arange(-10,10.5,0.5)
for offset in arange(-80,-19,30):
    wells.append(curve(x=offset+xx, y=yoffset+xx**2,z=0, radius=0.8))
wx = label(pos=(-80,-50,0), yoffset=-20, text="x", opacity=0)
wy = label(pos=(-50,-50,0), yoffset=-20, text="y", opacity=0)
wz = label(pos=(-20,-50,0), yoffset=-20, text="z", opacity=0)
wcolors = [(1,.8,0),(0,1,0),(.5,0,1)]
wells[0].color=wcolors[0]
wells[1].color=wcolors[1]
wells[2].color=wcolors[2]

levels=[]       ## three sets of energy levels (x,y,z)
kk = 0
for offset in arange(-80,-19,30):
    a = []
    for y in arange(10,100,20):
        xx = sqrt(y)
        a.append(cylinder(color=wcolors[kk], radius=0.8,
                          pos=(offset-xx,y+yoffset,0), axis = (2*xx,0,0)))
    levels.append(a)
    kk += 1

xoffset = 50
scene.center.x = 15
AO = vector(xoffset,0,0) # equil pos of atom
A = 4  # amplitude
omega = 2*pi/1
atom = sphere(pos=AO, radius=8, color=color.magenta)
slth = 45
springs = []
springends = [(-slth,0,0),(slth,0,0),(0,slth,0),(0,-slth,0),(0,0,slth),(0,0,-slth)]
scolors = [(1,.8,0),(1,.8,0),(0,1,0),(0,1,0),(.5,0,1),(.5,0,1)]
kk = 0
for loc in springends:
    v = vector(loc)
    springs.append(helix(pos=atom.pos+v, axis=(-v), radius=4,
                         color=scolors[kk], thickness=0.9))
    springs[-1].up = norm(vector(1,1,1))
    kk += 1
dl = 4
gray = 0.9
col = (gray,gray,gray)
box(pos=(xoffset-slth-dl/2,0,0), size=(dl,2*slth,2*slth), color=col)
box(pos=(xoffset+slth+dl/2,0,0), size=(dl,2*slth,2*slth), color=col)
box(pos=(xoffset,-slth-dl/2,0), size=(2*slth+2*dl,dl,2*slth), color=col)
box(pos=(xoffset,slth+dl/2,0), size=(2*slth+2*dl,dl,2*slth), color=col)
box(pos=(xoffset,0,-slth-dl/2), size=(2*slth+2*dl,2*slth+2*dl,dl), color=col)

picked = [0,0,0]
for lvl in levels:
    for ii in picked:
        lvl[ii].color =color.red
phi0 = pi/3.
phi1 = 2*pi/3
phi2 = pi
dt = 0.01
t = 0
while 1:
    rate(200)
    atom.pos = AO+vector((picked[0]+1)*A*sin(omega*t+phi0), (picked[1]+1)*A*sin(omega*t+phi1),
                         (picked[2]+1)*A*sin(omega*t+phi2))
    for sp in springs:
        sp.axis = atom.pos - sp.pos
    t = t+dt

    if scene.mouse.events:
        mm = scene.mouse.getevent()
        if mm.pick is not None:
            for well in levels:
                jj = levels.index(well)
                if mm.pick in well:
                    ii = well.index(mm.pick)
                    well[picked[jj]].color = wcolors[jj]
                    picked[jj] = ii
                    well[ii].color = color.red
##            print picked
##        else:
##            print 'None'