from visual import *

dt = 0.05
t = 0
counter = 0
m = 1.0

rand = random.random
randu = random.uniform
randc = random.choice

omega = 10.0   # trap frequency
alpha = 1.0    # confinement field strength

omega_L = 100  # laser frequency
beta = 0.001   # the parameter for the z-axis confinement
delv = 0.0001  # deltav per photon (hbar k/m)
c = 100.0     # in some units
vm = 0.1       # a big velocity

gamma = (vm/c)*omega_L
omega0 = omega_L+gamma/2  # define resonant frequency
dv = 0.5*vm/1000  # how big a dv kick per photon?

N0 = 10**2     # number of photons per time unit at resonance


def NFactor(omegap):
    return (gamma/2)**2/((gamma/2)**2 + (omegap-omega0)**2)

def Fz(s,t):
    z = s[2]
    vz = s[5]
    return -2*alpha*z*cos(omega*t) + beta*z
    
def Fy(s,t):
    y = s[1]
    vy = s[4]
    return 2*alpha*y*cos(omega*t) + beta*y
    
def Fx(s, t):
    x=s[0]
    vx=s[3]
    return 2*alpha*x*cos(omega*t) - beta*x

def F(s, t):
    return array([Fx(s,t), Fy(s,t), Fz(s,t)])

#def U(s, t):
#    x = 1#s[0]
#    y = 1#s[1]
#    z = 1#s[2]
#    print(x,y,z)
#    return 10#x**2*(-alpha*cos(omega*t) + beta*0.5) - (y**2 + z**2)*(alpha*cos(omega*t) + beta*0.5)

def v_betrag(s):
    return (s[3]**2 + s[4]**2 + s[5]**2)**(0.5)

def T(t):
    v = (v_betrag(s0)**2 + v_betrag(s1)**2 + v_betrag(s2)**2 +#
         v_betrag(s3)**2 + v_betrag(s4)**2 + v_betrag(s5)**2)/6
    U0 = -0.5*(Fx(s0,t)*s0[0] + Fy(s0,t)*s0[1] + Fz(s0,t)*s0[2])
    U1 = -0.5*(Fx(s1,t)*s1[0] + Fy(s1,t)*s1[1] + Fz(s1,t)*s1[2])
    U2 = -0.5*(Fx(s2,t)*s2[0] + Fy(s2,t)*s2[1] + Fz(s2,t)*s2[2])
    U3 = -0.5*(Fx(s3,t)*s3[0] + Fy(s3,t)*s3[1] + Fz(s3,t)*s3[2])
    U4 = -0.5*(Fx(s4,t)*s4[0] + Fy(s4,t)*s4[1] + Fz(s4,t)*s4[2])
    U5 = -0.5*(Fx(s5,t)*s5[0] + Fy(s5,t)*s5[1] + Fz(s5,t)*s5[2])
    U = (U0**2 + U1**2 + U2**2 + U3**2 + U4**2 + U5**2)**(0.5)
#    return v**2 + 2/3*(U(s0,t) + U(s1,t) + U(s2,t) + U(s3,t) + U(s4,t) + U(s5,t))
    print(v**2,U)
    return (v**2*10 + 0.001 * U)*100

def trap_derivs(s, t, params=None):
    v = s[3:]
    a = F(s, t)/m
    sp = zeros(6)
    sp[:3] = v
    sp[3:] = a
    return sp

def rk4(x, t, tau, derivsRK, param):
    
    half_tau = 0.5*tau
    
    F1 = derivsRK( x, t, param )
    t_half = t + half_tau
    xtemp = x + half_tau*F1
    
    F2 = derivsRK( xtemp, t_half, param)
    xtemp = x + half_tau*F2
    
    F3 = derivsRK( xtemp, t_half, param)
    t_full = t + tau
    xtemp = x + tau*F3

    F4 = derivsRK( xtemp, t_full, param)

    xout = x + (tau/6.0)*(F1 + F4 + 2.0*(F2 + F3))
    return xout

def move(s, ion):
    s = rk4(s, t, dt, trap_derivs, None)
    ion.pos=s[:3]

