import numpy as np
from scipy import sparse
import scipy.sparse.linalg as splin
import matplotlib.pyplot as plt
import scipy.interpolate as inter


def RK4(f,t,dt,initial):  # the classical Runge-Kutta method (stage 4 method)
    m     = initial.shape[0]
    n     = int(t/dt)
    tlist = np.linspace(0,t,n)
    dt    = tlist[1] - tlist[0]
    y     = np.zeros((m,n),float)
    y[:,0]= initial
    for i in range(0,n-1):
        k1       = f(y[:,i])                    # Eq. (5.25)
        k2       = f(y[:,i] + 0.5*dt*k1)        # Eq. (5.26)
        k3       = f(y[:,i] + 0.5*dt*k2)        # Eq. (5.27)
        k4       = f(y[:,i] +     dt*k3)        # Eq. (5.28)
        y[:,i+1] = y[:,i] + (dt/6.0)*(k1 + 2*k2+2*k3  +k4)  # Eq. (5.24)

    return tlist, y


x0 = -0.15
xN = 0.15

y0 = -0.05
yN = 0.05

g = 0.03
d = 0.05

U = 1


Mx = 60
My = Mx/3

dx = (xN - x0) / Mx #step size in x
dy = (yN - y0) / My #step size in y

x = np.linspace(x0 + dx,xN - dx,Mx - 1)
y = np.linspace(y0 + dy,yN - dy,My - 1)

lenx = len(x)   #length of x
leny = len(y)   #length of y

def setupA(J,L,dx,dy,x):
    N = J*L     #NxN matrix size
    dia = np.zeros(N,float) #diagonal elemnets, corners
    dia[range(0,L)] = -3
    dia[range(L,N - L)] = -4
    dia[range(N-L,N)] = -3

    a = -0.055
    b = -0.025
    c = 0.025
    d = 0.055

    i = 0
    for k in x:
        if (k >= a and k < b) or (k >= c and k < d):
            dia[(i - 1)*L] = -3 #top
            dia[(i*L) - 1] = -3 #bot
        i += 1
    up = np.ones(N,float)
    up[range(L,N,L)] = 0

    dn = np.ones(N,float)
    dn[range(L -1, N - 1,L)] = 0

    offL = np.ones(N,float)
    vecs = np.array([offL,dn,dia,up,offL]) #join vecotrs
    diags = np.array([-L, -1, 0, 1, L])
    A = sparse.spdiags(vecs, diags, N, N) #sparse matrix
    A = A.tocsr() #from sparse matrix to a CRS - compressed sparse row
    return A




def setupboundary(J,L):   #Dirichlet boundary conditions
     N = J*L # NxN matrix size
     a = -0.025
     c = 0.025
     b = np.zeros(N,float)
     i = 0
     for k in x:
         if k >= a and k < c:
             b[(i - 1)*L] = -U #top
             b[(i*L) - 1] = -U #bottom
         i += 1
     return b


A = setupA(lenx,leny,dx,dy,x)#sparse matrix
b = setupboundary(lenx,leny)

sol = splin.spsolve(A,b)    #solve sparse
sol = np.reshape(sol,(lenx,leny)) #reshape to 2D
sol = sol.transpose()   #transpose

dh = dx

Ex = np.zeros((leny,lenx))  #field in x
for i in range(0,leny):
    for j in range(1,lenx - 1):
        Ex[i,j] = -(sol[i,j + 1] - sol[i,j - 1])/(2*dh)

Ey = np.zeros((leny,lenx))  #field in y
for i in range(0,lenx):
    for j in range(1,leny - 1):
        Ey[j,i] = - (sol[j + 1,i] - sol[j - 1,i])/(2*dh)

Exint = inter.interp2d(x,y,Ex,kind = 'cubic')
Eyint = inter.interp2d(x,y,Ey,kind = 'cubic')


XX,YY = np.meshgrid(x,y)
fig,ax = plt.subplots(1,1,figsize = (15,5))
cont = ax.contourf(XX,YY,sol,60)
plt.colorbar(cont)
ax.quiver(XX,YY,Ex,Ey)


ax.set_title('Electrostatic Potential of an Einzel Lens at %s V' %(U), fontsize = 25) #%s as placeholder, afterwards %(name)
ax.set_xlim(x0,xN)
ax.set_ylim(y0,yN)

ax.set_xlabel('$x/m$',fontsize = 15)
ax.set_ylabel('$y/m$',fontsize = 15)
ax.set_aspect('equal')
plt.show()

#################


U = 1.80 #change to get it focused


A = setupA(lenx,leny,dx,dy,x)#sparse matrix
b = setupboundary(lenx,leny)

sol = splin.spsolve(A,b)    #solve sparse
sol = np.reshape(sol,(lenx,leny)) #reshape to 2D
sol = sol.transpose()   #transpose

dh = dx

Ex = np.zeros((leny,lenx))  #field in x
for i in range(0,leny):
    for j in range(1,lenx - 1):
        Ex[i,j] = -(sol[i,j + 1] - sol[i,j - 1])/(2*dh)

Ey = np.zeros((leny,lenx))  #field in y
for i in range(0,lenx):
    for j in range(1,leny - 1):
        Ey[j,i] = - (sol[j + 1,i] - sol[j - 1,i])/(2*dh)

Exint = inter.interp2d(x,y,Ex,kind = 'cubic')
Eyint = inter.interp2d(x,y,Ey,kind = 'cubic')

q = -1.602e-19  #charge
m = 9.1e-31 #mass
const = q/m

XX,YY = np.meshgrid(x,y)
fig,ax = plt.subplots(1,1,figsize = (15,5))
cont = ax.contourf(XX,YY,sol,60)
plt.colorbar(cont)
ax.quiver(XX,YY,Ex,Ey)

#trajectory

def PewPew(fun):
    dy0 = fun[2]    #velocity in x
    dy1 = fun[3]    #velocity in y
    dy2 = const*Exint(fun[0],fun[1])[0] #acceleration in x [0] to take the value instead of array
    dy3 = const*Eyint(fun[0],fun[1])[0] #acceleration in y
    return np.array([dy0,dy1,dy2,dy3])

xtra = -0.15
ytra = np.linspace(0.03,-0.03,13)
vxtra = 0.5e6
vytra = 0

tfinal = 0.3/vxtra
dt = tfinal/lenx



for i in ytra:
    initial = np.array([xtra,i,vxtra,vytra])
    tRK4,yRK4 = RK4(PewPew,tfinal,dt,initial)
    ax.plot(yRK4[0,:],yRK4[1,:],'-r')

ax.set_title('The Trajectory of Electrons through an Einzel Lens at %s V' %(U), fontsize = 25) #%s as placeholder, afterwards %(name)
ax.set_xlim(x0,xN)
ax.set_ylim(y0,yN)

ax.set_xlabel('$x/m$',fontsize = 15)
ax.set_ylabel('$y/m$',fontsize = 15)
ax.set_aspect('equal')
plt.show()