import numpy as np
import matplotlib.pyplot as plt
from scipy.sparse.linalg import eigs
import time
from scipy.sparse import dok_matrix

# Grid nodes
x = np.concatenate([ np.arange(-2.0, -0.649, 0.05),
                     np.arange(-0.6, 0.6, 0.01),
                     np.arange(0.65, 2.0, 0.05) ])
xel = 0.5*(x[0:-1] + x[1:])
Nx = len(x)
Ne = Nx-1   # i:th element has x[i], x[i+1]

# Permittivity
eps = np.where(abs(xel) <= 0.4, 12.0, 2.25)

k0 = 2*np.pi

diff2x = dok_matrix((Nx,Nx))
wavenumop = dok_matrix((Nx,Nx))
mass = dok_matrix((Nx, Nx))

# Matrix contruction
t1 = time.time()
for ie in range(Ne):
    deltax = x[ie+1] - x[ie]  # length of element
    diff2x[ie,ie] += -1/deltax
    diff2x[ie+1,ie+1] += -1/deltax
    diff2x[ie,ie+1] -= -1/deltax
    diff2x[ie+1,ie] -= -1/deltax
    
    wavenumop[ie,ie] += k0**2 * eps[ie] * deltax * 1/3
    wavenumop[ie+1,ie+1] += k0**2 * eps[ie] *deltax * 1/3
    wavenumop[ie,ie+1] += k0**2 * eps[ie] *deltax * 1/6
    wavenumop[ie+1,ie] += k0**2 * eps[ie] *deltax * 1/6
    
    mass[ie,ie] += deltax * 1/3
    mass[ie+1,ie+1] += deltax * 1/3
    mass[ie,ie+1] += deltax * 1/6
    mass[ie+1,ie] += deltax * 1/6
t2 = time.time()
print(f'Construction: {t2-t1}')

A = diff2x + wavenumop
A = A.tocsc()
mass = mass.tocsc()
# Eigenvectors
t1 = time.time()
val,vec = eigs(A, M=mass, which='LR', k=6)
t2 = time.time()
print(f'Solution: {t2-t1}')

beta = np.sqrt(val)
n_eff = beta/k0

# Plotting
plt.rcParams['font.size'] = 18
plt.figure()
plt.plot(x, vec)
plt.plot(x,np.zeros(x.shape),'ok')
plt.legend([ f'n_eff = {n_eff[i]}' for i in range(len(n_eff)) ]) 
plt.show()