s0 = zeros(6, float64) # Positions- und Geschwindigkeitsvektor
s1 = zeros(6, float64)
s2 = zeros(6, float64)
s3 = zeros(6, float64)
s4 = zeros(6, float64)
s5 = zeros(6, float64)
s6 = zeros(6, float64)
s7 = zeros(6, float64)
s8 = zeros(6, float64)
s9 = zeros(6, float64)
def reset():
	s0[:3] = array([randu(-0.8,0.8),randu(-0.8,0.8),randu(-0.8,0.8)])
	s1[:3] = array([randu(-0.8,0.8),randu(-0.8,0.8),randu(-0.8,0.8)])
	s2[:3] = array([randu(-0.8,0.8),randu(-0.8,0.8),randu(-0.8,0.8)])
	s3[:3] = array([randu(-0.8,0.8),randu(-0.8,0.8),randu(-0.8,0.8)])
	s4[:3] = array([randu(-0.8,0.8),randu(-0.8,0.8),randu(-0.8,0.8)])
	s5[:3] = array([randu(-0.8,0.8),randu(-0.8,0.8),randu(-0.8,0.8)])
	s6[:3] = array([randu(-0.8,0.8),randu(-0.8,0.8),randu(-0.8,0.8)])
	s7[:3] = array([randu(-0.8,0.8),randu(-0.8,0.8),randu(-0.8,0.8)])
	s8[:3] = array([randu(-0.8,0.8),randu(-0.8,0.8),randu(-0.8,0.8)])
	s9[:3] = array([randu(-0.8,0.8),randu(-0.8,0.8),randu(-0.8,0.8)])
	
	s0[3:] = array([randu(0.03,0.04)*randc([-1,1]),randu(0.03,0.04)*randc([-1,1]),
                randu(0.03,0.04)*randc([-1,1])])
	s1[3:] = array([randu(0.03,0.04)*randc([-1,1]),randu(0.03,0.04)*randc([-1,1]),
                randu(0.03,0.04)*randc([-1,1])])
	s2[3:] = array([randu(0.03,0.04)*randc([-1,1]),randu(0.03,0.04)*randc([-1,1]),
                randu(0.03,0.04)*randc([-1,1])])
	s3[3:] = array([randu(0.03,0.04)*randc([-1,1]),randu(0.03,0.04)*randc([-1,1]),
                randu(0.03,0.04)*randc([-1,1])])
	s4[3:] = array([randu(0.03,0.04)*randc([-1,1]),randu(0.03,0.04)*randc([-1,1]),
                randu(0.03,0.04)*randc([-1,1])])
	s5[3:] = array([randu(0.03,0.04)*randc([-1,1]),randu(0.03,0.04)*randc([-1,1]),
                randu(0.03,0.04)*randc([-1,1])])
	s6[3:] = array([randu(0.03,0.04)*randc([-1,1]),randu(0.03,0.04)*randc([-1,1]),
                randu(0.03,0.04)*randc([-1,1])])
	s7[3:] = array([randu(0.03,0.04)*randc([-1,1]),randu(0.03,0.04)*randc([-1,1]),
                randu(0.03,0.04)*randc([-1,1])])
	s8[3:] = array([randu(0.03,0.04)*randc([-1,1]),randu(0.03,0.04)*randc([-1,1]),
                randu(0.03,0.04)*randc([-1,1])])
	s9[3:] = array([randu(0.03,0.04)*randc([-1,1]),randu(0.03,0.04)*randc([-1,1]),
                randu(0.03,0.04)*randc([-1,1])])
	ion0.pos=s0[:3]
        ion1.pos=s1[:3]
        ion2.pos=s2[:3]
        ion3.pos=s3[:3]
        ion4.pos=s4[:3]
        ion5.pos=s5[:3]
        ion6.pos=s6[:3]
        ion7.pos=s7[:3]
        ion8.pos=s8[:3]
        ion9.pos=s9[:3]

ion1 = sphere(pos=s1[:3], radius=0.05, color=color.cyan)
ion2 = sphere(pos=s2[:3], radius=0.05, color=color.cyan)
ion3 = sphere(pos=s3[:3], radius=0.05, color=color.cyan)
ion4 = sphere(pos=s4[:3], radius=0.05, color=color.cyan)
ion5 = sphere(pos=s5[:3], radius=0.05, color=color.cyan)
ion6 = sphere(pos=s6[:3], radius=0.05, color=color.yellow)
ion7 = sphere(pos=s7[:3], radius=0.05, color=color.yellow)
ion8 = sphere(pos=s8[:3], radius=0.05, color=color.yellow)
ion9 = sphere(pos=s9[:3], radius=0.05, color=color.yellow)
ion0 = sphere(pos=s0[:3], radius=0.05, color=color.yellow)

reset()

