import numpy as N
import scipy.constants as S
import scipy.special as SS
import matplotlib.pyplot as plt


#Definitions

def rkstepn(f,x,y,h):
    n=len(y)
    k1=[]
    k2=[]
    k3=[]
    k4=[]
    yplus=[]
    for i in range(n):
        k1.append(f[i](x,y))
    for i in range(n):
        k2.append(f[i](x+h/2,map(lambda j: j+h/2*k1[i], y)))
    for i in range(n):
        k3.append(f[i](x+h/2,map(lambda j: j+h/2*k2[i], y)))
    for i in range(n):
        k4.append(f[i](x+h,map(lambda j: j+h*k3[i], y)))
    for i in range(n):
        yplus.append(y[i]+h/6*(k1[i]+2*k2[i]+2*k3[i]+k4[i]))
    return yplus

def restrain(x):
    if x>N.pi: return restrain(x-2*N.pi)
    elif x<-N.pi: return restrain(x+2*N.pi)
    else: return x

#alpha and beta
def func(x,thetas):
    return thetas[0]*N.arctan(thetas[1]*x+thetas[2])+thetas[3]

def alpha(x):
    return func(x, thetasa)

def beta(x):
    return func(x, thetasb)

# new equations: dphi dp dgamma
# y=[p,phi,xi]

def dp(t,y):
    return -alpha(y[2])*N.sin(y[1])

def dxi(t,y):
    return y[0]/N.sqrt(1+y[0]**2)

def dphi(t,y):
    return 2*N.pi*((y[0]/N.sqrt(1+y[0]**2))*(1/beta(y[2]))-1)

def f2(xi,p,xik,pk):
    def z(x): return x*wavelength
    def mu(m): return SS.jn_zeros(0,m+1)[m]
    def s(m): return SS.jn(1,mu(m)*r/a)**2/((mu(m)**4)*(SS.jn(1,mu(m))**2))
    def g(m,x): return N.sign(x)*(1-N.exp(-mu(m)*abs(x)/a))
    def b(m,x): return 2*g(m,x)-g(m,x+2*d)-g(m,x-2*d)
    result=0
    for m in range(1,11):
        result+=s(m)*b(m,N.sqrt(1+p**2))*(z(xik)-z(xi))
    return result

#Initial conditions
p0=0.404 #starting impulse is const
#phi0=0 #phase, will do -1.57, 0, 1.57
beta0=0.374 #initial speed of particles
phimin=-N.pi
phimax=N.pi
number=10
a=0.01 #aperture
d=0.002 #disk thickness
r=0.002 #disk radius
wavelength=0.12
current=4 #0,0.2,0.4
const=(current*wavelength**2*a**2*S.e)/(2*N.pi*S.epsilon_0*d**2*r**2*S.m_e*S.c**3)

#(S.e*wavelength)/(S.m_e*S.c**2)*(current*wavelength)/(S.c*n*N.pi*r**2*2*d)*(a**2)/(S.epsilon_0*d) #constant

thetasa=[ 0.71920001,  2.82683308, -3.39808784,  1.30348301]
thetasb=[-0.25580129, -1.66465201,  2.53513356,  0.66269095]
h=0.01
L=7.8

class Particle:
    def __init__(self, phi0):
        self.phi = restrain(phi0)
        self.t = -(self.phi/N.pi-1)/2
        self.p=p0
        self.xi=0
        self.states=[[self.t,self.p,self.phi,self.xi]]
        self.ts=[self.t]
        self.ps=[self.p]
        self.phis=[self.phi]
        self.xis=[self.xi]
    def newstate(self):
        self.phi=restrain(self.phi)
        self.states.append([self.t,self.p,self.phi,self.xi])
        self.ts.append(self.t)
        self.ps.append(self.p)
        self.phis.append(self.phi)
        self.xis.append(self.xi)

t0s=[.1,.2,.3,.4,.5,.6,.7,.8,.9,1]
particles=[]
for phi0 in N.linspace(phimin, phimax, number):
    particles.append(Particle(phi0))

t=0
particlesOut=[]
#while particles:
while t<100:
    for part in particles:
        if t>part.t:
            integral=0
            for part2 in particles:
                integral+=f2(part.xi,part.p,part2.xi,part2.p)
            integral*=const/len(particles)
            fs=[lambda t,y: dp(t,y)+integral,dphi,dxi]
            ys=[part.p,part.phi,part.xi]
            rk=rkstepn(fs,t,ys,h)
            part.t=t
            part.p=rk[0]
            part.phi=rk[1]
            part.xi=rk[2]
            part.newstate()
            if part.xi>=L:
                particlesOut.append(part)
                particles.remove(part)
    if not particles: break
    t+=h

for part in particlesOut:
    print str(part.phis[0]) +': ' + str(part.xi)

print "---"
for part in particles:
    print str(part.phis[0]) +': ' + str(part.xi)
#    for i in range(len(part.states)):
#        print '[' + str(part.states[i][3]) + ',' +  str(part.states[i][2]) + ']'

plt.figure()
for part in particlesOut:
    plt.plot(part.xis, part.phis, label=r'$\tau_0$='+str(part.ts[0]))
plt.xlabel(r'$\xi$',fontsize=16)
plt.ylabel(r'$phi$',fontsize=16)
plt.yticks([-N.pi,0,N.pi],[r'$-\pi$',r'$0$',r'$\pi$'])
#plt.legend()
plt.show()