from visual import *
from visual.controls import *
def change(): # Called by controls when button is clicked
    if t1.value:
        s1.color=color.green
    else:
        s1.color=color.yellow
c = controls() # Create controls window

# Create a button in the controls window:
t1 = toggle(pos=(40,-30), width=10, height=10, text0='Red', text1='Cyan', action=lambda: change())


scene0=display(title='Master Controls',width=450,height=350)
#define Accelerator Controls
scene0.select()
scene0.autoscale=1
injectionoff=box(pos=(-50,15,0),length=20,width=1,height=20,color=color.red)
injectionon=box(pos=(-50,-15,0),length=20,width=1,height=20,color=color.green)
rampoff=box(pos=(-20,15,0),length=20,width=1,height=20,color=color.red)
rampon=box(pos=(-20,-15,0),length=20,width=1,height=20,color=color.green)
ok4off=box(pos=(10,15,0),length=20,width=1,height=20,color=color.red)
ok4on=box(pos=(10,-15,0),length=20,width=1,height=20,color=color.green)
inonlbl=label(pos=injectionon.pos,text='ON',zoffset=5,space=0,height=30,opacity=0,color=color.white,linecolor=color.green)
inofflbl=label(pos=injectionoff.pos,text='OFF',zoffset=5,space=0,height=25,opacity=0,color=color.white,linecolor=color.red)
rmonlbl=label(pos=rampon.pos,text='ON',zoffset=5,space=0,height=30,opacity=0,color=color.white,linecolor=color.green)
rmofflbl=label(pos=rampoff.pos,text='OFF',zoffset=5,space=0,height=25,opacity=0,color=color.white,linecolor=color.red)
okonlbl=label(pos=ok4on.pos,text='ON',zoffset=5,space=0,height=30,opacity=0,color=color.white,linecolor=color.green)
okofflbl=label(pos=ok4off.pos,text='OFF',zoffset=5,space=0,height=25,opacity=0,color=color.white,linecolor=color.red)
linactxt=label(pos=injectionoff.pos,text='LINAC',yoffset=60,height=20,color=color.red,linecolor=color.red)
ramptxt=label(pos=rampoff.pos,text='RAMP',yoffset=60,height=20,color=color.red,linecolor=color.red)
wiggtxt=label(pos=ok4off.pos,text='OK-4',yoffset=60,height=20,color=color.red,linecolor=color.red)
gammabutt=box(pos=(50,15,0),length=20,width=1,height=20,color=color.blue)
gammalbl=label(pos=gammabutt.pos,text='Gamma',zoffset=5,space=0,height=15,opacity=0,color=color.white,linecolor=color.blue)
exitbutt=box(pos=(50,-15,0),length=20,width=1,height=20,color=color.red)
exitlbl=label(pos=exitbutt.pos,text='Exit',zoffset=5,space=0,height=30,opacity=0,color=color.white,linecolor=color.red)

scene1=display(title='Duke Free Electron Laser Lab',width=1000,height=1000)
scene1.select()
#scene1.autoscale=0
# Draw the Duke Storage Ring
# North Straight Section
ring=curve(pos=[(65,0,-20),(-65,0,-20)],radius=2,color=color.blue)
#South Straight Section
ring.append=curve(pos=[(120,0,20),(-100,0,20)],radius=2,color=(0.6,0.8,0.3))
# West Arc
te=arange(3.14/2.0,3.0*3.2/2.0,0.1)
ring.append=curve(x=-65+20*cos(te),y=0,z=20*sin(te),radius=2,color=color.blue)
#East Arc
tw=arange(-3.14/2.0,3.14/2.0,0.1)
ring.append=curve(x=65+20*cos(tw),y=0,z=20*sin(tw),radius=2,color=color.blue)
# Photon
photon=frame()
tp=arange(0,20,0.1)
efield1=curve(frame=photon,x=tp,y=0,z=1.0*sin(tp),radius=0.5,color=color.blue)
dir1=arrow(frame=photon,x=20,y=0,z=0,axis=(5,0,0),shaftwidth=1,color=color.red)
efield2=curve(frame=photon,x=tp,y=0,z=2+1.0*sin(tp),radius=0.5,color=color.blue)
dir2=arrow(frame=photon,x=20,y=0,z=2,axis=(5,0,0),shaftwidth=1,color=color.red)
efield3=curve(frame=photon,x=tp,z=-2+1.0*sin(tp),y=0,radius=0.5,color=color.blue)
dir3=arrow(frame=photon,x=20,y=0,z=-2,axis=(5,0,0),shaftwidth=1,color=color.red)
photon.pos=vector(-30,5,20)
for obj in photon.objects:
    obj.visible=0