laserbeam_x = cylinder(pos=(-2.5,0,0), axis=(5,0,0), radius=1.1, color=color.red,
                       opacity=0.3)
laserbeam_y = cylinder(pos=(0,-2.5,0), axis=(0,5,0), radius=1.1, color=color.red,
                       opacity=0.3)
laserbeam_z = cylinder(pos=(0,0,-2.5), axis=(0,0,5), radius=1.1, color=color.red,
                       opacity=0.3)

laserbeam_x.visible = False
laserbeam_y.visible = False
laserbeam_z.visible = False

originx = arrow(pos=(0,0,0), axis=(0.5,0,0), shaftwidth=0.025, color=color.white)
originy = arrow(pos=(0,0,0), axis=(0,0.5,0), shaftwidth=0.025, color=color.white)
originz = arrow(pos=(0,0,0), axis=(0,0,0.5), shaftwidth=0.025, color=color.white)

laser_x1 = box(pos=(-4,0,0), length=2, height=0.5, width=0.5, opacity=1, color=color.red)
laser_x2 = box(pos=( 4,0,0), length=2, height=0.5, width=0.5, opacity=1, color=color.red)
laser_y1 = box(pos=(0,-4,0), length=0.5, height=2, width=0.5, opacity=1, color=color.red)
laser_y2 = box(pos=(0, 4,0), length=0.5, height=2, width=0.5, opacity=1, color=color.red)
laser_z1 = box(pos=(0,0,-4), length=0.5, height=0.5, width=2, opacity=1, color=color.red)
laser_z2 = box(pos=(0,0, 4), length=0.5, height=0.5, width=2, opacity=1, color=color.red)

update = False
time = 0

