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

def guided_modes_2D(epsilon, k0, h, num_eigs):
    # d2u/dx2 + d2u/dy2 + k0^2 eps u = beta^2 u
    # Assume x and y are the 0th and 1st indices of epsilon
    Nx, Ny = epsilon.shape
    diff2x = 1/h**2*diags([np.ones(Nx*Ny), -2*np.ones(Nx*Ny), np.ones(Nx*Ny)], [-1, 0, 1], shape=(Nx*Ny, Nx*Ny))
    diff2x = diff2x.todok()
    # Dirichlet u = 0 boundary conditions cut out some of these values.
    for i_y in range(1,Ny):
        i = Nx*i_y # 'Left wall'
        diff2x[i,i-1] = 0
    for i_y in range(Ny-1):
        i = Nx*i_y + Nx-1 # 'Right wall'
        diff2x[i,i+1] = 0
    diff2x = diff2x.tocsr()
    diff2y = 1/h**2*diags([np.ones(Nx*Ny), -2*np.ones(Nx*Ny), np.ones(Nx*Ny)], [-Nx, 0, Nx], shape=(Nx*Ny, Nx*Ny))
    epsm = diags(epsilon.flatten(), 0, shape=(Nx*Ny,Nx*Ny))
    A = diff2x + diff2y + k0**2*epsm
    vals, vecs = eigs(A, k=num_eigs, which='LR')
    # Reshape the eigenvectors to 2D arrays so one can visualize these things better
    modefields = np.reshape(vecs.flatten(), (Nx, Ny, num_eigs))
    return vals, modefields

def eps_SiO2(wl):
    # Malitson J. Opt. Soc. Am. 55, 1205, 1965
    return (1 + 0.6961663*wl**2/(wl**2 - 0.0684043**2) +
           0.4079426*wl**2/(wl**2 - 0.1162414**2) +
           0.8974794*wl**2/(wl**2 - 9.896161**2))

#x = np.linspace(-50, 50, 200)
#y = np.linspace(-50, 50, 200)
x = np.linspace(-20, 20, 100)
y = np.linspace(-20, 20, 100)
dx = x[1]-x[0]
x,y = np.meshgrid(x,y,indexing='ij')

#omega = 2*np.pi*np.linspace(187.5e12, 250e12, 50)
#omega = 2*np.pi*np.linspace(180e12, 390e12, 20)
lambda0 = np.linspace(1.2, 1.6, 50)
#lambda0 = 2*np.pi*3e8/omega*1e6
k0 = 2*np.pi/lambda0
betas = np.zeros_like(lambda0)
vg_analytical = np.zeros_like(lambda0)
t1 = time.time()
for i in range(len(lambda0)):
    eps = eps_SiO2(lambda0[i]) + 0.01*(x**2 + y**2 <= 6**2)
    beta2, u = guided_modes_2D(eps, k0[i], dx, 1)
    betas[i] = np.sqrt(beta2[0])
    N_points = x.shape[0]*x.shape[1]
t2 = time.time()
print(f'Time: {t2-t1}')
omega = 2*np.pi*3e8/lambda0*1e6

# Group velocity
#   vg = 1/(dbeta/domega)      omega = 2 pi c / wl   domega = -2 pi c / wl^2 * dwl
#      = -2 pi c /(wl^2 dbeta/dwl)
# Put everything in SI units first
betas = betas*1e6
lambda0 = lambda0*1e-6
k0 = k0*1e6
vg = np.diff(omega) / np.diff(betas)

plt.figure()
plt.plot(lambda0, betas/k0)

plt.figure()
plt.plot(lambda0[0:-1]*1e6, vg/3e8)#, lambda0, vg_analytical/3e8)
plt.xlabel('Wavelength (microns)')
plt.ylabel('v_g / c')

plt.show()