# Gamma Ray
gray=frame()
tr=arange(0,20,0.1)
gfield=curve(frame=gray,x=tp,z=1.0*sin(2.0*tr),radius=0.5,color=color.red)
gdir=arrow(frame=gray,x=20,y=0,z=0,axis=(5,0,0),shaftwidth=1,color=color.green)
gray.pos=vector(0,5,20)
for obj in gray.objects:
    obj.visible=0
#Gamma-Ray Beam Pipe
#gammaline=curve(pos=[(65,0,20),(120,0,20)],radius=2)
#gammaline.color=color.yellow
#LINAC
linac=curve(pos=[(65,0,-20),(120,0,-20)],radius=2)
linac.color=color.cyan
linac.stat=0
linaccav1=cylinder(pos=(120,0,-20),length=5,color=color.green,radius=4)
linaccav2=cylinder(pos=(110,0,-20),length=5,color=color.green,radius=4)
linaccav3=cylinder(pos=(100,0,-20),length=5,color=color.green,radius=4)
# RF Cavity
RFcav=cylinder(pos=(2.5,0,-20),length=2,color=color.green,radius=6)
RFcav=cylinder(pos=(-2.5,0,-20),length=2,color=color.green,radius=6)
RFcav=cylinder(pos=(-7,0,-20),length=14,color=color.red,radius=4)
RFcav=cylinder(pos=(7,0,-20),length=5,color=color.red,radius=6)
NPole=0
Wigglerstartx=-30
Wigglerstopx=30
# Wigglers
while NPole<3:
    stepN=-30+NPole*10
    stepS=-25+NPole*10
    northpole=box(pos=(stepN,-10,20),length=3,height=5,width=8,color=color.yellow)
#    northlabel=label(pos=northpole.pos,text='N',zoffset=5,space=0,height=10,opacity=0,color=color.red,linecolor=northpole.color)
    southpole=box(pos=(stepS,-10,20),length=3,height=5,width=8,color=color.green)
#    southlabel=label(pos=southpole.pos,text='S',zoffset=5,space=0,height=10,opacity=0,color=color.blue,linecolor=southpole.color)
    northpole=box(pos=(stepN+35,-10,20),length=3,height=5,width=8,color=color.yellow)
#    northlabel=label(pos=northpole.pos,text='N',zoffset=5,space=0,height=10,opacity=0,color=color.red,linecolor=northpole.color)
    southpole=box(pos=(stepS+35,-10,20),length=3,height=5,width=8,color=color.green)
#    southlabel=label(pos=southpole.pos,text='S',zoffset=5,space=0,height=10,opacity=0,color=color.blue,linecolor=southpole.color)
    NPole=NPole+1