while 1:
    if update:
        
        s0 = rk4(s0, t, dt, trap_derivs, None)
        s1 = rk4(s1, t, dt, trap_derivs, None)
        s2 = rk4(s2, t, dt, trap_derivs, None)
        s3 = rk4(s3, t, dt, trap_derivs, None)
        s4 = rk4(s4, t, dt, trap_derivs, None)
        s5 = rk4(s5, t, dt, trap_derivs, None)
        s6 = rk4(s6, t, dt, trap_derivs, None)
        s7 = rk4(s7, t, dt, trap_derivs, None)
        s8 = rk4(s8, t, dt, trap_derivs, None)
        s9 = rk4(s9, t, dt, trap_derivs, None)
        ion0.pos=s0[:3]
        ion1.pos=s1[:3]
        ion2.pos=s2[:3]
        ion3.pos=s3[:3]
        ion4.pos=s4[:3]
        ion5.pos=s5[:3]
        ion6.pos=s6[:3]
        ion7.pos=s7[:3]
        ion8.pos=s8[:3]
        ion9.pos=s9[:3]
        t +=dt

    if laserbeam_x.visible == True:
        v1_x = s1[3]  # v_x determines doppler shift
        v2_x = s2[3]
        v3_x = s3[3]
        v4_x = s4[3]
        v5_x = s5[3]
        Nph1 = N0*dt*NFactor(omega_L*(1-v1_x/c))
        Nph2 = N0*dt*NFactor(omega_L*(1-v2_x/c))
        Nph3 = N0*dt*NFactor(omega_L*(1-v3_x/c))
        Nph4 = N0*dt*NFactor(omega_L*(1-v4_x/c))
        Nph5 = N0*dt*NFactor(omega_L*(1-v5_x/c))
        cth1 = 2*rand(Nph1)-1.0  # get cos(theta)
        cth2 = 2*rand(Nph2)-1.0
        cth3 = 2*rand(Nph3)-1.0
        cth4 = 2*rand(Nph4)-1.0
        cth5 = 2*rand(Nph5)-1.0
        sth1 = sqrt(1.0-cth1**2) # get sin(theta)
        sth2 = sqrt(1.0-cth2**2)
        sth3 = sqrt(1.0-cth3**2)
        sth4 = sqrt(1.0-cth4**2)
        sth5 = sqrt(1.0-cth5**2)
        phi1 = 2*pi*rand(Nph1)   # get phi
        phi2 = 2*pi*rand(Nph2)
        phi3 = 2*pi*rand(Nph3)
        phi4 = 2*pi*rand(Nph4)
        phi5 = 2*pi*rand(Nph5)
        sthcph1 = sth1*cos(phi1)  # get sin(theta)*cos(phi)
        sthcph2 = sth2*cos(phi2)
        sthcph3 = sth3*cos(phi3)
        sthcph4 = sth4*cos(phi4)
        sthcph5 = sth5*cos(phi5)
        sthsph1 = sth1*sin(phi1)  # get sin(theta)*sin(phi)
        sthsph2 = sth2*sin(phi2)
        sthsph3 = sth3*sin(phi3)
        sthsph4 = sth4*sin(phi4)
        sthsph5 = sth5*sin(phi5)

        dv1x = dv*(1.0-cth1).sum()
        dv2x = dv*(1.0-cth2).sum()
        dv3x = dv*(1.0-cth3).sum()
        dv4x = dv*(1.0-cth4).sum()
        dv5x = dv*(1.0-cth5).sum()
        dv1y = dv*(sthsph1).sum()
        dv2y = dv*(sthsph2).sum()
        dv3y = dv*(sthsph3).sum()
        dv4y = dv*(sthsph4).sum()
        dv5y = dv*(sthsph5).sum()
        dv1z = dv*sthcph1.sum()
        dv2z = dv*sthcph2.sum()
        dv3z = dv*sthcph3.sum()
        dv4z = dv*sthcph4.sum()
        dv5z = dv*sthcph5.sum()

        s1[3:] += array([dv1x,dv1y,dv1z])
        s2[3:] += array([dv2x,dv2y,dv2z])
        s3[3:] += array([dv3x,dv3y,dv3z])
        s4[3:] += array([dv4x,dv4y,dv4z])
        s5[3:] += array([dv5x,dv5y,dv5z])


    if laserbeam_y.visible == True:
        v1_y = s1[4]  # v_y determines doppler shift
        v2_y = s2[4]
        v3_y = s3[4]
        v4_y = s4[4]
        v5_y = s5[4]
        Nph1 = N0*dt*NFactor(omega_L*(1-v1_y/c))
        Nph2 = N0*dt*NFactor(omega_L*(1-v2_y/c))
        Nph3 = N0*dt*NFactor(omega_L*(1-v3_y/c))
        Nph4 = N0*dt*NFactor(omega_L*(1-v4_y/c))
        Nph5 = N0*dt*NFactor(omega_L*(1-v5_y/c))
        cth1 = 2*rand(Nph1)-1.0  # get cos(theta)
        cth2 = 2*rand(Nph2)-1.0
        cth3 = 2*rand(Nph3)-1.0
        cth4 = 2*rand(Nph4)-1.0
        cth5 = 2*rand(Nph5)-1.0
        sth1 = sqrt(1.0-cth1**2) # get sin(theta)
        sth2 = sqrt(1.0-cth2**2)
        sth3 = sqrt(1.0-cth3**2)
        sth4 = sqrt(1.0-cth4**2)
        sth5 = sqrt(1.0-cth5**2)
        phi1 = 2*pi*rand(Nph1)   # get phi
        phi2 = 2*pi*rand(Nph2)
        phi3 = 2*pi*rand(Nph3)
        phi4 = 2*pi*rand(Nph4)
        phi5 = 2*pi*rand(Nph5)
        sthcph1 = sth1*cos(phi1)  # get sin(theta)*cos(phi)
        sthcph2 = sth2*cos(phi2)
        sthcph3 = sth3*cos(phi3)
        sthcph4 = sth4*cos(phi4)
        sthcph5 = sth5*cos(phi5)
        sthsph1 = sth1*sin(phi1)  # get sin(theta)*sin(phi)
        sthsph2 = sth2*sin(phi2)
        sthsph3 = sth3*sin(phi3)
        sthsph4 = sth4*sin(phi4)
        sthsph5 = sth5*sin(phi5)

        dv1y = dv*(1.0-cth1).sum()
        dv2y = dv*(1.0-cth2).sum()
        dv3y = dv*(1.0-cth3).sum()
        dv4y = dv*(1.0-cth4).sum()
        dv5y = dv*(1.0-cth5).sum()
        dv1z = dv*(sthsph1).sum()
        dv2z = dv*(sthsph2).sum()
        dv3z = dv*(sthsph3).sum()
        dv4z = dv*(sthsph4).sum()
        dv5z = dv*(sthsph5).sum()
        dv1x = dv*sthcph1.sum()
        dv2x = dv*sthcph2.sum()
        dv3x = dv*sthcph3.sum()
        dv4x = dv*sthcph4.sum()
        dv5x = dv*sthcph5.sum()

        s1[3:] += array([dv1x,dv1y,dv1z])
        s2[3:] += array([dv2x,dv2y,dv2z])
        s3[3:] += array([dv3x,dv3y,dv3z])
        s4[3:] += array([dv4x,dv4y,dv4z])
        s5[3:] += array([dv5x,dv5y,dv5z])
        

    if laserbeam_z.visible == True:
        v1_z = s1[5]  # v_z determines doppler shift
        v2_z = s2[5]
        v3_z = s3[5]
        v4_z = s4[5]
        v5_z = s5[5]
        Nph1 = N0*dt*NFactor(omega_L*(1-v1_z/c))
        Nph2 = N0*dt*NFactor(omega_L*(1-v2_z/c))
        Nph3 = N0*dt*NFactor(omega_L*(1-v3_z/c))
        Nph4 = N0*dt*NFactor(omega_L*(1-v4_z/c))
        Nph5 = N0*dt*NFactor(omega_L*(1-v5_z/c))
        cth1 = 2*rand(Nph1)-1.0  # get cos(theta)
        cth2 = 2*rand(Nph2)-1.0
        cth3 = 2*rand(Nph3)-1.0
        cth4 = 2*rand(Nph4)-1.0
        cth5 = 2*rand(Nph5)-1.0
        sth1 = sqrt(1.0-cth1**2) # get sin(theta)
        sth2 = sqrt(1.0-cth2**2)
        sth3 = sqrt(1.0-cth3**2)
        sth4 = sqrt(1.0-cth4**2)
        sth5 = sqrt(1.0-cth5**2)
        phi1 = 2*pi*rand(Nph1)   # get phi
        phi2 = 2*pi*rand(Nph2)
        phi3 = 2*pi*rand(Nph3)
        phi4 = 2*pi*rand(Nph4)
        phi5 = 2*pi*rand(Nph5)
        sthcph1 = sth1*cos(phi1)  # get sin(theta)*cos(phi)
        sthcph2 = sth2*cos(phi2)
        sthcph3 = sth3*cos(phi3)
        sthcph4 = sth4*cos(phi4)
        sthcph5 = sth5*cos(phi5)
        sthsph1 = sth1*sin(phi1)  # get sin(theta)*sin(phi)
        sthsph2 = sth2*sin(phi2)
        sthsph3 = sth3*sin(phi3)
        sthsph4 = sth4*sin(phi4)
        sthsph5 = sth5*sin(phi5)

        dv1z = dv*(1.0-cth1).sum()
        dv2z = dv*(1.0-cth2).sum()
        dv3z = dv*(1.0-cth3).sum()
        dv4z = dv*(1.0-cth4).sum()
        dv5z = dv*(1.0-cth5).sum()
        dv1y = dv*(sthsph1).sum()
        dv2y = dv*(sthsph2).sum()
        dv3y = dv*(sthsph3).sum()
        dv4y = dv*(sthsph4).sum()
        dv5y = dv*(sthsph5).sum()
        dv1x = dv*sthcph1.sum()
        dv2x = dv*sthcph2.sum()
        dv3x = dv*sthcph3.sum()
        dv4x = dv*sthcph4.sum()
        dv5x = dv*sthcph5.sum()

        s1[3:] += array([dv1x,dv1y,dv1z])
        s2[3:] += array([dv2x,dv2y,dv2z])
        s3[3:] += array([dv3x,dv3y,dv3z])
        s4[3:] += array([dv4x,dv4y,dv4z])
        s5[3:] += array([dv5x,dv5y,dv5z])
        
    counter += 1
    if counter % 10==0:
        rate(30)
    if t > (time + 20):
        laserbeam_x.visible = False
        laserbeam_y.visible = False
        laserbeam_z.visible = False

    if scene.kb.keys: # event waiting to be processed?
        k = scene.kb.getkey() # get keyboard info
        if k in ['x','y','z','\n', 't','r']:
            if k=='\n':
                update ^= True
            elif k=='x':
                time = t
                laserbeam_x.visible ^= True
#                laserbeam_y.visible = False
#                laserbeam_z.visible = False
            elif k=='y':
                time = t
#                laserbeam_x.visible = False
                laserbeam_y.visible ^= True
#                laserbeam_z.visible = False
            elif k=='z':
                time = t
#                laserbeam_x.visible = False
#                laserbeam_y.visible = False
                laserbeam_z.visible ^= True
            elif k=='r':
                reset()
		update = False
	    elif k=='t':
#                U = U(s0,t) + U(s1,t) + U(s2,t) + U(s3,t) + U(s4,t) + U(s5,t)#
                U = T(t)
                print(U)