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

def guided_modes_1DTE(epsilon, k0, h):
    # d2u/dx2 + k0^2 eps u = beta^2 u
    N = len(epsilon)
    diff2x = 1/h**2*diags([np.ones(N), -2*np.ones(N), np.ones(N)], [-1, 0, 1], shape=(N,N))
    # Dirichlet u = 0 boundary conditions are automatically set in this case.
    epsm = diags(epsilon, 0, shape=(N,N))
    A = diff2x + k0**2*epsm
    vals, vecs = eigs(A, k=13, which='LR')
    return vals, vecs, A

# Dense matrix version of the above.
def guided_modes_1DTE_dense(epsilon, k0, h):
    # d2u/dx2 + k0^2 eps u = beta^2 u
    N = len(epsilon)
    #diff2x = 1/h**2*diags([np.ones(N), -2*np.ones(N), np.ones(N)], [-1, 0, 1], shape=(N,N))
    diff2x = np.diag(-2*np.ones(N), 0)
    diff2x += np.diag(np.ones(N-1), 1) + np.diag(np.ones(N-1), -1)
    diff2x /= h**2
    # Dirichlet u = 0 boundary conditions are automatically set in this case.
    #epsm = diags(epsilon, 0, shape=(N,N))
    epsm = np.diag(epsilon, 0)
    A = diff2x + k0**2*epsm
    #vals, vecs = eigs(A, k=13, which='LR')
    vals, vecs = eig(A)
    return vals, vecs, A


# If i is the index into the stacked vector, then i_x and i_y into the corresponding 2D array are
#   i_x = i % Nx
#   i_y = i // Nx
# and if we invert this relation,
#   i = Nx i_y + i_x
# Thus the bottom points (i_x, 0) are
#   i from 0 to Nx-1
# The top points (i_x, Ny-1) are
#   i from Nx (Ny-1) to Nx (Ny-1) + Nx-1
# The left points (0, i_y) are
#   i from 0 to Nx (Ny-1)
# The right points (Nx-1, i_y) are
#   i from Nx-1 to Nx (Ny-1) + Nx-1
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

# This version is like the above but it only uses 'diags' for the construction,
# no for loops needed.
def guided_modes_2D_alldiagonal(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
    offdiagonal = np.ones(Nx*Ny)
    offdiagonal[Nx-1::Nx] = 0
    diff2x = 1/h**2*diags([offdiagonal, -2*np.ones(Nx*Ny), offdiagonal], [-1, 0, 1], shape=(Nx*Ny, Nx*Ny))
    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 guided_modes_2D_periodic(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))
    # Now we need to change this to make it periodic.
    # For (0,0) the other side point is Nx (Ny-1)
    diff2y = diff2y.todok()
    for i_x in range(Nx):
        i = i_x
        diff2y[i,Nx*(Ny-1)+i] = 1/h**2 
        diff2y[Nx*(Ny-1)+i,i] = 1/h**2
    diff2y = diff2y.tocsc()
    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, diff2y

### 1D mode solver, graded index waveguide
#x = np.linspace(-120,120,400)
x = np.linspace(-100, 100, 500)
dx = x[1]-x[0]

eps = 2.25 + 0.015*np.exp(-x**2/15**2)
k0 = 2*np.pi/0.78

t1 = time.time()
beta2, u, A = guided_modes_1DTE(eps, k0, dx)
t2 = time.time()
print(f'Time: {t2-t1}')

t1 = time.time()
beta2dense, udense, Adense = guided_modes_1DTE_dense(eps, k0, dx)
t2 = time.time()
print(f'Time: {t2-t1}')

plt.figure()
for i in range(len(beta2)):
#for i in range(3):
    plt.plot(x, u[:,i])
plt.legend([f'beta/k0 = {np.sqrt(beta2[i])/k0}' for i in range(len(beta2))])

plt.figure()
plt.plot(x, u[:,0:3])
plt.legend([f'beta/k0 = {np.sqrt(beta2[i])/k0}' for i in range(3)])

print((np.sqrt(beta2)/k0)**2)

#plt.show()

### 2D mode solver, optical fiber
x = np.linspace(-30, 30, 100)
y = np.linspace(-30, 30, 100)
dx = x[1]-x[0]
x,y = np.meshgrid(x,y,indexing='ij')
#eps = 2.25 + 0.015*np.exp(-(x**2 + y**2)/15**2)
# I think a step index fiber might be more relevant here
k0 = 2*np.pi/0.78
eps = 2.11 + 0.01*(x**2 + y**2 <= 6**2)

t1 = time.time()
beta22, u2 = guided_modes_2D(eps, k0, dx, 8)
t2 = time.time()
print(f'Time: {t2-t1}')
print(np.sqrt(beta22)/k0)

fig,ax = plt.subplots(2,4)
for i in range(8):
    ax[i//4,i%4].pcolormesh(x, y, u2[:,:,i].real, vmin=-0.05, vmax=0.05)

#x = np.linspace(-20, 20, 130)
#y = np.linspace(-10, 10, 130)
x = np.linspace(-30, 30, 100)
y = np.linspace(-30, 30, 100)
dx = x[1]-x[0]
x,y = np.meshgrid(x,y,indexing='ij')
eps = 2.25 + 0.015*np.exp(-(y**2)/15**2)

t1 = time.time()
beta22, u2, diff2y = guided_modes_2D_periodic(eps, k0, dx, 20)
t2 = time.time()
print(f'Time: {t2-t1}')
print(np.sqrt(beta22)/k0)

fig,ax = plt.subplots(4,5)
for i in range(20):
    ax[i//5,i%5].pcolormesh(x, y, u2[:,:,i].real.T, vmin=-0.02, vmax=0.02)

plt.show()