plate=box(pos=(-18,-15,20),length=30,height=1,width=8,color=color.blue)
plate=box(pos=(18,-15,20),length=30,height=1,width=8,color=color.blue)
# Buncher
buncher=box(pos=(0,-10,20),length=5,height=5,width=8,color=color.red)
# Mirrors
mirror1=cylinder(pos=(-85,0,20),length=0.5,radius=6,color=color.red)
mirror2=cylinder(pos=(85,0,20),length=0.5,radius=6,color=color.red)
# Draw and add dynamics to the electron bunch
# Some End Points
nss=curve()
sss=curve()
warc=curve()
earc=curve()
nss.start=vector(68,0,-20)
nss.end=vector(-65,0,-20)
sss.start=vector(-65,0,20)
sss.end=vector(65,0,20)
warc.start=vector(-65,0,-20)
warc.end=vector(-65,0,20)
earc.start=vector(65,0,20)
earc.end=vector(65,0,-20)
#Lasing bunch
b1=sphere()
b1.radius=3
b1.start=vector(120,0,-20)
b1.stop=vector(-65,0,20)
b1.color=color.red
b1.pos=b1.start
b1.velocity=vector(-5,0,0)
#scattering bunch
s1=sphere()
s1.radius=3
s1.start=vector(120,0,-20)
s1.stop=vector(-65,0,20)
s1.color=color.blue
s1.pos=vector(0,0,-20)
dt=0.05
dx=0.05
ds=0.1
dp=0.0525
dz=0.05
warc_t_i=3.14/2.0
warc_t_f=3.14/2.0
earc_t_i=3.0*3.14/2.0
earc_t_f=3.0*3.14/2.0
warcs_t_i=3.14/2.0
warcs_t_f=3.14/2.0
earcs_t_i=3.0*3.14/2.0
earcs_t_f=3.0*3.14/2.0

dtheta=0.013
# Set Default Operation
Ramp=0
OK4=0
LINAC=0
twob=0
gamma=0
arcrate=500
srate=500
pushphoton=0
refm1=0
refm2=0
exit='n'
WigglerField='w'
while exit!='y':
    c.interact()
    if scene0.mouse.clicked:
        butt=scene0.mouse.getclick()
        whichbutt=scene0.mouse.pick
        if whichbutt==injectionon:
            LINAC=1
            print 'LINAC YES'
        elif whichbutt==injectionoff:
            LINAC=-1        
            print 'LINAC NO'

        elif whichbutt==rampon:
            Ramp=1
            arcrate=2000
            srate=2000
            print 'RAMP YES'
        elif whichbutt==rampoff:
            Ramp=-1
            arcrate=500
            srate=500
            print 'RAMP NO'

        elif whichbutt==ok4on:
            OK4=1
            print 'OK4 YES'
        elif whichbutt==ok4off:
            OK4=-1
            pushphoton=0
            for obj in photon.objects:
                obj.visible=0
            photon.pos=vector(0,5,20)
            photon.axis=vector(20,0,0)
            for obg in gray.objects:
                obg.visible=0
            refm1=0
            refm2=0
            print 'OK4 NO'

        elif whichbutt==gammabutt:
            twob=1
            triggergamma=0
            print 'Gamma Mode'

        elif whichbutt==exitbutt:
            exit='y'
            print 'EXIT'

# check cases
   
    rate(srate)
# Accelerate in LINAC
    if LINAC==1:
        if b1.pos.x>=nss.start.x:
            b1.pos.x=b1.pos.x-dx
# Accelerate in North Straight Section
    if b1.pos.x<=nss.start.x and b1.pos.x>=nss.end.x:
        b1.pos.x=b1.pos.x-dx
    if gamma==1 and s1.pos.x<=nss.start.x and s1.pos.x>=nss.end.x:
        s1.pos.x=s1.pos.x-dx
# Go through West Arc
    if b1.pos.x<=warc.start.x:
        rate(arcrate)
        warc_t_f=warc_t_f+dtheta
        if warc_t_f <=4.72:
            darcx=20.0*(cos(warc_t_i)-cos(warc_t_f))
            darcz=20.0*(sin(warc_t_i)-sin(warc_t_f))
            warc_t_i=warc_t_f
            b1.pos.x=b1.pos.x-darcx
            b1.pos.z=b1.pos.z+darcz
            LINAC=-1
    if gamma==1 and s1.pos.x<=warc.start.x:
        rate(arcrate)
        warcs_t_f=warcs_t_f+dtheta
        if warcs_t_f <=4.72:
            darcxs=20.0*(cos(warcs_t_i)-cos(warcs_t_f))
            darczs=20.0*(sin(warcs_t_i)-sin(warcs_t_f))
            warcs_t_i=warcs_t_f
            s1.pos.x=s1.pos.x-darcxs
            s1.pos.z=s1.pos.z+darczs
#GO Through SSS
    if b1.pos.x>=sss.start.x and b1.pos.z>0 and b1.pos.x<=sss.end.x:
        warc_t_f=3.14/2.0
        warc_t_i=3.14/2.0
        earc_t_i=3.0*3.14/2.0
        earc_t_f=3.0*3.14/2.0
        if abs(b1.pos.x-Wigglerstartx)<0.5 and OK4==1 and pushphoton!=1:
            pushphoton=1
            efield1.visible=1
            dir1.visible=1
        if abs(b1.pos.x-5.0)<0.5 and OK4==1 and pushphoton==1:
            efield2.visible=1
            dir2.visible=1
            efield3.visible=1
            dir3.visible=1    
        
            
# if gamma mode of operation then start the scattering bucnh
        if gamma==0 and twob==1 and abs(b1.pos.x-0)<0.5:
            gamma=1
            s1.color=color.yellow
# IF GOING THROUGH THE WIGGLER
        if b1.pos.x>=Wigglerstartx and b1.pos.x<=Wigglerstopx and OK4==1:
            b1.pos.x=b1.pos.x+ds
            b1.pos.z=b1.pos.z+0.1*cos(b1.pos.x-Wigglerstartx)
        else:
            b1.pos.x=b1.pos.x+ds
            b1.pos.z=20

    if gamma==1 and s1.pos.x>=sss.start.x and s1.pos.z>0 and s1.pos.x<=sss.end.x:
        warcs_t_f=3.14/2.0
        warcs_t_i=3.14/2.0
        earcs_t_i=3.0*3.14/2.0
        earcs_t_f=3.0*3.14/2.0
        if gamma==1 and s1.pos.x>=Wigglerstartx and s1.pos.x<=Wigglerstopx and OK4==1:
            s1.pos.x=s1.pos.x+ds
            s1.pos.z=s1.pos.z+0.1*cos(s1.pos.x-Wigglerstartx)
        else:
            s1.pos.x=s1.pos.x+ds
            s1.pos.z=20
        if abs(s1.pos.x-photon.pos.x)<0.1 and OK4==1:
            gfield.visible=1
            gdir.visible=1
            triggergamma=1
    if gamma==1 and triggergamma==1 and OK4==1:
        if gray.pos.x<120:
            gray.pos.x=gray.pos.x+dp
        if gray.pos.x>=100:
            gfield.visible=0
            gdir.visible=0
            gray.pos=vector(0,5,20)
            triggergamma=0
    if pushphoton==1:
# which mirror is reflecting 
        if abs(photon.pos.x-75.0)<.05:
            refm1=1
            refm2=0
            photon.axis=(-20,0,0)
        if abs(photon.pos.x+75.0)<0.5:
            refm2=1
            refm1=0
# now reflect
        if refm1==1:
            photon.pos.x=photon.pos.x-dp
        if refm1!=1 or refm2==1:
            photon.axis=(20,0,0)
            photon.pos.x=photon.pos.x+dp
            
#Go Through East Arc
    if LINAC==-1:
        if b1.pos.x>=sss.end.x:
            rate(arcrate)
            earc_t_f=earc_t_f+dtheta
        if earc_t_f <=7.86:
            dearcx=20.0*(cos(earc_t_i)-cos(earc_t_f))
            dearcz=20.0*(sin(earc_t_i)-sin(earc_t_f))
            earc_t_i=earc_t_f
            b1.pos.x=b1.pos.x-dearcx
            b1.pos.z=b1.pos.z+dearcz
        if gamma==1 and s1.pos.x>=sss.end.x:
            rate(arcrate)
            earcs_t_f=earcs_t_f+dtheta
        if earcs_t_f <=7.86:
            dearcxs=20.0*(cos(earcs_t_i)-cos(earcs_t_f))
            dearczs=20.0*(sin(earcs_t_i)-sin(earcs_t_f))
            earcs_t_i=earcs_t_f
            s1.pos.x=s1.pos.x-dearcxs
            s1.pos.z=s1.pos.z+dearczs